Files
LocalAI/core/services/agentpool/agent_jobs.go
T
mudler-agentandEttore Di Giacinto e4fa051ee3 refactor(distributed): put the NATS-only paths behind interfaces (#12395)
* feat(messaging): add shared subject rules

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(messaging): cover BroadcastRoots, ControlRoots and SubjectRoot

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(messaging): add Broadcaster and enforce subject rules in every carrier

Broadcaster is the fan-out half of MessagingClient. The NATS client and
the in-memory FakeBus now refuse a subject outside the served roots and any
wildcard other than a whole single token, and FakeBus shares MatchSubject
instead of its own copy. FakeBus Unsubscribe now removes its own
subscription instead of the first one with the same subject.

A shared conformance suite in messagingtest runs against both carriers.
The distributed e2e specs that used invented test.* subjects, and the one
that subscribed with a > filter, now use subjects from subjects.go.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: depend on Broadcaster where only publish and subscribe are used

Narrowed to messaging.Broadcaster: nodes/staging_progress.go,
nodes/install_progress_publisher.go, galleryop/operation.go,
galleryop/service.go, agentpool/user_services.go, agentpool/agent_jobs.go,
openresponses/store.go, openresponses/sync.go, syncstate/syncstate.go,
finetune/service.go, quantization/service.go and
failover/distsync/distsync.go. SubscribeJSON now takes a Broadcaster
because it only calls Subscribe, which lets the narrowed consumers use it.

Stayed wide: worker/supervisor.go, because its client field also serves
the SubscribeReply handlers in worker/lifecycle.go. The request/reply,
queue and wiring files (nodes/unloader.go, nodes/file_stager_s3.go,
jobs/dispatcher.go, agents/dispatcher.go, agents/events.go,
worker/file_staging.go, cli/agent_worker.go, http/app.go) are unchanged by
design.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): name the no-route condition and confine the carrier error

Consumers matched nats.ErrNoResponders, which names an absence, to demote a node. They now match ErrNoRoute, the control path maps the carrier's failure onto it, and timeouts and worker refusals are pinned as not being no-route.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs(nodes): state which FileStager implementations return ErrNoRoute

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): build backend clients through one node-aware seam

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: describe the distributed transport seams

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: correct comments that overclaim after the seams refactor

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(agent-worker): refuse an unserved LOCALAI_AGENT_SUBJECT at startup

The messaging client now refuses a subject whose root no carrier serves.
An agent worker started with a custom LOCALAI_AGENT_SUBJECT such as
tenant-a.agent.execute used to start and then wait on a subject the
frontend never publishes to. After the subject rules landed it exited
at subscribe time with an error that did not name the setting.

Behaviour change: the worker now checks LOCALAI_AGENT_SUBJECT before it
registers or connects, and exits with an error that names the variable
and says to use a served subject under the agent root, for example
agent.execute. The served roots are not widened: a custom root was
never delivered by the frontend, and a wider set would reopen the
drift the subject rules exist to close. The flag help and the agent
worker docs state the constraint.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(nodes): pin the reactions to ErrNoRoute

Three callers react to ErrNoRoute and had no spec: the reconciler's
upgrade drain falls back to the legacy forced install, the reconciler
marks the node unhealthy when a pending op has no route, and the
backend-op fan-out marks the node unhealthy. Each spec drives the real
caller with a scripted no-responders reply and reads the result from
the registry or the recorded requests.

A fourth spec pins the other side: a pending op that times out leaves
the node healthy and only counts the attempt, so mapping timeouts onto
ErrNoRoute would fail here.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(messaging): pin client subject checks and fail the carrier suite in CI

Add specs that call Publish, Request, Subscribe, QueueSubscribe,
SubscribeReply and QueueSubscribeReply on a client with no connection.
Each call must return ErrUnservedSubject for bogus.thing and
ErrUnsupportedWildcard for jobs.>. This proves that the subject check
runs before the connection is used, and needs no server.

The NATS conformance suite is the only check that runs the subject rules
against a real carrier. Before this change it skipped without output
when Docker was missing. Now it fails when CI is set, so a Linux runner
without Docker cannot hide it. It still skips on local runs and on macOS
CI, which has no Docker.

