Files
LocalAI/core/services/agentpool/agent_jobs.go
LocalAI [bot] d7d7721eae feat(distributed): SyncedMap component + migrate finetune/quant/agent-tasks to cross-replica state (#10542)
* feat(distributed): add SyncedMap cross-replica in-memory state component

Introduce core/services/syncstate.SyncedMap[K,V]: a thread-safe in-memory map
that keeps itself consistent across frontend replicas via NATS, with an optional
pluggable durable Store and hydrate-from-source convergence.

Several features keep process-local state surfaced to the API (finetune/quant
jobs, agent tasks, model configs) and each hand-wired the same in-memory + NATS
broadcast + read-through-store legs - or forgot to, reintroducing cross-replica
staleness. SyncedMap makes that consistency a configuration choice:

- local writes mutate the map, write through the Store, then broadcast a delta;
- the apply path is memory-only and never re-publishes or re-writes the Store
  (structural echo-loop guard, mirroring galleryop.mergeStatus);
- on Start and on NATS reconnect the map re-hydrates from the source (Store, else
  Loader); an optional periodic Reconcile repairs silent drift;
- standalone mode (nil NATS client) is a strict in-memory no-op.

Reconnect re-hydrate is wired via a new *messaging.Client.OnReconnect callback,
consumed through an optional type-assertion so MessagingClient stays minimal.
Adds messaging.SubjectSyncStateDelta and a reusable testutil.FakeBus (synchronous
in-process MessagingClient with wildcard matching) for adopter tests.

Component only; service migrations follow in subsequent commits.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-8 [Claude Code]

* refactor(finetune): back jobs with SyncedMap for cross-replica consistency

FineTuneService kept jobs in a process-local map and, although it wrote them to
Postgres, ListJobs/GetJob never read the store back and the wired natsClient was
never used - so in distributed mode a job created on one replica was invisible to
the others. Replace the map and the dead client with a syncstate.SyncedMap keyed
by job ID, value *schema.FineTuneJob (the exact REST shape, so responses are
unchanged).

- Add a Store adapter (core/services/finetune/syncstore.go) over FineTuneStore,
  plus FineTuneStore.ListAll (global hydrate; per-user List kept) and an
  idempotent Upsert (create-or-update; Create alone fails on dup key).
- Writes go through SyncedMap.Set/Delete (write-through + broadcast); reads use
  List/Get. The on-disk state.json path becomes the standalone Loader, keeping
  single-node restart recovery (stale->stopped / exporting->failed fixups).
- Fold SetNATSClient/SetFineTuneStore into NewFineTuneService; app.go passes the
  distributed NATS client + store when distributed, nil otherwise.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-8 [Claude Code]

* refactor(agentpool): back agent tasks with SyncedMap for cross-replica consistency

AgentJobService.ListTasks read the process-local tasks map only, while ListJobs
already read through the DB persister + dispatcher NATS - so in distributed mode
a task created on one replica was invisible to the others. Back tasks with a
syncstate.SyncedMap keyed by task ID (value schema.Task, the exact REST shape);
jobs are left untouched.

- Store adapter (task_syncstore.go) over the existing JobPersister
  (LoadTasks/SaveTask/DeleteTask); reads svc.persister/userID live so a persister
  swap needs no rebuild. No new persister methods required.
- Task reads -> SyncedMap.List/Get; create/update -> Set (write-through +
  broadcast); delete -> Delete. The file persister now owns its own task set so
  the write-through path does not re-enter the SyncedMap lock (deadlock guard).
- The distributed NATS client is not available at construction (start() precedes
  initDistributed), so it is injected via SetTaskSyncNATS, which rebuilds the
  still-empty map before Start/hydrate. Wired at the main, restart, and per-user
  (UserServicesManager) distributed sites.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-8 [Claude Code]

* refactor(quantization): back jobs with SyncedMap + durable QuantStore

QuantizationService kept jobs in a process-local map persisted only to a local
state.json, so in distributed mode jobs were neither visible across replicas nor
durable cluster-wide. Back jobs with a syncstate.SyncedMap keyed by job ID
(value *schema.QuantizationJob, the exact REST shape).

- New distributed.QuantStore (GORM, table quantization_jobs) mirroring
  FineTuneStore: Create/Get/ListAll/Upsert(idempotent)/Delete, registered for
  AutoMigrate via distributed.InitStores (Stores.Quant).
- New adapter (quantization/syncstore.go) over QuantStore implementing
  syncstate.Store, with record<->schema conversion.
- Reads go through List/Get, writes through Set/Delete (write-through +
  broadcast); state.json is kept as the standalone Loader for single-node restart
  recovery (stale-job fixups preserved).
- app.go passes the distributed NATS client + QuantStore when distributed, nil
  otherwise; Start/Close lifecycle mirrors finetune.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-8 [Claude Code]

* fix(syncstate): annotate gosec G118 false positive on lifeCtx

gosec flagged the WithCancel in Start as "cancellation function not called"
because the returned cancel is stored on the struct rather than called/deferred
in scope. It is invoked in Close (covered by tests), and lifeCtx must outlive
Start to drive the reconnect/reconcile goroutines. Suppress the verified false
positive with a justified #nosec G118.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-8 [Claude Code]

* test(distributed): e2e two-replica SyncedMap sync over real NATS + Postgres

Adds the real-infrastructure counterpart to the fake-bus unit tests, in the
existing distributed e2e suite (testcontainers NATS + PostgreSQL). Two SyncedMap
instances stand in for two frontend replicas - each with its OWN NATS connection
to a shared server and a SHARED Postgres store (the distributed-mode invariant) -
and assert, over the wire:

- a create on replica A is observed by replica B;
- an update and a delete propagate A -> B (delete prunes, which a reload cannot);
- a late-joining replica recovers a job it never received a delta for, via store
  hydrate on Start (the at-most-once gap a fake bus cannot exercise);
- a local Set is written through to the shared Postgres store.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-8 [Claude Code]

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
2026-06-27 23:23:51 +02:00

1503 lines
44 KiB
Go

package agentpool
import (
"bytes"
"cmp"
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"path/filepath"
"slices"
"strings"
"sync"
"sync/atomic"
"text/template"
"time"
"github.com/Masterminds/sprig/v3"
"github.com/google/uuid"
"github.com/mudler/cogito"
"github.com/mudler/cogito/clients"
"github.com/mudler/xlog"
"github.com/robfig/cron/v3"
"github.com/mudler/LocalAI/core/config"
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/jobs"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/syncstate"
"github.com/mudler/LocalAI/core/templates"
"github.com/mudler/LocalAI/pkg/httpclient"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/pkg/xsync"
)
// AgentJobService manages agent tasks and job execution
type AgentJobService struct {
appConfig *config.ApplicationConfig
modelLoader *model.ModelLoader
configLoader *config.ModelConfigLoader
evaluator *templates.Evaluator
// tasks is the cross-replica task store: an in-memory map kept consistent
// across replicas via NATS, with read-through to the configured persister
// (file in standalone, PostgreSQL in distributed). Unlike jobs - which already
// converge via the dispatcher + DB read-through - tasks previously read
// in-memory only, so ListTasks went stale on non-originating replicas.
tasks *syncstate.SyncedMap[string, schema.Task]
// taskNats is the distributed NATS client backing the tasks SyncedMap. It is
// not available at construction time, so it is injected via SetTaskSyncNATS
// during distributed wiring; nil keeps tasks in-memory-only (standalone).
taskNats messaging.MessagingClient
// Storage (in-memory primary, persister for secondary persistence)
jobs *xsync.SyncedMap[string, schema.Job]
persister JobPersister
userID string // Scoping: empty for global (main service), set for per-user instances
// Job execution channel
jobQueue chan JobExecution
// Cancellation support
cancellations *xsync.SyncedMap[string, context.CancelFunc]
// Cron scheduler
cronScheduler *cron.Cron
cronEntries *xsync.SyncedMap[string, cron.EntryID]
// Job retention
retentionDays int // From runtime settings, default: 30
// Distributed mode (nil when not in distributed mode)
dispatcher DistributedDispatcher
rawDBStore *jobs.JobStore // kept for DBStore() accessor
// Service lifecycle
ctx context.Context
cancel context.CancelFunc
}
// DistributedDispatcher is the interface for distributed job dispatching via NATS.
// Satisfied by *jobs.Dispatcher.
type DistributedDispatcher interface {
Enqueue(jobID, taskID, userID string) error
Cancel(jobID string) error
}
// SetDistributedBackends sets the dispatcher for distributed mode.
// When set, ExecuteJob/CancelJob delegate to NATS and Start() skips local workers.
func (s *AgentJobService) SetDistributedBackends(dispatcher DistributedDispatcher) {
s.dispatcher = dispatcher
}
// SetUserID sets the user ID for per-user scoping of DB queries.
func (s *AgentJobService) SetUserID(id string) {
s.userID = id
}
// SetDistributedJobStore sets the database-backed job store for persisting tasks/jobs.
// When set, the DB is the source of truth instead of JSON files.
func (s *AgentJobService) SetDistributedJobStore(store *jobs.JobStore) {
s.rawDBStore = store
s.persister = &dbJobPersister{store: store}
}
// SetTaskSyncNATS wires the distributed NATS client used to keep agent *tasks*
// consistent across replicas (jobs already converge via the dispatcher + DB
// read-through, so they are left untouched). The client is not available when the
// service is constructed, so it is injected here during distributed wiring and the
// tasks SyncedMap is rebuilt to pick it up. It is always called before Start /
// hydrate, while the map is still empty, so rebuilding loses no state. Passing nil
// (standalone) keeps the map in-memory-only with no broadcast.
func (s *AgentJobService) SetTaskSyncNATS(nats messaging.MessagingClient) {
s.taskNats = nats
s.buildTasksMap()
}
// buildTasksMap (re)constructs the cross-replica tasks SyncedMap from the current
// taskNats. The Store adapter reads s.persister/s.userID live, so a persister swap
// (SetDistributedJobStore) needs no rebuild; only the NATS client, fixed at
// New-time, forces one - hence SetTaskSyncNATS calls this.
func (s *AgentJobService) buildTasksMap() {
s.tasks = syncstate.New(syncstate.Config[string, schema.Task]{
Name: "agent.tasks",
Key: func(t schema.Task) string { return t.ID },
Nats: s.taskNats,
Store: &taskStoreAdapter{svc: s},
})
}
// Dispatcher returns the distributed dispatcher (nil if not in distributed mode).
func (s *AgentJobService) Dispatcher() DistributedDispatcher {
return s.dispatcher
}
// DBStore returns the database-backed job store (nil if not configured).
func (s *AgentJobService) DBStore() *jobs.JobStore {
return s.rawDBStore
}
// saveJobs persists jobs via the configured persister (file or DB).
func (s *AgentJobService) saveJobs(job schema.Job) {
if err := s.persister.SaveJob(s.userID, job); err != nil {
xlog.Warn("Failed to persist job", "error", err, "job_id", job.ID)
}
}
// LoadFromDB populates the in-memory maps from the database.
// LoadFromDB loads tasks and jobs from the configured persister.
// Kept for backward compatibility with user_services.go.
func (s *AgentJobService) LoadFromDB() {
s.loadFromPersister()
}
// loadFromPersister loads tasks and jobs from the configured persister into memory.
func (s *AgentJobService) loadFromPersister() {
if err := s.hydrateTasks(s.appConfig.Context); err != nil {
xlog.Warn("Failed to load tasks from persister", "error", err)
}
if loadedJobs, err := s.persister.LoadJobs(s.userID); err != nil {
xlog.Warn("Failed to load jobs from persister", "error", err)
} else {
for _, job := range loadedJobs {
s.jobs.Set(job.ID, job)
}
xlog.Info("Loaded jobs from persister", "count", len(loadedJobs))
}
}
// hydrateTasks loads tasks into the cross-replica SyncedMap and (re)schedules
// cron entries for enabled tasks. Hydration goes through the SyncedMap's Store
// read-through (Start), not Set, so it neither re-persists nor re-broadcasts the
// loaded tasks. Each service instance hydrates exactly once: the main service via
// Start -> loadFromPersister, per-user services via LoadFromDB or LoadTasksFromFile.
func (s *AgentJobService) hydrateTasks(ctx context.Context) error {
if err := s.tasks.Start(ctx); err != nil {
return err
}
tasks := s.tasks.List()
for _, task := range tasks {
if task.Enabled && task.Cron != "" {
if err := s.ScheduleCronTask(task); err != nil {
xlog.Warn("Failed to schedule cron task on load", "error", err, "task_id", task.ID)
}
}
}
xlog.Info("Loaded tasks from persister", "count", len(tasks))
return nil
}
// JobExecution represents a job to be executed
type JobExecution struct {
Job schema.Job
Task schema.Task
Ctx context.Context
Cancel context.CancelFunc
}
const (
JobImageType = "image"
JobVideoType = "video"
JobAudioType = "audio"
JobFileType = "file"
)
// NewAgentJobService creates a new AgentJobService instance
func NewAgentJobService(
appConfig *config.ApplicationConfig,
modelLoader *model.ModelLoader,
configLoader *config.ModelConfigLoader,
evaluator *templates.Evaluator,
) *AgentJobService {
// Determine storage directory: DataPath > DynamicConfigsDir
tasksFile := ""
jobsFile := ""
dataDir := appConfig.DataPath
if dataDir == "" {
dataDir = appConfig.DynamicConfigsDir
}
if dataDir != "" {
tasksFile = filepath.Join(dataDir, "agent_tasks.json")
jobsFile = filepath.Join(dataDir, "agent_jobs.json")
}
return NewAgentJobServiceWithPaths(appConfig, modelLoader, configLoader, evaluator, tasksFile, jobsFile)
}
// NewAgentJobServiceWithPaths creates a new AgentJobService with explicit file paths.
func NewAgentJobServiceWithPaths(
appConfig *config.ApplicationConfig,
modelLoader *model.ModelLoader,
configLoader *config.ModelConfigLoader,
evaluator *templates.Evaluator,
tasksFile, jobsFile string,
) *AgentJobService {
retentionDays := cmp.Or(appConfig.AgentJobRetentionDays, 30)
jobsMap := xsync.NewSyncedMap[string, schema.Job]()
s := &AgentJobService{
appConfig: appConfig,
modelLoader: modelLoader,
configLoader: configLoader,
evaluator: evaluator,
jobs: jobsMap,
persister: &fileJobPersister{
jobs: jobsMap,
tasksFile: tasksFile,
jobsFile: jobsFile,
taskSet: make(map[string]schema.Task),
},
jobQueue: make(chan JobExecution, 100), // Buffer for 100 jobs
cancellations: xsync.NewSyncedMap[string, context.CancelFunc](),
cronScheduler: cron.New(), // Support seconds in cron
cronEntries: xsync.NewSyncedMap[string, cron.EntryID](),
retentionDays: retentionDays,
}
// Build the cross-replica tasks map standalone (nil NATS); SetTaskSyncNATS
// rebuilds it with the distributed client once that is available, before Start.
s.buildTasksMap()
return s
}
// LoadTasksFromFile loads tasks from the persister into the in-memory map
// and schedules cron entries. Named "FromFile" for backward compat; in DB
// mode it loads from the database.
func (s *AgentJobService) LoadTasksFromFile() error {
return s.hydrateTasks(s.appConfig.Context)
}
// SaveTasksToFile flushes the current tasks map via the persister. File
// persister bulk-writes the JSON file atomically; DB persister no-ops
// because per-task SaveTask calls already wrote through.
func (s *AgentJobService) SaveTasksToFile() error {
return s.persister.FlushTasks()
}
// LoadJobsFromFile loads jobs from the persister into the in-memory map.
// Named "FromFile" for backward compat; in DB mode it loads from the
// database.
func (s *AgentJobService) LoadJobsFromFile() error {
jobs, err := s.persister.LoadJobs(s.userID)
if err != nil {
return err
}
for _, job := range jobs {
s.jobs.Set(job.ID, job)
}
return nil
}
// SaveJobsToFile flushes the current jobs map via the persister. File
// persister bulk-writes the JSON file atomically; DB persister no-ops
// because per-job SaveJob calls already wrote through.
func (s *AgentJobService) SaveJobsToFile() error {
return s.persister.FlushJobs()
}
// CreateTask creates a new task
func (s *AgentJobService) CreateTask(task schema.Task) (string, error) {
if task.Name == "" {
return "", fmt.Errorf("task name is required")
}
if task.Model == "" {
return "", fmt.Errorf("task model is required")
}
if task.Prompt == "" {
return "", fmt.Errorf("task prompt is required")
}
// Generate UUID
id := uuid.New().String()
task.ID = id
now := time.Now()
task.CreatedAt = now
task.UpdatedAt = now
if !task.Enabled {
task.Enabled = true // Default to enabled
}
// Store task: Set updates the in-memory map, write-throughs to the persister
// (file or DB), and broadcasts the create to peer replicas. Background ctx
// because CreateTask carries no request ctx (mirrors the finetune service).
if err := s.tasks.Set(context.Background(), task); err != nil {
return "", fmt.Errorf("failed to persist task: %w", err)
}
// Schedule cron if enabled and has cron expression
if task.Enabled && task.Cron != "" {
if err := s.ScheduleCronTask(task); err != nil {
xlog.Warn("Failed to schedule cron task", "error", err, "task_id", id)
}
}
return id, nil
}
// UpdateTask updates an existing task
func (s *AgentJobService) UpdateTask(id string, task schema.Task) error {
existing, ok := s.tasks.Get(id)
if !ok {
return fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
// Preserve ID and CreatedAt
task.ID = id
task.CreatedAt = existing.CreatedAt
task.UpdatedAt = time.Now()
// Unschedule old cron if it had one
if existing.Cron != "" {
s.UnscheduleCronTask(id)
}
// Store updated task: write-through + broadcast (see CreateTask).
if err := s.tasks.Set(context.Background(), task); err != nil {
return fmt.Errorf("failed to persist task: %w", err)
}
// Schedule new cron if enabled and has cron expression
if task.Enabled && task.Cron != "" {
if err := s.ScheduleCronTask(task); err != nil {
xlog.Warn("Failed to schedule cron task", "error", err, "task_id", id)
}
}
return nil
}
// DeleteTask deletes a task
func (s *AgentJobService) DeleteTask(id string) error {
if _, ok := s.tasks.Get(id); !ok {
return fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
// Unschedule cron
s.UnscheduleCronTask(id)
// Delete removes from the in-memory map, deletes from the persister, and
// broadcasts the removal to peer replicas.
if err := s.tasks.Delete(context.Background(), id); err != nil {
xlog.Warn("Failed to delete task from store", "error", err, "task_id", id)
}
return nil
}
// GetTask retrieves a task by ID
func (s *AgentJobService) GetTask(id string) (*schema.Task, error) {
task, ok := s.tasks.Get(id)
if !ok {
return nil, fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
return &task, nil
}
// ListTasks returns all tasks, sorted by creation date (newest first)
func (s *AgentJobService) ListTasks() []schema.Task {
tasks := s.tasks.List()
// Sort by CreatedAt descending (newest first), then by Name for stability
slices.SortFunc(tasks, func(a, b schema.Task) int {
if a.CreatedAt.Equal(b.CreatedAt) {
return cmp.Compare(a.Name, b.Name)
}
return b.CreatedAt.Compare(a.CreatedAt)
})
return tasks
}
// buildPrompt builds a prompt from a template with parameters
func (s *AgentJobService) buildPrompt(templateStr string, params map[string]string) (string, error) {
tmpl, err := template.New("prompt").Parse(templateStr)
if err != nil {
return "", fmt.Errorf("failed to parse prompt template: %w", err)
}
var buf bytes.Buffer
if err := tmpl.Execute(&buf, params); err != nil {
return "", fmt.Errorf("failed to execute prompt template: %w", err)
}
return buf.String(), nil
}
// ExecuteJob creates and queues a job for execution
// multimedia can be nil for backward compatibility
func (s *AgentJobService) ExecuteJob(taskID string, params map[string]string, triggeredBy string, multimedia *schema.MultimediaAttachment) (string, error) {
task, ok := s.tasks.Get(taskID)
if !ok {
return "", fmt.Errorf("%w: %s", ErrTaskNotFound, taskID)
}
if !task.Enabled {
return "", fmt.Errorf("%w: %s", ErrTaskDisabled, taskID)
}
// Create job
jobID := uuid.New().String()
now := time.Now()
job := schema.Job{
ID: jobID,
TaskID: taskID,
Status: schema.JobStatusPending,
Parameters: params,
CreatedAt: now,
TriggeredBy: triggeredBy,
}
// Handle multimedia: merge task-level (for cron) and job-level (for manual execution)
if triggeredBy == "cron" && len(task.MultimediaSources) > 0 {
// Fetch multimedia from task sources
job.Images = []string{}
job.Videos = []string{}
job.Audios = []string{}
job.Files = []string{}
for _, source := range task.MultimediaSources {
// Fetch content from URL with custom headers
dataURI, err := s.fetchMultimediaFromURL(source.URL, source.Headers, source.Type)
if err != nil {
xlog.Warn("Failed to fetch multimedia from task source", "error", err, "url", source.URL, "type", source.Type)
continue
}
// Add to appropriate slice based on type
switch source.Type {
case JobImageType:
job.Images = append(job.Images, dataURI)
case JobVideoType:
job.Videos = append(job.Videos, dataURI)
case JobAudioType:
job.Audios = append(job.Audios, dataURI)
case JobFileType:
job.Files = append(job.Files, dataURI)
}
}
}
// Override with job-level multimedia if provided (manual execution takes precedence)
if multimedia != nil {
if len(multimedia.Images) > 0 {
job.Images = multimedia.Images
}
if len(multimedia.Videos) > 0 {
job.Videos = multimedia.Videos
}
if len(multimedia.Audios) > 0 {
job.Audios = multimedia.Audios
}
if len(multimedia.Files) > 0 {
job.Files = multimedia.Files
}
}
// Store job
s.jobs.Set(jobID, job)
// Distributed mode: delegate to NATS dispatcher
if s.dispatcher != nil {
go func() { s.saveJobs(job) }()
xlog.Info("Enqueuing job via distributed dispatcher", "job_id", jobID, "task_id", taskID)
return jobID, s.dispatcher.Enqueue(jobID, taskID, "")
}
go func() { s.saveJobs(job) }()
// Create context for cancellation
ctx, cancel := context.WithCancel(context.Background())
s.cancellations.Set(jobID, cancel)
// Queue job
select {
case s.jobQueue <- JobExecution{
Job: job,
Task: task,
Ctx: ctx,
Cancel: cancel,
}:
default:
// Queue is full, update job status
job.Status = schema.JobStatusFailed
job.Error = "job queue is full"
s.jobs.Set(jobID, job)
return "", ErrJobQueueFull
}
return jobID, nil
}
// GetJob retrieves a job by ID
func (s *AgentJobService) GetJob(id string) (*schema.Job, error) {
// Try authoritative read from persister (DB returns fresh data; file returns nil)
if job, err := s.persister.GetJob(id); err == nil && job != nil {
s.jobs.Set(job.ID, *job) // sync back to in-memory
return job, nil
}
// Fall back to in-memory
job := s.jobs.Get(id)
if job.ID == "" {
return nil, fmt.Errorf("%w: %s", ErrJobNotFound, id)
}
return &job, nil
}
// ListJobs returns jobs, optionally filtered by task_id and status
func (s *AgentJobService) ListJobs(taskID *string, status *schema.JobStatus, limit int) []schema.Job {
// Try authoritative list from persister (DB returns fresh data; file returns nil)
taskFilter := ""
if taskID != nil {
taskFilter = *taskID
}
statusFilter := ""
if status != nil {
statusFilter = string(*status)
}
if result, err := s.persister.ListJobs(s.userID, taskFilter, statusFilter, limit); err == nil && result != nil {
for i := range result {
s.jobs.Set(result[i].ID, result[i]) // sync back
}
return result
}
// Fall back to in-memory filtering
allJobs := s.jobs.Values()
filtered := []schema.Job{}
for _, job := range allJobs {
if taskID != nil && job.TaskID != *taskID {
continue
}
if status != nil && job.Status != *status {
continue
}
filtered = append(filtered, job)
}
// Sort by CreatedAt descending (newest first)
slices.SortFunc(filtered, func(a, b schema.Job) int { return b.CreatedAt.Compare(a.CreatedAt) })
// Apply limit
if limit > 0 && limit < len(filtered) {
filtered = filtered[:limit]
}
return filtered
}
// CancelJob cancels a running job
func (s *AgentJobService) CancelJob(id string) error {
job := s.jobs.Get(id)
if job.ID == "" {
return fmt.Errorf("%w: %s", ErrJobNotFound, id)
}
if job.Status != schema.JobStatusPending && job.Status != schema.JobStatusRunning {
return fmt.Errorf("job cannot be cancelled: status is %s", job.Status)
}
// Distributed mode: delegate cancel via NATS
if s.dispatcher != nil {
return s.dispatcher.Cancel(id)
}
// Cancel context
if s.cancellations.Exists(id) {
cancel := s.cancellations.Get(id)
cancel()
s.cancellations.Delete(id)
}
// Update job status
now := time.Now()
job.Status = schema.JobStatusCancelled
job.CompletedAt = &now
s.jobs.Set(id, job)
go func() { s.saveJobs(job) }()
return nil
}
// DeleteJob deletes a job
func (s *AgentJobService) DeleteJob(id string) error {
if !s.jobs.Exists(id) {
return fmt.Errorf("%w: %s", ErrJobNotFound, id)
}
s.jobs.Delete(id)
if err := s.persister.DeleteJob(id); err != nil {
xlog.Warn("Failed to delete job from persister", "error", err, "job_id", id)
}
return nil
}
// traceCollector synchronises concurrent cogito callback writes to job traces.
type traceCollector struct {
mu sync.Mutex
jobID string
traces []schema.JobTrace
jobs *xsync.SyncedMap[string, schema.Job]
}
func (tc *traceCollector) add(trace schema.JobTrace) {
tc.mu.Lock()
defer tc.mu.Unlock()
tc.traces = append(tc.traces, trace)
if tc.jobs.Exists(tc.jobID) {
job := tc.jobs.Get(tc.jobID)
job.Traces = make([]schema.JobTrace, len(tc.traces))
copy(job.Traces, tc.traces)
tc.jobs.Set(tc.jobID, job)
}
}
type multimediaContent struct {
url string
mediaType string
}
func (mu multimediaContent) URL() string {
return mu.url
}
// fetchMultimediaFromURL fetches multimedia content from a URL with custom headers
// and converts it to a data URI string
func (s *AgentJobService) fetchMultimediaFromURL(url string, headers map[string]string, mediaType string) (string, error) {
// Create HTTP request
req, err := http.NewRequest("GET", url, nil)
if err != nil {
return "", fmt.Errorf("failed to create request: %w", err)
}
// Set custom headers
for key, value := range headers {
req.Header.Set(key, value)
}
// Execute request
client := httpclient.NewWithTimeout(30 * time.Second)
resp, err := client.Do(req)
if err != nil {
return "", fmt.Errorf("failed to fetch URL: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return "", fmt.Errorf("HTTP error: %d", resp.StatusCode)
}
// Read content
data, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("failed to read response: %w", err)
}
// Encode to base64
encoded := base64.StdEncoding.EncodeToString(data)
// Determine MIME type
mimeType := s.getMimeTypeForMediaType(mediaType)
if contentType := resp.Header.Get("Content-Type"); contentType != "" {
mimeType = contentType
}
// Return as data URI
return fmt.Sprintf("data:%s;base64,%s", mimeType, encoded), nil
}
// getMimeTypeForMediaType returns the default MIME type for a media type
func (s *AgentJobService) getMimeTypeForMediaType(mediaType string) string {
switch mediaType {
case JobImageType:
return "image/png"
case JobVideoType:
return "video/mp4"
case JobAudioType:
return "audio/mpeg"
case JobFileType:
return "application/octet-stream"
default:
return "application/octet-stream"
}
}
// convertToMultimediaContent converts a slice of strings (URLs or base64) to multimediaContent objects
func (s *AgentJobService) convertToMultimediaContent(items []string, mediaType string) ([]cogito.Multimedia, error) {
result := make([]cogito.Multimedia, 0, len(items))
for _, item := range items {
if item == "" {
continue
}
// Check if it's already a data URI
if strings.HasPrefix(item, "data:") {
result = append(result, multimediaContent{url: item, mediaType: mediaType})
continue
}
// Check if it's a URL
if strings.HasPrefix(item, "http://") || strings.HasPrefix(item, "https://") {
// Pass URL directly to cogito (it handles fetching)
result = append(result, multimediaContent{url: item, mediaType: mediaType})
continue
}
// Assume it's base64 without data URI prefix
// Add appropriate prefix based on media type
mimeType := s.getMimeTypeForMediaType(mediaType)
dataURI := fmt.Sprintf("data:%s;base64,%s", mimeType, item)
result = append(result, multimediaContent{url: dataURI, mediaType: mediaType})
}
return result, nil
}
// ExecuteJobInternal executes a job using cogito.
// Exported so the distributed Dispatcher can bridge NATS-dequeued jobs back into this service.
func (s *AgentJobService) ExecuteJobInternal(job schema.Job, task schema.Task, ctx context.Context) error {
// Update job status to running
now := time.Now()
job.Status = schema.JobStatusRunning
job.StartedAt = &now
s.jobs.Set(job.ID, job)
xlog.Info("Job started", "job_id", job.ID, "task_id", job.TaskID)
// Load model config
modelConfig, err := s.configLoader.LoadModelConfigFileByNameDefaultOptions(task.Model, s.appConfig)
if err != nil {
job.Status = schema.JobStatusFailed
job.Error = fmt.Sprintf("failed to load model config: %v", err)
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("failed to load model config: %w", err)
}
// Validate MCP configuration
if modelConfig.MCP.Servers == "" && modelConfig.MCP.Stdio == "" {
job.Status = schema.JobStatusFailed
job.Error = "no MCP servers configured for model"
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("no MCP servers configured for model: %s", task.Model)
}
// Get MCP config from model config
remote, stdio, err := modelConfig.MCP.MCPConfigFromYAML()
if err != nil {
job.Status = schema.JobStatusFailed
job.Error = fmt.Sprintf("failed to get MCP config: %v", err)
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("failed to get MCP config: %w", err)
}
// Get MCP sessions
sessions, err := mcpTools.SessionsFromMCPConfig(modelConfig.Name, remote, stdio)
if err != nil {
job.Status = schema.JobStatusFailed
job.Error = fmt.Sprintf("failed to get MCP sessions: %v", err)
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("failed to get MCP sessions: %w", err)
}
if len(sessions) == 0 {
job.Status = schema.JobStatusFailed
job.Error = "no working MCP servers found"
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("no working MCP servers found")
}
// Build prompt from template
prompt, err := s.buildPrompt(task.Prompt, job.Parameters)
if err != nil {
job.Status = schema.JobStatusFailed
job.Error = fmt.Sprintf("failed to build prompt: %v", err)
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("failed to build prompt: %w", err)
}
// Create cogito fragment
fragment := cogito.NewEmptyFragment()
// Collect all multimedia content
multimediaItems := []cogito.Multimedia{}
// Convert images
if len(job.Images) > 0 {
images, err := s.convertToMultimediaContent(job.Images, JobImageType)
if err != nil {
xlog.Warn("Failed to convert images", "error", err, "job_id", job.ID)
} else {
multimediaItems = append(multimediaItems, images...)
}
}
// Convert videos
if len(job.Videos) > 0 {
videos, err := s.convertToMultimediaContent(job.Videos, JobVideoType)
if err != nil {
xlog.Warn("Failed to convert videos", "error", err, "job_id", job.ID)
} else {
multimediaItems = append(multimediaItems, videos...)
}
}
// Convert audios
if len(job.Audios) > 0 {
audios, err := s.convertToMultimediaContent(job.Audios, JobAudioType)
if err != nil {
xlog.Warn("Failed to convert audios", "error", err, "job_id", job.ID)
} else {
multimediaItems = append(multimediaItems, audios...)
}
}
// Convert files
if len(job.Files) > 0 {
files, err := s.convertToMultimediaContent(job.Files, JobFileType)
if err != nil {
xlog.Warn("Failed to convert files", "error", err, "job_id", job.ID)
} else {
multimediaItems = append(multimediaItems, files...)
}
}
fragment = fragment.AddMessage("user", prompt, multimediaItems...)
// Get API address and key
_, port, err := net.SplitHostPort(s.appConfig.APIAddress)
if err != nil {
job.Status = schema.JobStatusFailed
job.Error = fmt.Sprintf("failed to parse API address: %v", err)
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("failed to parse API address: %w", err)
}
apiKey := ""
if len(s.appConfig.ApiKeys) > 0 {
apiKey = s.appConfig.ApiKeys[0]
}
// Create LLM client
defaultLLM := clients.NewLocalAILLM(modelConfig.Name, apiKey, "http://127.0.0.1:"+port)
// Initialize traces slice
job.Traces = []schema.JobTrace{}
tc := &traceCollector{
jobID: job.ID,
traces: job.Traces,
jobs: s.jobs,
}
// Build cogito options
cogitoOpts := modelConfig.BuildCogitoOptions()
cogitoOpts = append(
cogitoOpts,
cogito.WithContext(ctx),
cogito.WithMCPs(sessions...),
cogito.WithStatusCallback(func(status string) {
xlog.Debug("Status", "job_id", job.ID, "model", modelConfig.Name, "status", status)
// Store trace
trace := schema.JobTrace{
Type: "status",
Content: status,
Timestamp: time.Now(),
}
tc.add(trace)
}),
cogito.WithReasoningCallback(func(reasoning string) {
xlog.Debug("Reasoning", "job_id", job.ID, "model", modelConfig.Name, "reasoning", reasoning)
// Store trace
trace := schema.JobTrace{
Type: "reasoning",
Content: reasoning,
Timestamp: time.Now(),
}
tc.add(trace)
}),
cogito.WithToolCallBack(func(t *cogito.ToolChoice, state *cogito.SessionState) cogito.ToolCallDecision {
xlog.Debug("Tool call", "job_id", job.ID, "model", modelConfig.Name, "tool", t.Name, "reasoning", t.Reasoning, "arguments", t.Arguments)
// Store trace
arguments := make(map[string]any)
if t.Arguments != nil {
arguments = t.Arguments
}
trace := schema.JobTrace{
Type: "tool_call",
Content: t.Reasoning,
Timestamp: time.Now(),
ToolName: t.Name,
Arguments: arguments,
}
tc.add(trace)
return cogito.ToolCallDecision{
Approved: true,
}
}),
cogito.WithToolCallResultCallback(func(t cogito.ToolStatus) {
xlog.Debug("Tool call result", "job_id", job.ID, "model", modelConfig.Name, "tool", t.Name, "result", t.Result, "tool_arguments", t.ToolArguments)
// Store trace
arguments := make(map[string]any)
// Convert ToolArguments to map via JSON marshaling
if toolArgsBytes, err := json.Marshal(t.ToolArguments); err == nil {
var toolArgsMap map[string]any
if err := json.Unmarshal(toolArgsBytes, &toolArgsMap); err == nil {
arguments = toolArgsMap
}
}
arguments["result"] = t.Result
trace := schema.JobTrace{
Type: "tool_result",
Content: t.Result,
Timestamp: time.Now(),
ToolName: t.Name,
Arguments: arguments,
}
tc.add(trace)
}),
cogito.WithStreamCallback(func(ev cogito.StreamEvent) {
switch ev.Type {
case cogito.StreamEventReasoning:
trace := schema.JobTrace{
Type: "stream_reasoning",
Content: ev.Content,
Timestamp: time.Now(),
}
tc.add(trace)
case cogito.StreamEventContent:
trace := schema.JobTrace{
Type: "stream_content",
Content: ev.Content,
Timestamp: time.Now(),
}
tc.add(trace)
case cogito.StreamEventToolCall:
trace := schema.JobTrace{
Type: "stream_tool_call",
Content: ev.ToolArgs,
ToolName: ev.ToolName,
Timestamp: time.Now(),
}
tc.add(trace)
}
}),
)
// Execute tools
f, err := cogito.ExecuteTools(defaultLLM, fragment, cogitoOpts...)
// Re-read job from the store to pick up traces written by callbacks
if updated := s.jobs.Get(job.ID); updated.ID != "" {
job = updated
}
if err != nil && !errors.Is(err, cogito.ErrNoToolSelected) {
job.Status = schema.JobStatusFailed
job.Error = fmt.Sprintf("failed to execute tools: %v", err)
completedAt := time.Now()
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
return fmt.Errorf("failed to execute tools: %w", err)
}
// Extract traces from fragment.Status after execution
// This provides complete information about tool calls and results
// We use Status data to supplement/replace callback data for completeness
if f.Status != nil {
// Clear existing tool_call and tool_result traces (from callbacks) and replace with Status data
// Keep status and reasoning traces from callbacks
filteredTraces := []schema.JobTrace{}
for _, trace := range job.Traces {
if trace.Type != "tool_call" && trace.Type != "tool_result" {
filteredTraces = append(filteredTraces, trace)
}
}
job.Traces = filteredTraces
// Extract tool calls from Status.ToolsCalled
if len(f.Status.ToolsCalled) > 0 {
for _, toolCallInterface := range f.Status.ToolsCalled {
// Marshal to JSON and unmarshal to extract fields
if toolCallBytes, err := json.Marshal(toolCallInterface); err == nil {
var toolCallData map[string]any
if err := json.Unmarshal(toolCallBytes, &toolCallData); err == nil {
arguments := make(map[string]any)
if args, ok := toolCallData["arguments"].(map[string]any); ok {
arguments = args
}
reasoning := ""
if r, ok := toolCallData["reasoning"].(string); ok {
reasoning = r
}
name := ""
if n, ok := toolCallData["name"].(string); ok {
name = n
}
trace := schema.JobTrace{
Type: "tool_call",
Content: reasoning,
Timestamp: time.Now(),
ToolName: name,
Arguments: arguments,
}
job.Traces = append(job.Traces, trace)
}
}
}
}
// Extract tool results from Status.ToolResults
if len(f.Status.ToolResults) > 0 {
for _, toolResult := range f.Status.ToolResults {
arguments := make(map[string]any)
// Convert ToolArguments to map via JSON marshaling
if toolArgsBytes, err := json.Marshal(toolResult.ToolArguments); err == nil {
var toolArgsMap map[string]any
if err := json.Unmarshal(toolArgsBytes, &toolArgsMap); err == nil {
arguments = toolArgsMap
}
}
arguments["result"] = toolResult.Result
trace := schema.JobTrace{
Type: "tool_result",
Content: toolResult.Result,
Timestamp: time.Now(),
ToolName: toolResult.Name,
Arguments: arguments,
}
job.Traces = append(job.Traces, trace)
}
}
}
// Update job with result
completedAt := time.Now()
job.Status = schema.JobStatusCompleted
job.Result = f.LastMessage().Content
job.CompletedAt = &completedAt
s.jobs.Set(job.ID, job)
xlog.Info("Job completed", "job_id", job.ID, "status", job.Status)
go func() { s.saveJobs(job) }()
// Send webhooks (non-blocking)
go func() {
s.sendWebhooks(job, task)
}()
return nil
}
// worker processes jobs from the queue
func (s *AgentJobService) worker(ctx context.Context) {
for {
select {
case <-ctx.Done():
return
case exec := <-s.jobQueue:
// Check if job was cancelled before execution
select {
case <-exec.Ctx.Done():
job := exec.Job
now := time.Now()
job.Status = schema.JobStatusCancelled
job.CompletedAt = &now
s.jobs.Set(job.ID, job)
s.cancellations.Delete(job.ID)
continue
default:
}
// Execute job
err := s.ExecuteJobInternal(exec.Job, exec.Task, exec.Ctx)
if err != nil {
xlog.Error("Job execution failed", "error", err, "job_id", exec.Job.ID)
}
// Clean up cancellation
s.cancellations.Delete(exec.Job.ID)
}
}
}
// ScheduleCronTask schedules a task to run on a cron schedule
func (s *AgentJobService) ScheduleCronTask(task schema.Task) error {
if task.Cron == "" {
return nil // No cron expression
}
// Parse cron expression (support standard 5-field format)
// Convert to 6-field format if needed (with seconds)
cronExpr := task.Cron
// Use cron parameters if provided, otherwise use empty map
cronParams := task.CronParameters
if cronParams == nil {
cronParams = map[string]string{}
}
entryID, err := s.cronScheduler.AddFunc(cronExpr, func() {
// Create job for cron execution with configured parameters
// Multimedia will be fetched from task sources in ExecuteJob
_, err := s.ExecuteJob(task.ID, cronParams, "cron", nil)
if err != nil {
xlog.Error("Failed to execute cron job", "error", err, "task_id", task.ID)
}
})
if err != nil {
return fmt.Errorf("failed to parse cron expression: %w", err)
}
s.cronEntries.Set(task.ID, entryID)
xlog.Info("Scheduled cron task", "task_id", task.ID, "cron", cronExpr)
return nil
}
// UnscheduleCronTask removes a task from the cron scheduler
func (s *AgentJobService) UnscheduleCronTask(taskID string) {
if s.cronEntries.Exists(taskID) {
entryID := s.cronEntries.Get(taskID)
s.cronScheduler.Remove(entryID)
s.cronEntries.Delete(taskID)
xlog.Info("Unscheduled cron task", "task_id", taskID)
}
}
// sendWebhooks sends webhook notifications to all configured webhooks
func (s *AgentJobService) sendWebhooks(job schema.Job, task schema.Task) {
// Collect all webhook configs from new format
webhookConfigs := task.Webhooks
if len(webhookConfigs) == 0 {
return // No webhooks configured
}
xlog.Info("Sending webhooks", "job_id", job.ID, "webhook_count", len(webhookConfigs))
// Send all webhooks concurrently and track results
var wg sync.WaitGroup
errors := make(chan webhookError, len(webhookConfigs))
var successCount atomic.Int32
for _, webhookConfig := range webhookConfigs {
wg.Go(func() {
if err := s.sendWebhook(job, task, webhookConfig); err != nil {
errors <- webhookError{
URL: webhookConfig.URL,
Error: err.Error(),
}
} else {
successCount.Add(1)
}
})
}
wg.Wait()
close(errors)
// Collect errors
var webhookErrors []string
for err := range errors {
webhookErrors = append(webhookErrors, fmt.Sprintf("%s: %s", err.URL, err.Error))
}
// Update job with webhook status
job = s.jobs.Get(job.ID)
if job.ID == "" {
return
}
now := time.Now()
if len(webhookErrors) == 0 {
// All webhooks succeeded
job.WebhookSent = true
job.WebhookSentAt = &now
job.WebhookError = ""
} else if successCount.Load() > 0 {
// Some succeeded, some failed
job.WebhookSent = true
job.WebhookSentAt = &now
job.WebhookError = fmt.Sprintf("Some webhooks failed (%d/%d succeeded): %s", int(successCount.Load()), len(webhookConfigs), strings.Join(webhookErrors, "; "))
} else {
// All failed
job.WebhookSent = false
job.WebhookError = fmt.Sprintf("All webhooks failed: %s", strings.Join(webhookErrors, "; "))
}
s.jobs.Set(job.ID, job)
go func() { s.saveJobs(job) }()
}
// webhookError represents a webhook delivery error
type webhookError struct {
URL string
Error string
}
// sendWebhook sends a single webhook notification
// Returns an error if the webhook delivery failed
func (s *AgentJobService) sendWebhook(job schema.Job, task schema.Task, webhookConfig schema.WebhookConfig) error {
// Build payload
payload, err := s.buildWebhookPayload(job, task, webhookConfig)
if err != nil {
xlog.Error("Failed to build webhook payload", "error", err, "job_id", job.ID, "webhook_url", webhookConfig.URL)
return fmt.Errorf("failed to build payload: %w", err)
}
xlog.Debug("Sending webhook", "job_id", job.ID, "webhook_url", webhookConfig.URL, "payload", string(payload))
// Determine HTTP method (default to POST)
method := webhookConfig.Method
if method == "" {
method = "POST"
}
// Create HTTP request
req, err := http.NewRequest(method, webhookConfig.URL, bytes.NewBuffer(payload))
if err != nil {
xlog.Error("Failed to create webhook request", "error", err, "job_id", job.ID, "webhook_url", webhookConfig.URL)
return fmt.Errorf("failed to create request: %w", err)
}
// Set headers
req.Header.Set("Content-Type", "application/json")
for key, value := range webhookConfig.Headers {
req.Header.Set(key, value)
}
// Execute with retry
client := httpclient.NewWithTimeout(30 * time.Second)
err = s.executeWithRetry(client, req)
if err != nil {
xlog.Error("Webhook delivery failed", "error", err, "job_id", job.ID, "webhook_url", webhookConfig.URL)
return fmt.Errorf("webhook delivery failed: %w", err)
}
xlog.Info("Webhook delivered successfully", "job_id", job.ID, "webhook_url", webhookConfig.URL)
return nil
}
// buildWebhookPayload builds webhook payload (default or template)
func (s *AgentJobService) buildWebhookPayload(job schema.Job, task schema.Task, webhookConfig schema.WebhookConfig) ([]byte, error) {
if webhookConfig.PayloadTemplate != "" {
// Use custom template
return s.buildPayloadFromTemplate(job, task, webhookConfig.PayloadTemplate)
}
// Use default format
// Include Error field (empty string if no error)
payload := map[string]any{
"job_id": job.ID,
"task_id": job.TaskID,
"task_name": task.Name,
"status": string(job.Status),
"result": job.Result,
"error": job.Error, // Empty string if no error
"parameters": job.Parameters,
"started_at": job.StartedAt,
"completed_at": job.CompletedAt,
}
return json.Marshal(payload)
}
// buildPayloadFromTemplate builds payload from template
func (s *AgentJobService) buildPayloadFromTemplate(job schema.Job, task schema.Task, templateStr string) ([]byte, error) {
// Create template context
// Available variables:
// - .Job - Job object with all fields
// - .Task - Task object
// - .Result - Job result (if successful)
// - .Error - Error message (if failed, empty string if successful)
// - .Status - Job status string
ctx := map[string]any{
"Job": job,
"Task": task,
"Result": job.Result,
"Error": job.Error,
"Parameters": job.Parameters,
"Status": string(job.Status),
}
// Add json function for template
funcMap := template.FuncMap{
"json": func(v any) string {
b, _ := json.Marshal(v)
return string(b)
},
}
tmpl, err := template.New("payload").Funcs(funcMap).Funcs(sprig.FuncMap()).Parse(templateStr)
if err != nil {
return nil, err
}
var buf bytes.Buffer
if err := tmpl.Execute(&buf, ctx); err != nil {
return nil, err
}
return buf.Bytes(), nil
}
// executeWithRetry executes HTTP request with retry logic
func (s *AgentJobService) executeWithRetry(client *http.Client, req *http.Request) error {
maxRetries := 3
backoff := []time.Duration{1 * time.Second, 2 * time.Second, 4 * time.Second}
// Save body bytes once before the loop so retries can re-use them
var savedBody []byte
if req.Body != nil {
var err error
savedBody, err = io.ReadAll(req.Body)
req.Body.Close()
if err != nil {
return fmt.Errorf("reading request body: %w", err)
}
}
var lastErr error
for i := range maxRetries {
// Reset body for each attempt
if savedBody != nil {
req.Body = io.NopCloser(bytes.NewReader(savedBody))
req.ContentLength = int64(len(savedBody))
}
var resp *http.Response
resp, lastErr = client.Do(req)
if lastErr == nil && resp.StatusCode >= 200 && resp.StatusCode < 300 {
resp.Body.Close()
return nil
}
if resp != nil {
resp.Body.Close()
}
if i < maxRetries-1 {
time.Sleep(backoff[i])
}
}
return fmt.Errorf("failed after %d retries: %w", maxRetries, lastErr)
}
// CleanupOldJobs removes jobs older than retention period
func (s *AgentJobService) CleanupOldJobs() error {
cutoff := time.Now().AddDate(0, 0, -s.retentionDays)
allJobs := s.jobs.Values()
removed := 0
for _, job := range allJobs {
if job.CreatedAt.Before(cutoff) {
s.jobs.Delete(job.ID)
removed++
}
}
// Also clean up from the persister (DB cleans via SQL, file re-serializes)
retention := time.Duration(s.retentionDays) * 24 * time.Hour
if dbRemoved, err := s.persister.CleanupOldJobs(retention); err != nil {
xlog.Warn("Failed to cleanup old jobs from persister", "error", err)
} else if dbRemoved > 0 {
removed += int(dbRemoved)
}
if removed > 0 {
xlog.Info("Cleaned up old jobs", "removed", removed, "retention_days", s.retentionDays)
// For file persister, re-serialize the map after deletions
s.saveJobs(schema.Job{})
}
return nil
}
// Start starts the background service
func (s *AgentJobService) Start(ctx context.Context) error {
// Create service context
s.ctx, s.cancel = context.WithCancel(ctx)
// Update retention days from config
s.retentionDays = cmp.Or(s.appConfig.AgentJobRetentionDays, 30)
// Load tasks and jobs from persister (DB or files)
s.loadFromPersister()
// In distributed mode, the Dispatcher handles worker pool + cron leader election.
// Skip local worker pool and cron scheduler.
if s.dispatcher != nil {
xlog.Info("AgentJobService started (distributed mode — local workers/cron disabled)", "retention_days", s.retentionDays)
return nil
}
// Start cron scheduler
s.cronScheduler.Start()
// Start worker pool (5 workers)
workerCount := 5
for range workerCount {
go s.worker(s.ctx)
}
// Schedule daily cleanup at midnight
_, err := s.cronScheduler.AddFunc("0 0 * * *", func() {
if err := s.CleanupOldJobs(); err != nil {
xlog.Error("Failed to cleanup old jobs", "error", err)
}
})
if err != nil {
xlog.Warn("Failed to schedule daily cleanup", "error", err)
}
// Run initial cleanup
if err := s.CleanupOldJobs(); err != nil {
xlog.Warn("Failed to run initial cleanup", "error", err)
}
xlog.Info("AgentJobService started", "retention_days", s.retentionDays)
return nil
}
// Stop stops the agent job service
func (s *AgentJobService) Stop() error {
if s.cancel != nil {
s.cancel()
s.cancel = nil
}
if s.cronScheduler != nil {
s.cronScheduler.Stop()
}
// Release the tasks SyncedMap subscription / background workers.
if s.tasks != nil {
if err := s.tasks.Close(); err != nil {
xlog.Warn("Error closing tasks sync map", "error", err)
}
}
xlog.Info("AgentJobService stopped")
return nil
}
// UpdateRetentionDays updates the retention days setting
func (s *AgentJobService) UpdateRetentionDays(days int) {
s.retentionDays = cmp.Or(days, 30)
xlog.Info("Updated agent job retention days", "retention_days", s.retentionDays)
}