package worker import ( "context" "encoding/json" "fmt" "log/slog" "strings" "time" "git.toki-labs.com/toki/gito/services/core/internal/config" "git.toki-labs.com/toki/gito/services/core/internal/core" "git.toki-labs.com/toki/gito/services/core/internal/events" "git.toki-labs.com/toki/gito/services/core/internal/gitengine" "git.toki-labs.com/toki/gito/services/core/internal/storage" ) type OperationPicker interface { PickQueuedOperation(ctx context.Context, now time.Time) (core.Operation, bool, error) } type Clock interface { Now() time.Time } type realClock struct{} func (realClock) Now() time.Time { return time.Now() } type AgentRunInput struct { OperationID string RepoID string Branch string WorkspacePath string Instruction string PolicyContext map[string]any ExpectedRevision string CredentialRefs map[string]string } type AgentRunState string const ( AgentRunStateSucceeded AgentRunState = "succeeded" AgentRunStateFailed AgentRunState = "failed" AgentRunStateCancelled AgentRunState = "cancelled" ) type AgentRunResult struct { State AgentRunState Message string } type AgentRunGitEvidence struct { StatusClean bool StatusOutput string BeforeRevision string AfterRevision string ChangedFiles []gitengine.ChangedFile } type AgentRunInputLoader interface { LoadAgentRunInput(ctx context.Context, op core.Operation) (AgentRunInput, error) } type AgentRunInvoker interface { InvokeAgentRun(ctx context.Context, input AgentRunInput) (AgentRunResult, error) } type RunnerDependencies struct { OperationStore storage.OperationStore OperationEventStore storage.OperationEventStore AgentRunInputLoader AgentRunInputLoader AgentRunInvoker AgentRunInvoker GitRunner gitengine.CommandRunner } type Runner struct { cfg config.Config logger *slog.Logger picker OperationPicker clock Clock operationStore storage.OperationStore operationEventStore storage.OperationEventStore agentRunInputLoader AgentRunInputLoader agentRunInvoker AgentRunInvoker gitRunner gitengine.CommandRunner } func NewRunner(cfg config.Config, logger *slog.Logger, picker OperationPicker, clock Clock) *Runner { return NewRunnerWithDependencies(cfg, logger, picker, clock, RunnerDependencies{}) } func NewRunnerWithDependencies(cfg config.Config, logger *slog.Logger, picker OperationPicker, clock Clock, deps RunnerDependencies) *Runner { if clock == nil { clock = realClock{} } if logger == nil { logger = slog.Default() } gitRunner := deps.GitRunner if gitRunner == nil { gitRunner = gitengine.CLI{} } return &Runner{ cfg: cfg, logger: logger, picker: picker, clock: clock, operationStore: deps.OperationStore, operationEventStore: deps.OperationEventStore, agentRunInputLoader: deps.AgentRunInputLoader, agentRunInvoker: deps.AgentRunInvoker, gitRunner: gitRunner, } } func (r *Runner) Run() error { return r.RunOnce(context.Background()) } func (r *Runner) RunOnce(ctx context.Context) error { if !r.cfg.WorkerEnabled { r.logger.Info("worker disabled") return nil } if r.picker == nil { r.logger.Warn("worker picker not configured") return nil } op, ok, err := r.picker.PickQueuedOperation(ctx, r.clock.Now()) if err != nil { r.logger.Error("failed to pick operation", "error", err) return fmt.Errorf("pick operation: %w", err) } if !ok { r.logger.Debug("no pending operation found") return nil } r.logger.Info("picked operation", "operation_id", op.ID, "type", op.Type) if op.Type == core.OperationAgentRun { return r.runAgentRunOperation(ctx, op) } return nil } func (r *Runner) runAgentRunOperation(ctx context.Context, op core.Operation) error { if r.operationStore == nil || r.operationEventStore == nil || r.agentRunInputLoader == nil || r.agentRunInvoker == nil { return fmt.Errorf("agent_run handler is not fully configured") } startedAt := r.clock.Now().UTC() started := storage.OperationEvent{ OperationID: op.ID, Event: events.Event{ ID: eventID(op.ID, events.OperationStarted, startedAt), Type: events.OperationStarted, Subject: operationEventSubject(op.ID), }, } startPayload, err := agentRunEventPayload(op, core.OperationRunning, startedAt, nil) if err != nil { return err } started.Event.Payload = startPayload started.Event.CreatedAt = startedAt if err := r.operationEventStore.AppendEvent(ctx, started); err != nil { return fmt.Errorf("append agent_run started event: %w", err) } agentStarted := storage.OperationEvent{ OperationID: op.ID, Event: events.Event{ ID: eventID(op.ID, events.AgentRunStarted, startedAt), Type: events.AgentRunStarted, Subject: operationEventSubject(op.ID), }, } agentRunStartedPayload, err := agentRunEventPayload(op, core.OperationRunning, startedAt, nil) if err != nil { return err } agentStarted.Event.Payload = agentRunStartedPayload agentStarted.Event.CreatedAt = startedAt if err := r.operationEventStore.AppendEvent(ctx, agentStarted); err != nil { return fmt.Errorf("append agent_run_started event: %w", err) } input, err := r.agentRunInputLoader.LoadAgentRunInput(ctx, op) if err != nil { return fmt.Errorf("load agent_run input: %w", err) } result, err := r.agentRunInvoker.InvokeAgentRun(ctx, input) terminalState := core.OperationFailed completionMessage := strings.TrimSpace(result.Message) // If backend invocation fails, fail the operation with a stable, sanitized message. if err != nil { completionMessage = "agent run invocation failed" } else { terminalState, err = toTerminalOperationState(result.State) if err != nil { return fmt.Errorf("unsupported agent_run result state: %w", err) } } if completionMessage == "" { completionMessage = string(terminalState) } completionExtra := map[string]any{"message": completionMessage} if terminalState == core.OperationSucceeded { evidence, evidenceErr := r.collectAgentRunGitEvidence(input) if evidenceErr != nil { r.logger.Warn("agent_run git evidence collection failed", "operation_id", op.ID, "error", evidenceErr) terminalState = core.OperationFailed completionMessage = "git evidence collection failed" completionExtra = map[string]any{"message": completionMessage} } else { completionExtra = agentRunGitEvidencePayload(evidence) if !evidence.StatusClean { terminalState = core.OperationFailed completionMessage = "workspace dirty after agent run" } completionExtra["message"] = completionMessage } } terminalAt := r.clock.Now().UTC() var updatedOperation core.Operation switch terminalState { case core.OperationSucceeded: updatedOperation, err = r.operationStore.SucceedOperation(ctx, op.ID, terminalAt) case core.OperationFailed: updatedOperation, err = r.operationStore.FailOperation(ctx, op.ID, terminalAt) case core.OperationCancelled: updatedOperation, err = r.operationStore.CancelOperation(ctx, op.ID, terminalAt) default: err = fmt.Errorf("unsupported terminal state %q", terminalState) } if err != nil { return fmt.Errorf("mark agent_run operation terminal: %w", err) } completionEventType := events.OperationCompleted if terminalState == core.OperationFailed { completionEventType = events.OperationFailed } completedAt := storage.OperationEvent{ OperationID: op.ID, Event: events.Event{ ID: eventID(op.ID, completionEventType, terminalAt), Type: completionEventType, Subject: operationEventSubject(op.ID), }, } completedPayload, err := agentRunEventPayload(updatedOperation, terminalState, terminalAt, completionExtra) if err != nil { return err } completedAt.Event.Payload = completedPayload completedAt.Event.CreatedAt = terminalAt if err := r.operationEventStore.AppendEvent(ctx, completedAt); err != nil { return fmt.Errorf("append agent_run completion event: %w", err) } agentCompleted := storage.OperationEvent{ OperationID: op.ID, Event: events.Event{ ID: eventID(op.ID, events.AgentRunCompleted, terminalAt), Type: events.AgentRunCompleted, Subject: operationEventSubject(op.ID), }, } agentCompletedPayload, err := agentRunEventPayload(updatedOperation, terminalState, terminalAt, completionExtra) if err != nil { return err } agentCompleted.Event.Payload = agentCompletedPayload agentCompleted.Event.CreatedAt = terminalAt if err := r.operationEventStore.AppendEvent(ctx, agentCompleted); err != nil { return fmt.Errorf("append agent_run_completed event: %w", err) } return nil } func (r *Runner) collectAgentRunGitEvidence(input AgentRunInput) (AgentRunGitEvidence, error) { workdir := strings.TrimSpace(input.WorkspacePath) before := strings.TrimSpace(input.ExpectedRevision) if workdir == "" || before == "" { return AgentRunGitEvidence{}, fmt.Errorf("workspace_path and expected_revision are required") } status, err := gitengine.Status(r.gitRunner, workdir) if err != nil { return AgentRunGitEvidence{}, fmt.Errorf("git status: %w", err) } after, err := gitengine.HeadRevision(r.gitRunner, workdir) if err != nil { return AgentRunGitEvidence{}, fmt.Errorf("git head revision: %w", err) } changedFiles, err := gitengine.ChangedFilesWithStatus(r.gitRunner, workdir, before, after) if err != nil { return AgentRunGitEvidence{}, fmt.Errorf("git changed files: %w", err) } return AgentRunGitEvidence{ StatusClean: status.Clean, StatusOutput: status.Output, BeforeRevision: before, AfterRevision: after, ChangedFiles: changedFiles, }, nil } func agentRunGitEvidencePayload(evidence AgentRunGitEvidence) map[string]any { changedFiles := make([]any, 0, len(evidence.ChangedFiles)) for _, file := range evidence.ChangedFiles { changedFiles = append(changedFiles, map[string]any{ "path": file.Path, "change_type": file.ChangeType, }) } return map[string]any{ "status_clean": evidence.StatusClean, "status_output": evidence.StatusOutput, "before_revision": evidence.BeforeRevision, "after_revision": evidence.AfterRevision, "changed_files": changedFiles, } } func operationEventSubject(operationID string) string { return "operation:" + operationID } func eventID(operationID, eventType string, at time.Time) string { return fmt.Sprintf("agent-run:%s:%s:%d", operationID, eventType, at.UnixNano()) } func agentRunEventPayload(op core.Operation, state core.OperationState, at time.Time, extra map[string]any) ([]byte, error) { payload := map[string]any{ "operation_id": op.ID, "repo_id": op.RepoID, "type": op.Type, "state": state, "at": at.Format(time.RFC3339Nano), } for key, value := range extra { payload[key] = value } encoded, err := json.Marshal(payload) if err != nil { return nil, fmt.Errorf("serialize agent run event payload: %w", err) } return encoded, nil } func toTerminalOperationState(state AgentRunState) (core.OperationState, error) { switch state { case AgentRunStateSucceeded: return core.OperationSucceeded, nil case AgentRunStateFailed: return core.OperationFailed, nil case AgentRunStateCancelled: return core.OperationCancelled, nil default: return "", fmt.Errorf("invalid agent run state %q", state) } }