Add SubjectNodeBackendInstallProgress to the list of constructors that
must build served subjects, and ask contributors to extend the list.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: state what ErrNoRoute may change, and group the distributed guides

The seams note said MarkUnhealthy was the only state change allowed
on ErrNoRoute. A pending backend op still records the failed attempt,
counts toward the reconciler's retry limit and is dead-lettered after
the maximum attempts. The note now says that MarkUnhealthy is the only
change to the node's own state, and that the per-op accounting is not
a verdict about the node.

The note also documents that the NATS conformance run fails under CI
when Docker is missing. The distributed-seams row moves next to the
distributed-state row in the topics table. The liveness ping spec
header now says no route is a reason to skip the worker, not proof
that the worker is gone.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): give the backend client factory the node id

Mechanical: the method gains a nodeID parameter and the eight test fakes are updated. No behaviour change.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): drop the optional node-aware factory

The node id is now in the main method, so the optional interface and its helper had no behaviour of their own.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): dial backend probes through the client factory

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): dial workers' file servers through a per-node dialer

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(http): proxy backend logs through the per-node worker dialer

The admin backend-logs proxy (list, lines and the WebSocket stream) now reaches a worker through the same per-node dialer as the HTTP file stager, so every frontend-to-worker dial goes through one seam. The shared direct dialer keeps alive for 15s where the proxy used 30s. Harmless for requests bounded at 15s.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(http): keep the backend-logs proxy independent of the admin connection

The proxy request had no context before the dialer change and is bounded only by its 15s timeout. Keep it that way.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: move the worker control payloads to workerctl

Mechanical move of the request and reply structs, the install progress event and the file payloads out of messaging. The verbs no longer belong to one carrier. No alias is left behind.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(worker): serve the lifecycle verbs through a controlServer

The worker registers one handler per verb and a NATS server maps each verb to its subject. Registration errors now name the verb. node.stop is served with SubscribeReply, which is identical on the wire because the handler never replies.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(worker): report install progress through the control sink

Install and upgrade now emit download progress through the sink the control server hands them. The debounce and the terminal flush stay in the handler path, built over that sink by the new nodes.NewDebouncedInstallProgressSink, which replaces NewDebouncedInstallProgressPublisher. The subject and payload on the wire are unchanged. The supervisor no longer holds the bus, and installFn and upgradeFn let specs drive both verbs without a gallery.

The malformed-request log lines are restored for install, upgrade, backend.delete, model.unload, model.stop and model.delete, with the reply bytes unchanged. The signal adapter is renamed noReply, which also lets worker.go import os/signal without an alias again.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(worker): serve the file-staging verbs through a controlServer

An empty list-dir answer is now {} rather than {"files":null}, because the typed reply omits an empty Files slice. The frontend decodes both to a nil slice in nodes/file_stager_s3.go.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(messaging): add WorkQueue and the NATS producer

This is the producer side of the competing-consumer seam. The work kinds map one to one to today's subjects and queue groups: task to jobs.new and mcp-ci to jobs.mcp-ci.new (both in group workers), agent-run to agent.execute (group agent-workers). Enqueue publishes the payload as Publish does today, with one JSON marshal. FakeBus now records queue groups and keeps reply handlers so later specs can pin and drive them.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(messaging): add the NATS WorkConsumer

An in-flight limit of one runs the handler inline on the delivery goroutine, as the MCP CI consumer does today. Any other limit spawns per delivery, as the agent consumer does. Queue groups are unchanged.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: publish queued work through WorkQueue

The job dispatcher, the agent pool and the agent scheduler enqueue through messaging.WorkQueue; the NATS implementation publishes to the same subjects as before. DistributedServices builds the queue next to the NATS client and hands it to the dispatcher and the agent pool, whose distributed mode switch now reads a non-nil WorkQueue. The unused AgentPoolService.SetNATSClient is removed.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: consume queued work through WorkConsumer

The agent dispatcher and the MCP CI consumer register through messaging.WorkConsumer. The NATS implementation keeps the inline one-at-a-time model for MCP CI and the per-delivery model for agent runs. handleMCPCIJob reports on the events publisher the carrier hands it instead of a captured client.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: delete the consumers nothing in production reached

