Atlas/internal/service/invoke/async_worker.go

package invoke

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	"github.com/rose_cat707/Atlas/internal/provider"
	"github.com/rose_cat707/Atlas/internal/store"
)

type asyncWorker struct {
	svc *Service
}

func newAsyncWorker(svc *Service) *asyncWorker {
	return &asyncWorker{svc: svc}
}

func (w *asyncWorker) run() {
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		w.pollOnce()
	}
}

func (w *asyncWorker) pollOnce() {
	ctx := context.Background()
	tasks, err := w.svc.store.ListPendingTasks(ctx)
	if err != nil {
		return
	}
	for _, task := range tasks {
		w.pollTask(ctx, task)
	}
}

func (w *asyncWorker) pollTask(ctx context.Context, task store.Task) {
	cfg, err := w.svc.store.GetProvider(ctx, task.ProviderID)
	if err != nil {
		return
	}
	pcfg := &provider.Config{
		ID:       cfg.ID,
		Protocol: cfg.Protocol,
		BaseURL:  cfg.BaseURL,
		APIKey:   cfg.APIKey,
	}
	path := fmt.Sprintf("/v1/predictions/%s", task.ExternalID)
	res, err := w.svc.router.Proxy(ctx, pcfg, "GET", path, nil, nil)
	if err != nil {
		return
	}
	if res.StatusCode >= 400 {
		task.Status = "failed"
		task.Error = string(res.Body)
		w.svc.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error {
			return st.UpdateTask(ctx, &task)
		})
		w.svc.notifyTask(task)
		return
	}
	var pred struct {
		Status   string `json:"status"`
		Output   any    `json:"output"`
		Error    string `json:"error"`
		Progress int    `json:"progress"`
	}
	_ = json.Unmarshal(res.Body, &pred)
	task.Progress = pred.Progress
	switch pred.Status {
	case "succeeded":
		task.Status = "succeeded"
		now := time.Now().UTC().Format(time.RFC3339)
		task.CompletedAt = &now
		out, _ := json.Marshal(map[string]any{"output": pred.Output})
		task.ResultJSON = string(out)
		pid := task.ProviderID
		gen := &store.Generation{
			SessionID:   task.SessionID,
			ServiceType: task.ServiceType,
			ProviderID:  &pid,
			InputJSON:   task.InputJSON,
			OutputJSON:  task.ResultJSON,
			TaskID:      &task.ID,
		}
		w.svc.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error {
			_, err := st.InsertGeneration(ctx, gen)
			return err
		})
	case "failed", "canceled":
		task.Status = pred.Status
		task.Error = pred.Error
		now := time.Now().UTC().Format(time.RFC3339)
		task.CompletedAt = &now
	case "processing", "starting":
		task.Status = "running"
	default:
		task.Status = "running"
	}
	w.svc.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error {
		return st.UpdateTask(ctx, &task)
	})
	w.svc.notifyTask(task)
}

func (s *Service) CancelTask(ctx context.Context, sessionID string, taskID int64) error {
	task, err := s.store.GetTask(ctx, sessionID, taskID)
	if err != nil {
		return err
	}
	cfg, err := s.store.GetProvider(ctx, task.ProviderID)
	if err != nil {
		return err
	}
	pcfg := &provider.Config{BaseURL: cfg.BaseURL, APIKey: cfg.APIKey}
	path := fmt.Sprintf("/v1/predictions/%s/cancel", task.ExternalID)
	meta := provider.CallMeta{
		SessionID: sessionID,
		Method:    "POST",
		Path:      "/v1/predictions/:id/cancel",
	}
	res, err := s.router.ProxyWithMeta(ctx, pcfg, "POST", path, nil, nil, meta)
	if err != nil {
		return err
	}
	if res.StatusCode >= 400 {
		return fmt.Errorf("cancel failed: %s", string(res.Body))
	}
	task.Status = "canceled"
	now := time.Now().UTC().Format(time.RFC3339)
	task.CompletedAt = &now
	if err := s.store.UpdateTask(ctx, task); err != nil {
		return err
	}
	s.notifyTask(*task)
	return nil
}