jobs.new has a producer and no production consumer, and the agent dispatcher's Dispatch was only called from tests. Publishing jobs.new is unchanged.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(mcp): send MCP requests to agent workers through AgentControl

Timeouts still honour only the deadline, not cancellation, exactly as today. The NATS no-responders error maps to ErrNoRoute and a timeout does not.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(agent-worker): serve MCP requests and backend.stop through agentRPCServer

The agent worker's MCP tool and discovery reply subscriptions and its backend stop listener move behind an unexported agentRPCServer interface, served on NATS by nodes.NATSAgentRPCServer. The handlers become typed mcp.ToolHandler and mcp.DiscoveryHandler values that answer every failure with a reply carrying Error.

Queue group (agent-workers), inline execution on the delivery goroutine, the background handler context, the unmarshal error reply texts and the reply-less backend stop subscription are unchanged. The backend stop handler takes the decoded backend name, so it can still close that backend's MCP sessions.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(messaging): remove helpers that only tests used

BroadcastRoots, ControlRoots and SubjectRoot had no production caller. The roots spec now asserts every served root through ValidateSubject instead. MatchSubject moves back into the test support package, the only place that used it, with its table. NATSAgentRPCServer drops the subscription list it stored and never read, and NewNATSAgentRPCServer gets a doc comment.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(mcp): round trip the agent RPC server over a real NATS server

One spec sends a tool request and a discovery request through NATSAgentControl to NATSAgentRPCServer and checks that the handlers see the decoded requests and the replies come back. It also puts an undecodable body on the tool subject and checks the server answers with an unmarshal error instead of leaving the requester to time out.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: describe the distributed transport seams

The developer note now lists the final seams: fan-out, queues, both halves of the control verbs and of agent RPC, and the dial. It records the open items a second carrier has to handle.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test: pin the in-flight limit each queue consumer asks for

The work queue specs pin what Consume does for a given limit, but nothing
pinned which limit each production consumer passes. Changing the agent
worker's MCP CI limit from 1 to 0 would have let MCP CI jobs run
concurrently on each worker with every test green.

Move the MCP CI Consume call into startMCPCIConsumer with the same wiring
and pin that it asks for (WorkMCPCI, 1). Pin that NATSDispatcher.Start asks
for (WorkAgentRun, maxConcurrent) for several limits.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: remove helpers the branch left without a caller

SubjectJobCancelWildcard lost its last subscriber when the frontend
stopped listening on jobs.*.cancel; the NATS permissions and conformance
suite spell the subject out, so nothing reads the constant.

decodeBackendStopRequest returned a stopAll flag that production dropped
and only a test read. decodeBackendStop is now the single decoder with the
same semantics: an empty body is stop-all, an empty Backend is stop-all,
malformed JSON is an error. stopBackends still derives stop-all from
Backend, so no reply changes.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(messaging): keep an explicitly empty agent queue a plain subscription

Before the work queue seam the agent worker passed LOCALAI_AGENT_QUEUE
straight to QueueSubscribe, so an explicitly empty value made a plain
subscription and every agent worker ran every agent run.
WithAgentRunRoute replaced an empty queue with agent-workers, which
silently changed that.

Keep the queue as given once the option is applied. An empty subject still
falls back to agent.execute, since it never had a meaning of its own. The
flag default stays agent-workers, so only an explicitly empty value
reaches this.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: correct comments and record the PR B notes

Fix the recordingFactory comment (it also records the parallel flag),
document that a negative maxInFlight is unbounded and that Unsubscribe from
a handler deadlocks, and say a permanently undecodable payload returns nil.

Record controlHandler's undecodable return as a kept exception, and add the
second carrier notes to the developer note: the reconciler has no
ClientFactory option, the logs proxy honours HTTP_PROXY, verbs one carrier
serves need an opt-out, terminal replies come from the result event, and
agent runs publish through the NATS-bound EventBridge, which is not an
additive change.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
2026-10-02 23:58:41 +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.Broadcaster
// 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.Broadcaster) {
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)
}