From 49c1863a61b0f80f42396da17402fbc21dce8e89 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 25 Feb 2026 16:08:54 +0800 Subject: [PATCH] feat: add orchestration ledger and periodic audit supervisor --- config/config.example.json | 28 ++ docs/design/orchestration-audit.md | 112 +++++++ pkg/agent/audit.go | 463 +++++++++++++++++++++++++++++ pkg/agent/audit_test.go | 89 ++++++ pkg/agent/loop.go | 127 +++++++- pkg/agent/orchestration_test.go | 43 +++ pkg/config/config.go | 71 +++-- pkg/config/config_test.go | 68 +++++ pkg/config/defaults.go | 25 ++ pkg/tools/spawn_test.go | 205 +++++++++++++ pkg/tools/subagent.go | 372 ++++++++++++++++++++--- pkg/tools/subagent_tool_test.go | 123 ++++++++ pkg/tools/task_ledger.go | 341 +++++++++++++++++++++ pkg/tools/task_ledger_test.go | 94 ++++++ pkg/tools/toolloop.go | 37 ++- 15 files changed, 2138 insertions(+), 60 deletions(-) create mode 100644 docs/design/orchestration-audit.md create mode 100644 pkg/agent/audit.go create mode 100644 pkg/agent/audit_test.go create mode 100644 pkg/agent/orchestration_test.go create mode 100644 pkg/tools/task_ledger.go create mode 100644 pkg/tools/task_ledger_test.go diff --git a/config/config.example.json b/config/config.example.json index 2d10022b6..35d00babb 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -277,6 +277,34 @@ "enabled": false, "monitor_usb": true }, + "orchestration": { + "enabled": false, + "max_spawn_depth": 3, + "max_parallel_workers": 4, + "max_tasks_per_agent": 20, + "default_task_timeout_seconds": 180, + "retry_limit_per_task": 2 + }, + "audit": { + "enabled": false, + "interval_minutes": 30, + "lookback_minutes": 180, + "min_confidence": 0.75, + "inconsistency_policy": "strict", + "auto_remediation": "safe_only", + "notify_channel": "last_active", + "supervisor": { + "enabled": false, + "model": { + "primary": "gpt-5.2", + "fallbacks": [ + "claude-sonnet-4.6" + ] + }, + "temperature": 0.0, + "max_tokens": 2048 + } + }, "gateway": { "host": "127.0.0.1", "port": 18790 diff --git a/docs/design/orchestration-audit.md b/docs/design/orchestration-audit.md new file mode 100644 index 000000000..a68c6360f --- /dev/null +++ b/docs/design/orchestration-audit.md @@ -0,0 +1,112 @@ +# Orchestration and Audit Runtime Design + +This document describes how PicoClaw's subagent orchestration and periodic audit pipeline work after the `orchestration` and `audit` extensions. + +## Goals + +- Align subagent orchestration semantics with OpenClaw-style multi-level delegation. +- Keep existing behavior backward compatible when features are disabled. +- Provide periodic checks for missed tasks, low-quality outputs, and execution inconsistencies. + +## Config Keys and Runtime Effects + +### `orchestration` + +- `enabled` + - Feature gate for orchestration controls. Existing subagent tools still work; limits are always read from this section. +- `max_spawn_depth` + - Enforced in `SubagentManager.SpawnTask`. + - Nested spawns are detected through `sender_id=subagent:`. + - When depth is exceeded, spawn is rejected with `max spawn depth reached`. +- `max_parallel_workers` + - Enforced as max concurrent running tasks per manager. +- `max_tasks_per_agent` + - Enforced as max active (non-terminal) tasks per manager. +- `default_task_timeout_seconds` + - Used as default deadline metadata in task ledger entries. +- `retry_limit_per_task` + - Used by audit logic to detect failed tasks that still have retry budget. + +### `audit` + +- `enabled` + - Starts a background audit loop in `AgentLoop.Run`. +- `interval_minutes` + - Periodic audit cadence. +- `lookback_minutes` + - Task window scanned in each cycle. +- `min_confidence` + - Threshold for supervisor model score. +- `inconsistency_policy` + - `strict` mode flags completed tasks with no tool evidence. +- `auto_remediation` + - `safe_only` records low-risk remediation actions in ledger. +- `notify_channel` + - Destination for audit report: + - `last_active`: last recorded user channel/chat. + - `channel:chat_id`: explicit destination. + - `channel`: uses last active chat id with an overridden channel. + +### `audit.supervisor` + +- `enabled` + - Enables model-based review in addition to deterministic rule checks. +- `model.primary` / `model.fallbacks` + - Model alias resolved through `model_list`. +- `temperature`, `max_tokens` + - Passed into supervisor model calls. + +## Task Ledger + +`TaskLedger` persists orchestration records under: + +- `/tasks/ledger.json` + +Each entry tracks: + +- task identity and lineage (`id`, `parent_task_id`, `agent_id`) +- routing context (`origin_channel`, `origin_chat_id`) +- execution state and result (`status`, `result`, `error`) +- timing (`created_at_ms`, `updated_at_ms`, `deadline_at_ms`) +- evidence and remediation arrays + +## Subagent Lifecycle + +1. `spawn`/`sessions_spawn` creates a task entry (`created` event). +2. Manager resolves execution profile: + - default provider/model/tools, or + - target agent profile when `agent_id` is provided. +3. Tool loop runs and records per-tool traces. +4. Task emits final event (`completed`, `failed`, or `cancelled`). +5. System inbound message is published for main-loop post-processing. + +## Parent/Child and Cascade Semantics + +- Nested spawns capture `parent_task_id`. +- If a parent task fails or is cancelled, manager cancels descendants (task tree) by context cancellation. +- Descendants transition to cancelled when they observe context cancellation. + +## Audit Rules + +Deterministic checks: + +- `missed` + - planned task overdue + - running task timeout + - failed task with retry budget +- `quality` + - completed task with empty result +- `inconsistency` + - completed task with zero evidence in `strict` mode + +Optional model checks: + +- Supervisor model receives task JSON and returns structured score/issues. +- Findings are merged into deterministic report. + +## Backward Compatibility + +- Existing tools and loop behavior remain unchanged when `audit.enabled=false`. +- New fields are additive and optional. +- `spawn` and `subagent` retain previous parameter contract; `agent_id` is additive for `subagent`. + diff --git a/pkg/agent/audit.go b/pkg/agent/audit.go new file mode 100644 index 000000000..bec87083a --- /dev/null +++ b/pkg/agent/audit.go @@ -0,0 +1,463 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/constants" + "github.com/sipeed/picoclaw/pkg/logger" + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/tools" +) + +type AuditFinding struct { + TaskID string `json:"task_id"` + Category string `json:"category"` + Severity string `json:"severity"` + Message string `json:"message"` + Recommendation string `json:"recommendation,omitempty"` +} + +type AuditReport struct { + GeneratedAt time.Time `json:"generated_at"` + Lookback time.Duration `json:"lookback"` + TotalTasks int `json:"total_tasks"` + Findings []AuditFinding `json:"findings"` +} + +func (r *AuditReport) FormatMessage() string { + if r == nil { + return "Task audit report unavailable." + } + + var b strings.Builder + b.WriteString("Task Audit Report\n") + b.WriteString(fmt.Sprintf("Generated: %s\n", r.GeneratedAt.Format("2006-01-02 15:04:05"))) + b.WriteString(fmt.Sprintf("Lookback: %dm\n", int(r.Lookback.Minutes()))) + b.WriteString(fmt.Sprintf("Tasks scanned: %d\n", r.TotalTasks)) + b.WriteString(fmt.Sprintf("Findings: %d\n", len(r.Findings))) + + if len(r.Findings) == 0 { + b.WriteString("\nNo issues detected.") + return b.String() + } + + limit := len(r.Findings) + if limit > 10 { + limit = 10 + } + b.WriteString("\nTop findings:\n") + for i := 0; i < limit; i++ { + f := r.Findings[i] + b.WriteString(fmt.Sprintf( + "%d. [%s/%s] task=%s - %s\n", + i+1, + strings.ToUpper(f.Severity), + f.Category, + f.TaskID, + f.Message, + )) + if f.Recommendation != "" { + b.WriteString(fmt.Sprintf(" Action: %s\n", f.Recommendation)) + } + } + if len(r.Findings) > limit { + b.WriteString(fmt.Sprintf("... and %d more findings.", len(r.Findings)-limit)) + } + return b.String() +} + +type supervisorReview struct { + Score float64 `json:"score"` + Issues []struct { + Category string `json:"category"` + Severity string `json:"severity"` + Message string `json:"message"` + } `json:"issues"` +} + +func (al *AgentLoop) runAuditLoop(ctx context.Context) { + if al.cfg == nil || !al.cfg.Audit.Enabled { + return + } + + intervalMinutes := al.cfg.Audit.IntervalMinutes + if intervalMinutes <= 0 { + intervalMinutes = 30 + } + interval := time.Duration(intervalMinutes) * time.Minute + ticker := time.NewTicker(interval) + defer ticker.Stop() + + // Run once shortly after startup. + timer := time.NewTimer(5 * time.Second) + defer timer.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-timer.C: + al.executeAuditCycle(ctx) + case <-ticker.C: + al.executeAuditCycle(ctx) + } + } +} + +func (al *AgentLoop) executeAuditCycle(ctx context.Context) { + report, err := al.RunTaskAudit(ctx) + if err != nil { + logger.WarnCF("audit", "Task audit failed", map[string]any{"error": err.Error()}) + return + } + if report == nil { + return + } + + logger.InfoCF("audit", "Task audit completed", map[string]any{ + "tasks_scanned": report.TotalTasks, + "findings": len(report.Findings), + }) + + if len(report.Findings) == 0 { + return + } + + al.applyAutoRemediation(report) + al.publishAuditReport(report) +} + +func (al *AgentLoop) RunTaskAudit(ctx context.Context) (*AuditReport, error) { + if al.taskLedger == nil || al.cfg == nil { + return nil, nil + } + + lookback := time.Duration(al.cfg.Audit.LookbackMinutes) * time.Minute + if lookback <= 0 { + lookback = 3 * time.Hour + } + + records := al.taskLedger.ListSince(time.Now().Add(-lookback)) + report := &AuditReport{ + GeneratedAt: time.Now(), + Lookback: lookback, + TotalTasks: len(records), + Findings: make([]AuditFinding, 0), + } + nowMS := time.Now().UnixMilli() + + timeoutSeconds := al.cfg.Orchestration.DefaultTaskTimeoutSeconds + if timeoutSeconds <= 0 { + timeoutSeconds = 180 + } + timeoutMS := int64(timeoutSeconds) * 1000 + retryLimit := al.cfg.Orchestration.RetryLimitPerTask + if retryLimit < 0 { + retryLimit = 0 + } + inconsistencyPolicy := strings.ToLower(strings.TrimSpace(al.cfg.Audit.InconsistencyPolicy)) + if inconsistencyPolicy == "" { + inconsistencyPolicy = "strict" + } + + for _, record := range records { + switch record.Status { + case tools.TaskStatusPlanned: + overdue := false + if record.DeadlineAtMS != nil && nowMS > *record.DeadlineAtMS { + overdue = true + } + if !overdue && nowMS-record.CreatedAtMS > timeoutMS { + overdue = true + } + if overdue { + report.Findings = append(report.Findings, AuditFinding{ + TaskID: record.ID, + Category: "missed", + Severity: "high", + Message: "Task is still planned but appears overdue.", + Recommendation: "Rerun or escalate this task.", + }) + } + case tools.TaskStatusRunning: + if nowMS-record.UpdatedAtMS > timeoutMS { + report.Findings = append(report.Findings, AuditFinding{ + TaskID: record.ID, + Category: "missed", + Severity: "high", + Message: "Task is running past expected timeout.", + Recommendation: "Cancel and retry with a narrower scope.", + }) + } + case tools.TaskStatusCompleted: + if strings.TrimSpace(record.Result) == "" { + report.Findings = append(report.Findings, AuditFinding{ + TaskID: record.ID, + Category: "quality", + Severity: "medium", + Message: "Task completed but produced an empty result.", + Recommendation: "Re-run task and require explicit output fields.", + }) + } + + if len(record.Evidence) == 0 && inconsistencyPolicy == "strict" { + report.Findings = append(report.Findings, AuditFinding{ + TaskID: record.ID, + Category: "inconsistency", + Severity: "medium", + Message: "No execution evidence was captured for a completed task.", + Recommendation: "Re-run with trace capture enabled.", + }) + } + case tools.TaskStatusFailed: + if record.RetryCount < retryLimit { + report.Findings = append(report.Findings, AuditFinding{ + TaskID: record.ID, + Category: "missed", + Severity: "medium", + Message: "Task failed and still has retry budget.", + Recommendation: "Retry this task automatically or manually.", + }) + } + } + } + + modelFindings, err := al.supervisorModelAudit(ctx, records) + if err != nil { + logger.WarnCF("audit", "Supervisor model audit skipped", map[string]any{"error": err.Error()}) + } else { + report.Findings = append(report.Findings, modelFindings...) + } + + return report, nil +} + +func (al *AgentLoop) supervisorModelAudit( + ctx context.Context, + records []tools.TaskLedgerEntry, +) ([]AuditFinding, error) { + if al.cfg == nil || !al.cfg.Audit.Supervisor.Enabled { + return nil, nil + } + modelCfg := al.cfg.Audit.Supervisor.Model + if modelCfg == nil || strings.TrimSpace(modelCfg.Primary) == "" { + return nil, fmt.Errorf("audit supervisor model is not configured") + } + + provider, modelID, err := al.createProviderForModelAlias(modelCfg.Primary) + if err != nil { + return nil, err + } + if provider == nil || strings.TrimSpace(modelID) == "" { + return nil, fmt.Errorf("unable to initialize supervisor provider") + } + if closable, ok := provider.(providers.StatefulProvider); ok { + defer closable.Close() + } + + minConfidence := al.cfg.Audit.MinConfidence + if minConfidence <= 0 { + minConfidence = 0.75 + } + + options := map[string]any{} + if al.cfg.Audit.Supervisor.Temperature != nil { + options["temperature"] = *al.cfg.Audit.Supervisor.Temperature + } + if al.cfg.Audit.Supervisor.MaxTokens > 0 { + options["max_tokens"] = al.cfg.Audit.Supervisor.MaxTokens + } + + findings := make([]AuditFinding, 0) + for _, record := range records { + if record.Status != tools.TaskStatusCompleted { + continue + } + review, err := al.reviewTaskWithSupervisor(ctx, provider, modelID, options, record) + if err != nil { + continue + } + + if review.Score < minConfidence { + findings = append(findings, AuditFinding{ + TaskID: record.ID, + Category: "quality", + Severity: "medium", + Message: fmt.Sprintf("Supervisor confidence %.2f is below threshold %.2f.", review.Score, minConfidence), + Recommendation: "Rerun task with stricter acceptance criteria.", + }) + } + for _, issue := range review.Issues { + category := strings.TrimSpace(strings.ToLower(issue.Category)) + if category == "" { + category = "quality" + } + severity := strings.TrimSpace(strings.ToLower(issue.Severity)) + if severity == "" { + severity = "medium" + } + findings = append(findings, AuditFinding{ + TaskID: record.ID, + Category: category, + Severity: severity, + Message: issue.Message, + Recommendation: "Investigate and re-run affected parts of the task.", + }) + } + } + return findings, nil +} + +func (al *AgentLoop) reviewTaskWithSupervisor( + ctx context.Context, + provider providers.LLMProvider, + modelID string, + options map[string]any, + record tools.TaskLedgerEntry, +) (*supervisorReview, error) { + taskJSON, err := json.MarshalIndent(record, "", " ") + if err != nil { + return nil, err + } + + systemPrompt := `You are a strict operations auditor. +Review the task execution data and output ONLY JSON: +{"score":0.0,"issues":[{"category":"quality|inconsistency|missed","severity":"low|medium|high","message":"..."}]}` + userPrompt := fmt.Sprintf("Task data:\n%s", string(taskJSON)) + + resp, err := provider.Chat(ctx, []providers.Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userPrompt}, + }, nil, modelID, options) + if err != nil { + return nil, err + } + if strings.TrimSpace(resp.Content) == "" { + return nil, fmt.Errorf("empty supervisor response") + } + + parsed, err := parseSupervisorReview(resp.Content) + if err != nil { + return nil, err + } + return parsed, nil +} + +func parseSupervisorReview(raw string) (*supervisorReview, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, fmt.Errorf("empty review content") + } + var review supervisorReview + if err := json.Unmarshal([]byte(raw), &review); err == nil { + return &review, nil + } + + start := strings.Index(raw, "{") + end := strings.LastIndex(raw, "}") + if start < 0 || end <= start { + return nil, fmt.Errorf("no json object in review response") + } + if err := json.Unmarshal([]byte(raw[start:end+1]), &review); err != nil { + return nil, err + } + return &review, nil +} + +func (al *AgentLoop) applyAutoRemediation(report *AuditReport) { + if report == nil || len(report.Findings) == 0 || al.taskLedger == nil { + return + } + + mode := strings.ToLower(strings.TrimSpace(al.cfg.Audit.AutoRemediation)) + if mode == "" || mode == "disabled" || mode == "off" || mode == "none" { + return + } + if mode != "safe_only" { + return + } + + for _, finding := range report.Findings { + if finding.Category != "missed" { + continue + } + _ = al.taskLedger.AddRemediation(finding.TaskID, tools.TaskRemediation{ + Action: "notify", + Status: "queued", + Note: finding.Message, + }) + } +} + +func (al *AgentLoop) publishAuditReport(report *AuditReport) { + if report == nil || len(report.Findings) == 0 || al.bus == nil { + return + } + channel, chatID := al.resolveAuditDestination() + if channel == "" || chatID == "" || constants.IsInternalChannel(channel) { + return + } + + al.bus.PublishOutbound(bus.OutboundMessage{ + Channel: channel, + ChatID: chatID, + Content: report.FormatMessage(), + }) +} + +func (al *AgentLoop) resolveAuditDestination() (string, string) { + if al.cfg == nil { + return "", "" + } + + notify := strings.TrimSpace(al.cfg.Audit.NotifyChannel) + last := "" + if al.state != nil { + last = al.state.GetLastChannel() + } + lastChannel, lastChatID := splitChannelChat(last) + + if notify == "" || notify == "last_active" { + return lastChannel, lastChatID + } + if strings.Contains(notify, ":") { + return splitChannelChat(notify) + } + if lastChatID != "" { + return notify, lastChatID + } + return "", "" +} + +func splitChannelChat(value string) (string, string) { + parts := strings.SplitN(strings.TrimSpace(value), ":", 2) + if len(parts) != 2 { + return "", "" + } + channel := strings.TrimSpace(parts[0]) + chatID := strings.TrimSpace(parts[1]) + if channel == "" || chatID == "" { + return "", "" + } + return channel, chatID +} + +func (al *AgentLoop) createProviderForModelAlias(modelAlias string) (providers.LLMProvider, string, error) { + if al.cfg == nil { + return nil, "", fmt.Errorf("config is nil") + } + modelCfg, err := al.cfg.GetModelConfig(modelAlias) + if err != nil { + return nil, "", err + } + cfgCopy := *modelCfg + if cfgCopy.Workspace == "" { + cfgCopy.Workspace = al.cfg.WorkspacePath() + } + return providers.CreateProviderFromConfig(&cfgCopy) +} diff --git a/pkg/agent/audit_test.go b/pkg/agent/audit_test.go new file mode 100644 index 000000000..a4ba55d66 --- /dev/null +++ b/pkg/agent/audit_test.go @@ -0,0 +1,89 @@ +package agent + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/tools" +) + +func TestRunTaskAudit_DetectsMissedAndQualityIssues(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = t.TempDir() + cfg.Audit.Enabled = true + cfg.Audit.LookbackMinutes = 360 + cfg.Audit.Supervisor.Enabled = false + cfg.Orchestration.DefaultTaskTimeoutSeconds = 1 + cfg.Orchestration.RetryLimitPerTask = 2 + + loop := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) + ledger := loop.GetTaskLedger() + if ledger == nil { + t.Fatal("expected task ledger") + } + + oldTS := time.Now().Add(-10 * time.Minute).UnixMilli() + _ = ledger.UpsertTask(tools.TaskLedgerEntry{ + ID: "task-planned-old", + Status: tools.TaskStatusPlanned, + Intent: "do something", + CreatedAtMS: oldTS, + UpdatedAtMS: oldTS, + }) + _ = ledger.UpsertTask(tools.TaskLedgerEntry{ + ID: "task-completed-empty", + Status: tools.TaskStatusCompleted, + CreatedAtMS: oldTS, + UpdatedAtMS: oldTS, + Result: "", + }) + + report, err := loop.RunTaskAudit(context.Background()) + if err != nil { + t.Fatalf("RunTaskAudit error: %v", err) + } + if report == nil { + t.Fatal("expected non-nil report") + } + if len(report.Findings) < 2 { + t.Fatalf("expected at least 2 findings, got %d", len(report.Findings)) + } + + hasMissed := false + hasQuality := false + for _, f := range report.Findings { + if f.TaskID == "task-planned-old" && f.Category == "missed" { + hasMissed = true + } + if f.TaskID == "task-completed-empty" && f.Category == "quality" { + hasQuality = true + } + } + if !hasMissed { + t.Fatal("expected missed finding for overdue planned task") + } + if !hasQuality { + t.Fatal("expected quality finding for empty completed task") + } +} + +func TestParseSupervisorReview_EmbeddedJSON(t *testing.T) { + raw := "review result:\n{\"score\":0.42,\"issues\":[{\"category\":\"quality\",\"severity\":\"high\",\"message\":\"missing evidence\"}]}" + review, err := parseSupervisorReview(raw) + if err != nil { + t.Fatalf("parseSupervisorReview error: %v", err) + } + if review.Score != 0.42 { + t.Fatalf("score = %v, want 0.42", review.Score) + } + if len(review.Issues) != 1 { + t.Fatalf("issues len = %d, want 1", len(review.Issues)) + } + if !strings.EqualFold(review.Issues[0].Category, "quality") { + t.Fatalf("issue category = %q", review.Issues[0].Category) + } +} diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 19b6dd49f..525a18cd4 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -10,6 +10,7 @@ import ( "context" "encoding/json" "fmt" + "path/filepath" "strings" "sync" "sync/atomic" @@ -34,6 +35,7 @@ type AgentLoop struct { cfg *config.Config registry *AgentRegistry state *state.Manager + taskLedger *tools.TaskLedger running atomic.Bool summarizing sync.Map fallback *providers.FallbackChain @@ -63,21 +65,25 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers // Create state manager using default agent's workspace for channel recording defaultAgent := registry.GetDefaultAgent() var stateManager *state.Manager + ledgerPath := filepath.Join(cfg.WorkspacePath(), "tasks", "ledger.json") if defaultAgent != nil { stateManager = state.NewManager(defaultAgent.Workspace) + ledgerPath = filepath.Join(defaultAgent.Workspace, "tasks", "ledger.json") } + taskLedger := tools.NewTaskLedger(ledgerPath) al := &AgentLoop{ bus: msgBus, cfg: cfg, registry: registry, state: stateManager, + taskLedger: taskLedger, summarizing: sync.Map{}, fallback: fallbackChain, } // Register shared tools to all agents. - registerSharedTools(cfg, msgBus, registry, provider, al) + registerSharedTools(cfg, msgBus, registry, provider, al, taskLedger) return al } @@ -89,6 +95,7 @@ func registerSharedTools( registry *AgentRegistry, provider providers.LLMProvider, sessionsExecutor tools.SessionsSendExecutor, + taskLedger *tools.TaskLedger, ) { for _, agentID := range registry.ListAgentIDs() { agent, ok := registry.GetAgent(agentID) @@ -155,9 +162,23 @@ func registerSharedTools( // Spawn/session tools with allowlist checker. subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace, msgBus) subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) + subagentManager.SetLimits( + cfg.Orchestration.MaxParallelWorkers, + cfg.Orchestration.MaxTasksPerAgent, + cfg.Orchestration.MaxSpawnDepth, + ) + subagentManager.SetTools(agent.Tools) + currentAgentID := agentID + subagentManager.SetExecutionResolver(func(targetAgentID string) (tools.SubagentExecutionConfig, error) { + return resolveSubagentExecution(cfg, registry, provider, currentAgentID, targetAgentID) + }) + if taskLedger != nil { + subagentManager.SetEventHandler(func(event tools.SubagentTaskEvent) { + handleSubagentTaskEvent(taskLedger, cfg, event) + }) + } spawnTool := tools.NewSpawnTool(subagentManager) sessionsSpawnTool := tools.NewSessionsSpawnTool(subagentManager) - currentAgentID := agentID allowlist := func(targetAgentID string) bool { return registry.CanSpawnSubagent(currentAgentID, targetAgentID) } @@ -179,8 +200,106 @@ func registerSharedTools( } } +func resolveSubagentExecution( + cfg *config.Config, + registry *AgentRegistry, + fallbackProvider providers.LLMProvider, + parentAgentID, targetAgentID string, +) (tools.SubagentExecutionConfig, error) { + selectedAgentID := parentAgentID + if strings.TrimSpace(targetAgentID) != "" { + selectedAgentID = targetAgentID + } + + targetAgent, ok := registry.GetAgent(selectedAgentID) + if !ok || targetAgent == nil { + return tools.SubagentExecutionConfig{}, fmt.Errorf("target agent %q not found", selectedAgentID) + } + + execution := tools.SubagentExecutionConfig{ + Provider: fallbackProvider, + Model: targetAgent.Model, + Tools: targetAgent.Tools, + } + + modelCfg, err := cfg.GetModelConfig(targetAgent.Model) + if err != nil { + if execution.Provider != nil { + return execution, nil + } + return tools.SubagentExecutionConfig{}, err + } + + cfgCopy := *modelCfg + if cfgCopy.Workspace == "" { + cfgCopy.Workspace = targetAgent.Workspace + } + + resolvedProvider, resolvedModel, err := providers.CreateProviderFromConfig(&cfgCopy) + if err != nil { + if execution.Provider != nil { + return execution, nil + } + return tools.SubagentExecutionConfig{}, err + } + if resolvedProvider != nil { + execution.Provider = resolvedProvider + } + if resolvedModel != "" { + execution.Model = resolvedModel + } + return execution, nil +} + +func handleSubagentTaskEvent(ledger *tools.TaskLedger, cfg *config.Config, event tools.SubagentTaskEvent) { + if ledger == nil { + return + } + task := event.Task + status := tools.TaskStatus(task.Status) + if status == "" { + status = tools.TaskStatusPlanned + } + + var deadline *int64 + if cfg != nil && cfg.Orchestration.DefaultTaskTimeoutSeconds > 0 { + d := task.Created + int64(cfg.Orchestration.DefaultTaskTimeoutSeconds)*1000 + deadline = &d + } + + _ = ledger.UpsertTask(tools.TaskLedgerEntry{ + ID: task.ID, + ParentTaskID: task.ParentTaskID, + AgentID: task.AgentID, + Source: "spawn", + Intent: task.Task, + OriginChannel: task.OriginChannel, + OriginChatID: task.OriginChatID, + Status: status, + CreatedAtMS: task.Created, + DeadlineAtMS: deadline, + Result: task.Result, + Error: event.Err, + }) + + for _, tr := range event.Trace { + _ = ledger.AddEvidence(task.ID, tools.TaskEvidence{ + TimestampMS: event.Timestamp, + Iteration: tr.Iteration, + ToolName: tr.ToolName, + Arguments: tr.Arguments, + ResultPreview: utils.Truncate(tr.Result, 400), + IsError: tr.IsError, + DurationMS: tr.DurationMS, + }) + } +} + func (al *AgentLoop) Run(ctx context.Context) error { al.running.Store(true) + if al.cfg != nil && al.cfg.Audit.Enabled { + go al.runAuditLoop(ctx) + } for al.running.Load() { select { @@ -241,6 +360,10 @@ func (al *AgentLoop) SetChannelManager(cm *channels.Manager) { al.channelManager = cm } +func (al *AgentLoop) GetTaskLedger() *tools.TaskLedger { + return al.taskLedger +} + // RecordLastChannel records the last active channel for this workspace. // This uses the atomic state save mechanism to prevent data loss on crash. func (al *AgentLoop) RecordLastChannel(channel string) error { diff --git a/pkg/agent/orchestration_test.go b/pkg/agent/orchestration_test.go new file mode 100644 index 000000000..edf9950b5 --- /dev/null +++ b/pkg/agent/orchestration_test.go @@ -0,0 +1,43 @@ +package agent + +import ( + "path/filepath" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/tools" +) + +func TestHandleSubagentTaskEvent_PersistsParentTaskID(t *testing.T) { + ledger := tools.NewTaskLedger(filepath.Join(t.TempDir(), "tasks", "ledger.json")) + cfg := config.DefaultConfig() + cfg.Orchestration.DefaultTaskTimeoutSeconds = 60 + + now := time.Now().UnixMilli() + handleSubagentTaskEvent(ledger, cfg, tools.SubagentTaskEvent{ + Type: tools.SubagentTaskRunning, + Task: tools.SubagentTask{ + ID: "subagent-2", + ParentTaskID: "subagent-1", + Task: "child task", + AgentID: "worker", + OriginChannel: "telegram", + OriginChatID: "chat-1", + Status: "running", + Created: now, + }, + Timestamp: now, + }) + + entry, ok := ledger.Get("subagent-2") + if !ok { + t.Fatal("expected task in ledger") + } + if entry.ParentTaskID != "subagent-1" { + t.Fatalf("ParentTaskID = %q, want %q", entry.ParentTaskID, "subagent-1") + } + if entry.Status != tools.TaskStatusRunning { + t.Fatalf("Status = %q, want %q", entry.Status, tools.TaskStatusRunning) + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 1b4aec388..ebfeae78a 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -47,16 +47,18 @@ func (f *FlexibleStringSlice) UnmarshalJSON(data []byte) error { } type Config struct { - Agents AgentsConfig `json:"agents"` - Bindings []AgentBinding `json:"bindings,omitempty"` - Session SessionConfig `json:"session,omitempty"` - Channels ChannelsConfig `json:"channels"` - Providers ProvidersConfig `json:"providers,omitempty"` - ModelList []ModelConfig `json:"model_list"` // New model-centric provider configuration - Gateway GatewayConfig `json:"gateway"` - Tools ToolsConfig `json:"tools"` - Heartbeat HeartbeatConfig `json:"heartbeat"` - Devices DevicesConfig `json:"devices"` + Agents AgentsConfig `json:"agents"` + Bindings []AgentBinding `json:"bindings,omitempty"` + Session SessionConfig `json:"session,omitempty"` + Channels ChannelsConfig `json:"channels"` + Providers ProvidersConfig `json:"providers,omitempty"` + ModelList []ModelConfig `json:"model_list"` // New model-centric provider configuration + Gateway GatewayConfig `json:"gateway"` + Tools ToolsConfig `json:"tools"` + Heartbeat HeartbeatConfig `json:"heartbeat"` + Devices DevicesConfig `json:"devices"` + Orchestration OrchestrationConfig `json:"orchestration,omitempty"` + Audit AuditConfig `json:"audit,omitempty"` } // MarshalJSON implements custom JSON marshaling for Config @@ -167,17 +169,17 @@ type SessionConfig struct { } type AgentDefaults struct { - Workspace string `json:"workspace" env:"PICOCLAW_AGENTS_DEFAULTS_WORKSPACE"` - RestrictToWorkspace bool `json:"restrict_to_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE"` - Provider string `json:"provider" env:"PICOCLAW_AGENTS_DEFAULTS_PROVIDER"` - ModelName string `json:"model_name,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL_NAME"` - Model string `json:"model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` // Deprecated: use model_name instead - ModelFallbacks []string `json:"model_fallbacks,omitempty"` - ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"` - ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"` - MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` - Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` - MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` + Workspace string `json:"workspace" env:"PICOCLAW_AGENTS_DEFAULTS_WORKSPACE"` + RestrictToWorkspace bool `json:"restrict_to_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE"` + Provider string `json:"provider" env:"PICOCLAW_AGENTS_DEFAULTS_PROVIDER"` + ModelName string `json:"model_name,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL_NAME"` + Model string `json:"model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` // Deprecated: use model_name instead + ModelFallbacks []string `json:"model_fallbacks,omitempty"` + ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"` + ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"` + MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` + Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` + MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` Compaction AgentCompactionConfig `json:"compaction,omitempty"` ContextPruning AgentContextPruningConfig `json:"context_pruning,omitempty"` BootstrapSnapshot AgentBootstrapSnapshotConfig `json:"bootstrap_snapshot,omitempty"` @@ -354,6 +356,33 @@ type DevicesConfig struct { MonitorUSB bool `json:"monitor_usb" env:"PICOCLAW_DEVICES_MONITOR_USB"` } +type OrchestrationConfig struct { + Enabled bool `json:"enabled" env:"PICOCLAW_ORCHESTRATION_ENABLED"` + MaxSpawnDepth int `json:"max_spawn_depth" env:"PICOCLAW_ORCHESTRATION_MAX_SPAWN_DEPTH"` + MaxParallelWorkers int `json:"max_parallel_workers" env:"PICOCLAW_ORCHESTRATION_MAX_PARALLEL_WORKERS"` + MaxTasksPerAgent int `json:"max_tasks_per_agent" env:"PICOCLAW_ORCHESTRATION_MAX_TASKS_PER_AGENT"` + DefaultTaskTimeoutSeconds int `json:"default_task_timeout_seconds" env:"PICOCLAW_ORCHESTRATION_DEFAULT_TASK_TIMEOUT_SECONDS"` + RetryLimitPerTask int `json:"retry_limit_per_task" env:"PICOCLAW_ORCHESTRATION_RETRY_LIMIT_PER_TASK"` +} + +type AuditConfig struct { + Enabled bool `json:"enabled" env:"PICOCLAW_AUDIT_ENABLED"` + IntervalMinutes int `json:"interval_minutes" env:"PICOCLAW_AUDIT_INTERVAL_MINUTES"` + LookbackMinutes int `json:"lookback_minutes" env:"PICOCLAW_AUDIT_LOOKBACK_MINUTES"` + Supervisor AuditSupervisorConfig `json:"supervisor"` + MinConfidence float64 `json:"min_confidence" env:"PICOCLAW_AUDIT_MIN_CONFIDENCE"` + InconsistencyPolicy string `json:"inconsistency_policy" env:"PICOCLAW_AUDIT_INCONSISTENCY_POLICY"` + AutoRemediation string `json:"auto_remediation" env:"PICOCLAW_AUDIT_AUTO_REMEDIATION"` + NotifyChannel string `json:"notify_channel" env:"PICOCLAW_AUDIT_NOTIFY_CHANNEL"` +} + +type AuditSupervisorConfig struct { + Enabled bool `json:"enabled" env:"PICOCLAW_AUDIT_SUPERVISOR_ENABLED"` + Model *AgentModelConfig `json:"model,omitempty"` + Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AUDIT_SUPERVISOR_TEMPERATURE"` + MaxTokens int `json:"max_tokens,omitempty" env:"PICOCLAW_AUDIT_SUPERVISOR_MAX_TOKENS"` +} + type ProvidersConfig struct { Anthropic ProviderConfig `json:"anthropic"` OpenAI OpenAIProviderConfig `json:"openai"` diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index f4b254dcd..b8176d0a9 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -361,6 +361,74 @@ func TestDefaultConfig_OpenAIWebSearchEnabled(t *testing.T) { } } +func TestDefaultConfig_OrchestrationAndAuditDefaults(t *testing.T) { + cfg := DefaultConfig() + + if cfg.Orchestration.MaxSpawnDepth != 3 { + t.Fatalf("MaxSpawnDepth = %d, want 3", cfg.Orchestration.MaxSpawnDepth) + } + if cfg.Orchestration.DefaultTaskTimeoutSeconds != 180 { + t.Fatalf( + "DefaultTaskTimeoutSeconds = %d, want 180", + cfg.Orchestration.DefaultTaskTimeoutSeconds, + ) + } + if cfg.Audit.IntervalMinutes != 30 { + t.Fatalf("Audit.IntervalMinutes = %d, want 30", cfg.Audit.IntervalMinutes) + } + if cfg.Audit.LookbackMinutes != 180 { + t.Fatalf("Audit.LookbackMinutes = %d, want 180", cfg.Audit.LookbackMinutes) + } + if cfg.Audit.Supervisor.Model == nil || cfg.Audit.Supervisor.Model.Primary == "" { + t.Fatal("Audit supervisor model should be initialized in defaults") + } +} + +func TestConfig_UnmarshalAuditAndOrchestration(t *testing.T) { + jsonData := `{ + "orchestration": { + "enabled": true, + "max_spawn_depth": 4, + "max_parallel_workers": 2 + }, + "audit": { + "enabled": true, + "interval_minutes": 15, + "lookback_minutes": 60, + "min_confidence": 0.85, + "supervisor": { + "enabled": true, + "model": { + "primary": "gpt-5.2", + "fallbacks": ["claude-sonnet-4.6"] + }, + "max_tokens": 1024 + } + } + }` + + cfg := DefaultConfig() + if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { + t.Fatalf("unmarshal: %v", err) + } + + if !cfg.Orchestration.Enabled { + t.Fatal("orchestration.enabled should be true") + } + if cfg.Orchestration.MaxSpawnDepth != 4 { + t.Fatalf("max_spawn_depth = %d, want 4", cfg.Orchestration.MaxSpawnDepth) + } + if !cfg.Audit.Enabled { + t.Fatal("audit.enabled should be true") + } + if cfg.Audit.MinConfidence != 0.85 { + t.Fatalf("min_confidence = %v, want 0.85", cfg.Audit.MinConfidence) + } + if cfg.Audit.Supervisor.Model == nil || cfg.Audit.Supervisor.Model.Primary != "gpt-5.2" { + t.Fatalf("unexpected supervisor model: %+v", cfg.Audit.Supervisor.Model) + } +} + func TestDefaultConfig_MemoryVectorDefaults(t *testing.T) { cfg := DefaultConfig() mv := cfg.Agents.Defaults.MemoryVector diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 638cb85fb..cd7426b2a 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -351,5 +351,30 @@ func DefaultConfig() *Config { Enabled: false, MonitorUSB: true, }, + Orchestration: OrchestrationConfig{ + Enabled: false, + MaxSpawnDepth: 3, + MaxParallelWorkers: 4, + MaxTasksPerAgent: 20, + DefaultTaskTimeoutSeconds: 180, + RetryLimitPerTask: 2, + }, + Audit: AuditConfig{ + Enabled: false, + IntervalMinutes: 30, + LookbackMinutes: 180, + MinConfidence: 0.75, + InconsistencyPolicy: "strict", + AutoRemediation: "safe_only", + NotifyChannel: "last_active", + Supervisor: AuditSupervisorConfig{ + Enabled: false, + Model: &AgentModelConfig{ + Primary: "gpt-5.2", + Fallbacks: []string{"claude-sonnet-4.6"}, + }, + MaxTokens: 2048, + }, + }, } } diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index 0646c82a9..403b75538 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -2,10 +2,55 @@ package tools import ( "context" + "fmt" "strings" "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/providers" ) +type blockingProvider struct { + block <-chan struct{} +} + +func (p *blockingProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-p.block: + return &providers.LLMResponse{Content: "done"}, nil + } +} + +func (p *blockingProvider) GetDefaultModel() string { + return "blocking-model" +} + +func extractSpawnTaskID(msg string) (string, error) { + marker := "(id: " + idx := strings.LastIndex(msg, marker) + if idx < 0 { + return "", fmt.Errorf("task id marker not found") + } + rest := msg[idx+len(marker):] + end := strings.Index(rest, ")") + if end < 0 { + return "", fmt.Errorf("task id closing bracket not found") + } + id := strings.TrimSpace(rest[:end]) + if id == "" { + return "", fmt.Errorf("empty task id") + } + return id, nil +} + func TestSpawnTool_Execute_EmptyTask(t *testing.T) { provider := &MockLLMProvider{} manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) @@ -77,3 +122,163 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) { t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM) } } + +func TestSpawnTool_Execute_RespectsMaxTaskLimit(t *testing.T) { + block := make(chan struct{}) + provider := &blockingProvider{block: block} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager.SetLimits(0, 1, 0) + tool := NewSpawnTool(manager) + + ctx := context.Background() + + first := tool.Execute(ctx, map[string]any{"task": "task-1"}) + if first == nil || first.IsError { + t.Fatalf("first spawn should succeed, got %+v", first) + } + + second := tool.Execute(ctx, map[string]any{"task": "task-2"}) + if second == nil { + t.Fatal("second spawn result should not be nil") + } + if !second.IsError { + t.Fatalf("second spawn should fail due to max task limit, got %+v", second) + } + if !strings.Contains(second.ForLLM, "max task limit reached") { + t.Fatalf("unexpected error message: %s", second.ForLLM) + } + + // Unblock the first task to avoid goroutine leak. + close(block) + time.Sleep(10 * time.Millisecond) +} + +func TestSpawnTool_Execute_RespectsMaxDepth(t *testing.T) { + provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager.SetLimits(0, 0, 1) + tool := NewSpawnTool(manager) + + registry := NewToolRegistry() + registry.Register(tool) + + result := registry.ExecuteWithContext( + context.Background(), + "spawn", + map[string]any{"task": "nested task"}, + "cli", + "direct", + "subagent:parent-task", + nil, + ) + if result == nil { + t.Fatal("result should not be nil") + } + if !result.IsError { + t.Fatalf("expected depth-limit error, got %+v", result) + } + if !strings.Contains(result.ForLLM, "max spawn depth reached") { + t.Fatalf("unexpected error message: %s", result.ForLLM) + } +} + +func TestSpawnTool_Execute_SetsParentTaskIDAndDepth(t *testing.T) { + provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + tool := NewSpawnTool(manager) + + registry := NewToolRegistry() + registry.Register(tool) + + parent := registry.ExecuteWithContext( + context.Background(), + "spawn", + map[string]any{"task": "parent task"}, + "cli", + "direct", + "", + nil, + ) + if parent == nil || parent.IsError { + t.Fatalf("parent spawn failed: %+v", parent) + } + parentID, err := extractSpawnTaskID(parent.ForLLM) + if err != nil { + t.Fatalf("extract parent id failed: %v", err) + } + + child := registry.ExecuteWithContext( + context.Background(), + "spawn", + map[string]any{"task": "child task"}, + "cli", + "direct", + "subagent:"+parentID, + nil, + ) + if child == nil || child.IsError { + t.Fatalf("child spawn failed: %+v", child) + } + childID, err := extractSpawnTaskID(child.ForLLM) + if err != nil { + t.Fatalf("extract child id failed: %v", err) + } + + childTask, ok := manager.GetTask(childID) + if !ok { + t.Fatalf("child task %q not found", childID) + } + if childTask.ParentTaskID != parentID { + t.Fatalf("child ParentTaskID = %q, want %q", childTask.ParentTaskID, parentID) + } + if childTask.Depth != 2 { + t.Fatalf("child depth = %d, want 2", childTask.Depth) + } +} + +func TestSubagentManager_CancelTaskTree(t *testing.T) { + manager := NewSubagentManager(&MockLLMProvider{}, "test-model", "/tmp/test", nil) + + parentCtx, parentCancel := context.WithCancel(context.Background()) + defer parentCancel() + childCtx, childCancel := context.WithCancel(context.Background()) + defer childCancel() + grandCtx, grandCancel := context.WithCancel(context.Background()) + defer grandCancel() + otherCtx, otherCancel := context.WithCancel(context.Background()) + defer otherCancel() + + manager.mu.Lock() + manager.tasks["parent"] = &SubagentTask{ID: "parent", Status: "running"} + manager.tasks["child"] = &SubagentTask{ID: "child", ParentTaskID: "parent", Status: "running"} + manager.tasks["grand"] = &SubagentTask{ID: "grand", ParentTaskID: "child", Status: "running"} + manager.tasks["other"] = &SubagentTask{ID: "other", ParentTaskID: "", Status: "running"} + manager.taskCancels["parent"] = parentCancel + manager.taskCancels["child"] = childCancel + manager.taskCancels["grand"] = grandCancel + manager.taskCancels["other"] = otherCancel + manager.mu.Unlock() + + manager.cancelTaskTree("parent") + + select { + case <-childCtx.Done(): + case <-time.After(100 * time.Millisecond): + t.Fatal("expected child context to be cancelled") + } + select { + case <-grandCtx.Done(): + case <-time.After(100 * time.Millisecond): + t.Fatal("expected grandchild context to be cancelled") + } + select { + case <-otherCtx.Done(): + t.Fatal("unrelated task should not be cancelled") + default: + } + select { + case <-parentCtx.Done(): + t.Fatal("parent task is not cancelled by cancelTaskTree") + default: + } +} diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 960a41e5e..3aeb482ad 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "strings" "sync" "time" @@ -12,6 +13,7 @@ import ( type SubagentTask struct { ID string + ParentTaskID string Task string Label string AgentID string @@ -20,8 +22,37 @@ type SubagentTask struct { Status string Result string Created int64 + Depth int } +type SubagentExecutionConfig struct { + Provider providers.LLMProvider + Model string + Tools *ToolRegistry +} + +type SubagentExecutionResolver func(targetAgentID string) (SubagentExecutionConfig, error) + +type SubagentTaskEventType string + +const ( + SubagentTaskCreated SubagentTaskEventType = "created" + SubagentTaskRunning SubagentTaskEventType = "running" + SubagentTaskCompleted SubagentTaskEventType = "completed" + SubagentTaskFailed SubagentTaskEventType = "failed" + SubagentTaskCancelled SubagentTaskEventType = "cancelled" +) + +type SubagentTaskEvent struct { + Type SubagentTaskEventType + Task SubagentTask + Trace []ToolExecutionTrace + Err string + Timestamp int64 +} + +type SubagentTaskEventHandler func(event SubagentTaskEvent) + type SubagentManager struct { tasks map[string]*SubagentTask mu sync.RWMutex @@ -36,6 +67,12 @@ type SubagentManager struct { hasMaxTokens bool hasTemperature bool nextID int + resolver SubagentExecutionResolver + eventHandler SubagentTaskEventHandler + maxConcurrent int + maxTasks int + maxDepth int + taskCancels map[string]context.CancelFunc } func NewSubagentManager( @@ -52,6 +89,7 @@ func NewSubagentManager( tools: NewToolRegistry(), maxIterations: 10, nextID: 1, + taskCancels: make(map[string]context.CancelFunc), } } @@ -80,6 +118,33 @@ func (sm *SubagentManager) RegisterTool(tool Tool) { sm.tools.Register(tool) } +// SetExecutionResolver sets a resolver used to pick provider/model/tools for each task. +// If not set, the manager falls back to its default provider/model/tools. +func (sm *SubagentManager) SetExecutionResolver(resolver SubagentExecutionResolver) { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.resolver = resolver +} + +// SetEventHandler sets a callback that receives task lifecycle events. +func (sm *SubagentManager) SetEventHandler(handler SubagentTaskEventHandler) { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.eventHandler = handler +} + +// SetLimits sets orchestration limits for this manager. +// maxConcurrent <= 0 means unlimited. +// maxTasks <= 0 means unlimited. +// maxDepth <= 0 means unlimited. +func (sm *SubagentManager) SetLimits(maxConcurrent, maxTasks, maxDepth int) { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.maxConcurrent = maxConcurrent + sm.maxTasks = maxTasks + sm.maxDepth = maxDepth +} + func (sm *SubagentManager) Spawn( ctx context.Context, task, label, agentID, originChannel, originChatID string, @@ -105,11 +170,51 @@ func (sm *SubagentManager) SpawnTask( sm.mu.Lock() defer sm.mu.Unlock() + if sm.maxTasks > 0 { + active := 0 + for _, item := range sm.tasks { + if !isTerminalTaskStatus(item.Status) { + active++ + } + } + if active >= sm.maxTasks { + return nil, fmt.Errorf("max task limit reached (%d)", sm.maxTasks) + } + } + if sm.maxConcurrent > 0 { + running := 0 + for _, item := range sm.tasks { + if item.Status == "running" { + running++ + } + } + if running >= sm.maxConcurrent { + return nil, fmt.Errorf("max concurrent task limit reached (%d)", sm.maxConcurrent) + } + } + + depth := 1 + parentTaskID := "" + senderID := toolExecutionSenderID(ctx) + if strings.HasPrefix(strings.ToLower(senderID), "subagent:") { + parentID := strings.TrimSpace(senderID[len("subagent:"):]) + parentTaskID = parentID + if parent, ok := sm.tasks[parentID]; ok && parent.Depth > 0 { + depth = parent.Depth + 1 + } else { + depth = 2 + } + } + if sm.maxDepth > 0 && depth > sm.maxDepth { + return nil, fmt.Errorf("max spawn depth reached (%d)", sm.maxDepth) + } + taskID := fmt.Sprintf("subagent-%d", sm.nextID) sm.nextID++ subagentTask := &SubagentTask{ ID: taskID, + ParentTaskID: parentTaskID, Task: task, Label: label, AgentID: agentID, @@ -117,19 +222,36 @@ func (sm *SubagentManager) SpawnTask( OriginChatID: originChatID, Status: "running", Created: time.Now().UnixMilli(), + Depth: depth, } sm.tasks[taskID] = subagentTask + taskCtx, cancel := context.WithCancel(ctx) + sm.taskCancels[taskID] = cancel + snapshot := *subagentTask // Start task in background with context cancellation support. - go sm.runTask(ctx, subagentTask, callback) - - snapshot := *subagentTask + go sm.runTask(taskCtx, subagentTask, callback) + go sm.emitEvent(SubagentTaskEvent{ + Type: SubagentTaskCreated, + Task: snapshot, + Timestamp: time.Now().UnixMilli(), + }) return &snapshot, nil } func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, callback AsyncCallback) { + defer sm.clearTaskCancel(task.ID) + + sm.mu.Lock() task.Status = "running" task.Created = time.Now().UnixMilli() + snapshot := *task + sm.mu.Unlock() + sm.emitEvent(SubagentTaskEvent{ + Type: SubagentTaskRunning, + Task: snapshot, + Timestamp: time.Now().UnixMilli(), + }) // Build system prompt for subagent systemPrompt := `You are a subagent. Complete the given task independently and report the result. @@ -151,23 +273,79 @@ After completing the task, provide a clear summary of what was done.` select { case <-ctx.Done(): sm.mu.Lock() - task.Status = "canceled" - task.Result = "Task canceled before execution" + task.Status = "cancelled" + task.Result = "Task cancelled before execution" + snapshot = *task sm.mu.Unlock() + sm.emitEvent(SubagentTaskEvent{ + Type: SubagentTaskCancelled, + Task: snapshot, + Timestamp: time.Now().UnixMilli(), + }) + sm.cancelTaskTree(task.ID) return default: } // Run tool loop with access to tools sm.mu.RLock() - tools := sm.tools maxIter := sm.maxIterations maxTokens := sm.maxTokens temperature := sm.temperature hasMaxTokens := sm.hasMaxTokens hasTemperature := sm.hasTemperature + resolver := sm.resolver + defaultTools := sm.tools + defaultProvider := sm.provider + defaultModel := sm.defaultModel sm.mu.RUnlock() + execution := SubagentExecutionConfig{ + Provider: defaultProvider, + Model: defaultModel, + Tools: defaultTools, + } + if resolver != nil { + resolved, err := resolver(task.AgentID) + if err != nil { + sm.mu.Lock() + task.Status = "failed" + task.Result = fmt.Sprintf("Error: %v", err) + snapshot = *task + sm.mu.Unlock() + sm.emitEvent(SubagentTaskEvent{ + Type: SubagentTaskFailed, + Task: snapshot, + Err: err.Error(), + Timestamp: time.Now().UnixMilli(), + }) + if callback != nil { + callback(ctx, &ToolResult{ + ForLLM: task.Result, + IsError: true, + Err: err, + }) + } + if sm.bus != nil { + sm.publishTaskAnnouncement(snapshot, task.Status) + } + sm.cancelTaskTree(task.ID) + return + } + if resolved.Provider != nil { + execution.Provider = resolved.Provider + } + if resolved.Model != "" { + execution.Model = resolved.Model + } + if resolved.Tools != nil { + execution.Tools = resolved.Tools + } + } + if execution.Tools == nil { + execution.Tools = NewToolRegistry() + } + var llmOptions map[string]any if hasMaxTokens || hasTemperature { llmOptions = map[string]any{} @@ -180,31 +358,34 @@ After completing the task, provide a clear summary of what was done.` } loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - Model: sm.defaultModel, - Tools: tools, + Provider: execution.Provider, + Model: execution.Model, + Tools: execution.Tools, MaxIterations: maxIter, LLMOptions: llmOptions, + SenderID: fmt.Sprintf("subagent:%s", task.ID), }, messages, task.OriginChannel, task.OriginChatID) - sm.mu.Lock() var result *ToolResult - defer func() { - sm.mu.Unlock() - // Call callback if provided and result is set - if callback != nil && result != nil { - callback(ctx, result) - } - }() + event := SubagentTaskEvent{ + Task: *task, + Timestamp: time.Now().UnixMilli(), + } + if loopResult != nil { + event.Trace = loopResult.Trace + } if err != nil { + sm.mu.Lock() task.Status = "failed" task.Result = fmt.Sprintf("Error: %v", err) - // Check if it was canceled - if ctx.Err() != nil { - task.Status = "canceled" - task.Result = "Task canceled during execution" - } + // Check if it was cancelled + if ctx.Err() != nil { + task.Status = "cancelled" + task.Result = "Task cancelled during execution" + } + snapshot = *task + sm.mu.Unlock() result = &ToolResult{ ForLLM: task.Result, ForUser: "", @@ -213,9 +394,20 @@ After completing the task, provide a clear summary of what was done.` Async: false, Err: err, } + event.Task = snapshot + event.Err = err.Error() + if task.Status == "cancelled" { + event.Type = SubagentTaskCancelled + } else { + event.Type = SubagentTaskFailed + } + sm.cancelTaskTree(task.ID) } else { + sm.mu.Lock() task.Status = "completed" task.Result = loopResult.Content + snapshot = *task + sm.mu.Unlock() result = &ToolResult{ ForLLM: fmt.Sprintf( "Subagent '%s' completed (iterations: %d): %s", @@ -228,18 +420,19 @@ After completing the task, provide a clear summary of what was done.` IsError: false, Async: false, } + event.Type = SubagentTaskCompleted + event.Task = snapshot + } + sm.emitEvent(event) + + // Call callback if provided and result is set. + if callback != nil && result != nil { + callback(ctx, result) } // Send announce message back to main agent if sm.bus != nil { - announceContent := fmt.Sprintf("Task '%s' completed.\n\nResult:\n%s", task.Label, task.Result) - sm.bus.PublishInbound(bus.InboundMessage{ - Channel: "system", - SenderID: fmt.Sprintf("subagent:%s", task.ID), - // Format: "original_channel:original_chat_id" for routing back - ChatID: fmt.Sprintf("%s:%s", task.OriginChannel, task.OriginChatID), - Content: announceContent, - }) + sm.publishTaskAnnouncement(snapshot, task.Status) } } @@ -247,7 +440,11 @@ func (sm *SubagentManager) GetTask(taskID string) (*SubagentTask, bool) { sm.mu.RLock() defer sm.mu.RUnlock() task, ok := sm.tasks[taskID] - return task, ok + if !ok { + return nil, false + } + snapshot := *task + return &snapshot, true } func (sm *SubagentManager) ListTasks() []*SubagentTask { @@ -256,11 +453,85 @@ func (sm *SubagentManager) ListTasks() []*SubagentTask { tasks := make([]*SubagentTask, 0, len(sm.tasks)) for _, task := range sm.tasks { - tasks = append(tasks, task) + snapshot := *task + tasks = append(tasks, &snapshot) } return tasks } +func (sm *SubagentManager) emitEvent(event SubagentTaskEvent) { + sm.mu.RLock() + handler := sm.eventHandler + sm.mu.RUnlock() + if handler == nil { + return + } + handler(event) +} + +func (sm *SubagentManager) publishTaskAnnouncement(task SubagentTask, status string) { + if sm.bus == nil { + return + } + state := "completed" + if status == "failed" { + state = "failed" + } else if status == "cancelled" { + state = "cancelled" + } + announceContent := fmt.Sprintf("Task '%s' %s.\n\nResult:\n%s", task.Label, state, task.Result) + sm.bus.PublishInbound(bus.InboundMessage{ + Channel: "system", + SenderID: fmt.Sprintf("subagent:%s", task.ID), + // Format: "original_channel:original_chat_id" for routing back + ChatID: fmt.Sprintf("%s:%s", task.OriginChannel, task.OriginChatID), + Content: announceContent, + }) +} + +func (sm *SubagentManager) clearTaskCancel(taskID string) { + sm.mu.Lock() + defer sm.mu.Unlock() + delete(sm.taskCancels, taskID) +} + +func (sm *SubagentManager) cancelTaskTree(parentTaskID string) { + queue := []string{strings.TrimSpace(parentTaskID)} + for len(queue) > 0 { + current := queue[0] + queue = queue[1:] + + var children []string + var cancels []context.CancelFunc + + sm.mu.RLock() + for taskID, task := range sm.tasks { + if task.ParentTaskID != current { + continue + } + children = append(children, taskID) + if cancel, ok := sm.taskCancels[taskID]; ok { + cancels = append(cancels, cancel) + } + } + sm.mu.RUnlock() + + for _, cancel := range cancels { + cancel() + } + queue = append(queue, children...) + } +} + +func isTerminalTaskStatus(status string) bool { + switch strings.ToLower(strings.TrimSpace(status)) { + case "completed", "failed", "cancelled": + return true + default: + return false + } +} + // SubagentTool executes a subagent task synchronously and returns the result. // Unlike SpawnTool which runs tasks asynchronously, SubagentTool waits for completion // and returns the result directly in the ToolResult. @@ -298,6 +569,10 @@ func (t *SubagentTool) Parameters() map[string]any { "type": "string", "description": "Optional short label for the task (for display)", }, + "agent_id": map[string]any{ + "type": "string", + "description": "Optional target agent ID to delegate the task to", + }, }, "required": []string{"task"}, } @@ -315,6 +590,7 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe } label, _ := args["label"].(string) + agentID, _ := args["agent_id"].(string) if t.manager == nil { return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil")) @@ -335,14 +611,38 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe // Use RunToolLoop to execute with tools (same as async SpawnTool) sm := t.manager sm.mu.RLock() - tools := sm.tools maxIter := sm.maxIterations maxTokens := sm.maxTokens temperature := sm.temperature hasMaxTokens := sm.hasMaxTokens hasTemperature := sm.hasTemperature + resolver := sm.resolver + execution := SubagentExecutionConfig{ + Provider: sm.provider, + Model: sm.defaultModel, + Tools: sm.tools, + } sm.mu.RUnlock() + if resolver != nil { + resolved, resolveErr := resolver(agentID) + if resolveErr != nil { + return ErrorResult(fmt.Sprintf("Subagent execution failed: %v", resolveErr)).WithError(resolveErr) + } + if resolved.Provider != nil { + execution.Provider = resolved.Provider + } + if resolved.Model != "" { + execution.Model = resolved.Model + } + if resolved.Tools != nil { + execution.Tools = resolved.Tools + } + } + if execution.Tools == nil { + execution.Tools = NewToolRegistry() + } + var llmOptions map[string]any if hasMaxTokens || hasTemperature { llmOptions = map[string]any{} @@ -355,9 +655,9 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe } loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - Model: sm.defaultModel, - Tools: tools, + Provider: execution.Provider, + Model: execution.Model, + Tools: execution.Tools, MaxIterations: maxIter, LLMOptions: llmOptions, }, messages, t.originChannel, t.originChatID) diff --git a/pkg/tools/subagent_tool_test.go b/pkg/tools/subagent_tool_test.go index 59bfdffae..46100d14e 100644 --- a/pkg/tools/subagent_tool_test.go +++ b/pkg/tools/subagent_tool_test.go @@ -3,7 +3,9 @@ package tools import ( "context" "strings" + "sync" "testing" + "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/providers" @@ -14,6 +16,24 @@ type MockLLMProvider struct { lastOptions map[string]any } +type staticMockProvider struct { + content string +} + +func (m *staticMockProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + return &providers.LLMResponse{Content: m.content}, nil +} + +func (m *staticMockProvider) GetDefaultModel() string { + return "static-model" +} + func (m *MockLLMProvider) Chat( ctx context.Context, messages []providers.Message, @@ -348,3 +368,106 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) { t.Error("ForLLM should contain reference to original task") } } + +func TestSubagentManager_ExecutionResolverUsedBySpawn(t *testing.T) { + defaultProvider := &staticMockProvider{content: "default-provider-result"} + manager := NewSubagentManager(defaultProvider, "default-model", t.TempDir(), nil) + + var resolvedAgentID string + manager.SetExecutionResolver(func(targetAgentID string) (SubagentExecutionConfig, error) { + resolvedAgentID = targetAgentID + return SubagentExecutionConfig{ + Provider: &staticMockProvider{content: "resolved-provider-result"}, + Model: "resolved-model", + Tools: NewToolRegistry(), + }, nil + }) + + done := make(chan *ToolResult, 1) + _, err := manager.Spawn( + context.Background(), + "test task", + "resolver-test", + "agent-worker", + "cli", + "direct", + func(ctx context.Context, result *ToolResult) { + done <- result + }, + ) + if err != nil { + t.Fatalf("Spawn failed: %v", err) + } + + select { + case result := <-done: + if result == nil { + t.Fatal("expected callback result") + } + if result.IsError { + t.Fatalf("expected success result, got error: %s", result.ForLLM) + } + if !strings.Contains(result.ForUser, "resolved-provider-result") { + t.Fatalf("expected resolved provider result, got: %s", result.ForUser) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for spawn callback") + } + + if resolvedAgentID != "agent-worker" { + t.Fatalf("resolver target agent id = %q, want %q", resolvedAgentID, "agent-worker") + } +} + +func TestSubagentManager_EventHandlerLifecycle(t *testing.T) { + manager := NewSubagentManager(&staticMockProvider{content: "ok"}, "default-model", t.TempDir(), nil) + manager.SetExecutionResolver(func(targetAgentID string) (SubagentExecutionConfig, error) { + return SubagentExecutionConfig{ + Provider: &staticMockProvider{content: "ok"}, + Model: "model", + Tools: NewToolRegistry(), + }, nil + }) + + var mu sync.Mutex + seen := map[SubagentTaskEventType]bool{} + manager.SetEventHandler(func(event SubagentTaskEvent) { + mu.Lock() + defer mu.Unlock() + seen[event.Type] = true + }) + + done := make(chan struct{}, 1) + _, err := manager.Spawn( + context.Background(), + "event task", + "event-test", + "", + "cli", + "direct", + func(ctx context.Context, result *ToolResult) { + done <- struct{}{} + }, + ) + if err != nil { + t.Fatalf("Spawn failed: %v", err) + } + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for spawn callback") + } + + mu.Lock() + defer mu.Unlock() + if !seen[SubagentTaskCreated] { + t.Fatalf("expected created event, got %+v", seen) + } + if !seen[SubagentTaskRunning] { + t.Fatalf("expected running event, got %+v", seen) + } + if !seen[SubagentTaskCompleted] { + t.Fatalf("expected completed event, got %+v", seen) + } +} diff --git a/pkg/tools/task_ledger.go b/pkg/tools/task_ledger.go new file mode 100644 index 000000000..ebe9c8113 --- /dev/null +++ b/pkg/tools/task_ledger.go @@ -0,0 +1,341 @@ +package tools + +import ( + "encoding/json" + "errors" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" +) + +type TaskStatus string + +const ( + TaskStatusPlanned TaskStatus = "planned" + TaskStatusRunning TaskStatus = "running" + TaskStatusCompleted TaskStatus = "completed" + TaskStatusFailed TaskStatus = "failed" + TaskStatusCancelled TaskStatus = "cancelled" + TaskStatusDegraded TaskStatus = "degraded" +) + +type TaskEvidence struct { + TimestampMS int64 `json:"timestamp_ms"` + Iteration int `json:"iteration,omitempty"` + ToolName string `json:"tool_name,omitempty"` + Arguments map[string]any `json:"arguments,omitempty"` + ResultPreview string `json:"result_preview,omitempty"` + IsError bool `json:"is_error,omitempty"` + DurationMS int64 `json:"duration_ms,omitempty"` +} + +type TaskRemediation struct { + Action string `json:"action"` + Status string `json:"status"` + Note string `json:"note,omitempty"` + CreatedAtMS int64 `json:"created_at_ms"` +} + +type TaskLedgerEntry struct { + ID string `json:"id"` + ParentTaskID string `json:"parent_task_id,omitempty"` + AgentID string `json:"agent_id,omitempty"` + Source string `json:"source,omitempty"` + Intent string `json:"intent,omitempty"` + OriginChannel string `json:"origin_channel,omitempty"` + OriginChatID string `json:"origin_chat_id,omitempty"` + Status TaskStatus `json:"status"` + CreatedAtMS int64 `json:"created_at_ms"` + UpdatedAtMS int64 `json:"updated_at_ms"` + DeadlineAtMS *int64 `json:"deadline_at_ms,omitempty"` + RetryCount int `json:"retry_count,omitempty"` + Result string `json:"result,omitempty"` + Error string `json:"error,omitempty"` + Evidence []TaskEvidence `json:"evidence,omitempty"` + Remediations []TaskRemediation `json:"remediations,omitempty"` +} + +type taskLedgerStore struct { + Version int `json:"version"` + Tasks []TaskLedgerEntry `json:"tasks"` +} + +type TaskLedger struct { + path string + tasks map[string]*TaskLedgerEntry + mu sync.RWMutex +} + +func NewTaskLedger(path string) *TaskLedger { + l := &TaskLedger{ + path: path, + tasks: make(map[string]*TaskLedgerEntry), + } + _ = l.load() + return l +} + +func (l *TaskLedger) CreateTask(entry TaskLedgerEntry) error { + if strings.TrimSpace(entry.ID) == "" { + return errors.New("task id is required") + } + + l.mu.Lock() + defer l.mu.Unlock() + + now := time.Now().UnixMilli() + entry.ID = strings.TrimSpace(entry.ID) + if _, exists := l.tasks[entry.ID]; exists { + return errors.New("task id already exists") + } + if entry.Status == "" { + entry.Status = TaskStatusPlanned + } + if entry.CreatedAtMS == 0 { + entry.CreatedAtMS = now + } + entry.UpdatedAtMS = now + + snapshot := entry + l.tasks[entry.ID] = &snapshot + return l.saveLocked() +} + +func (l *TaskLedger) UpsertTask(entry TaskLedgerEntry) error { + if strings.TrimSpace(entry.ID) == "" { + return errors.New("task id is required") + } + + l.mu.Lock() + defer l.mu.Unlock() + + now := time.Now().UnixMilli() + entry.ID = strings.TrimSpace(entry.ID) + if existing, ok := l.tasks[entry.ID]; ok { + if entry.ParentTaskID != "" { + existing.ParentTaskID = entry.ParentTaskID + } + if entry.AgentID != "" { + existing.AgentID = entry.AgentID + } + if entry.Source != "" { + existing.Source = entry.Source + } + if entry.Intent != "" { + existing.Intent = entry.Intent + } + if entry.OriginChannel != "" { + existing.OriginChannel = entry.OriginChannel + } + if entry.OriginChatID != "" { + existing.OriginChatID = entry.OriginChatID + } + if entry.Status != "" { + existing.Status = entry.Status + } + if entry.DeadlineAtMS != nil { + existing.DeadlineAtMS = entry.DeadlineAtMS + } + if entry.Result != "" { + existing.Result = entry.Result + } + if entry.Error != "" { + existing.Error = entry.Error + } + existing.UpdatedAtMS = now + return l.saveLocked() + } + + if entry.Status == "" { + entry.Status = TaskStatusPlanned + } + if entry.CreatedAtMS == 0 { + entry.CreatedAtMS = now + } + entry.UpdatedAtMS = now + + snapshot := entry + l.tasks[entry.ID] = &snapshot + return l.saveLocked() +} + +func (l *TaskLedger) SetStatus(taskID string, status TaskStatus, result, taskErr string) error { + l.mu.Lock() + defer l.mu.Unlock() + + entry, ok := l.tasks[strings.TrimSpace(taskID)] + if !ok { + return errors.New("task not found") + } + entry.Status = status + entry.UpdatedAtMS = time.Now().UnixMilli() + entry.Result = result + entry.Error = taskErr + return l.saveLocked() +} + +func (l *TaskLedger) IncrementRetry(taskID string) error { + l.mu.Lock() + defer l.mu.Unlock() + + entry, ok := l.tasks[strings.TrimSpace(taskID)] + if !ok { + return errors.New("task not found") + } + entry.RetryCount++ + entry.UpdatedAtMS = time.Now().UnixMilli() + return l.saveLocked() +} + +func (l *TaskLedger) AddEvidence(taskID string, evidence TaskEvidence) error { + l.mu.Lock() + defer l.mu.Unlock() + + entry, ok := l.tasks[strings.TrimSpace(taskID)] + if !ok { + return errors.New("task not found") + } + if evidence.TimestampMS == 0 { + evidence.TimestampMS = time.Now().UnixMilli() + } + if evidence.Arguments != nil { + clonedArgs := make(map[string]any, len(evidence.Arguments)) + for k, v := range evidence.Arguments { + clonedArgs[k] = v + } + evidence.Arguments = clonedArgs + } + entry.Evidence = append(entry.Evidence, evidence) + entry.UpdatedAtMS = evidence.TimestampMS + return l.saveLocked() +} + +func (l *TaskLedger) AddRemediation(taskID string, remediation TaskRemediation) error { + l.mu.Lock() + defer l.mu.Unlock() + + entry, ok := l.tasks[strings.TrimSpace(taskID)] + if !ok { + return errors.New("task not found") + } + if remediation.CreatedAtMS == 0 { + remediation.CreatedAtMS = time.Now().UnixMilli() + } + entry.Remediations = append(entry.Remediations, remediation) + entry.UpdatedAtMS = remediation.CreatedAtMS + return l.saveLocked() +} + +func (l *TaskLedger) Get(taskID string) (TaskLedgerEntry, bool) { + l.mu.RLock() + defer l.mu.RUnlock() + + entry, ok := l.tasks[strings.TrimSpace(taskID)] + if !ok { + return TaskLedgerEntry{}, false + } + return cloneTaskEntry(*entry), true +} + +func (l *TaskLedger) List() []TaskLedgerEntry { + l.mu.RLock() + defer l.mu.RUnlock() + + return l.listLocked(0) +} + +func (l *TaskLedger) ListSince(since time.Time) []TaskLedgerEntry { + l.mu.RLock() + defer l.mu.RUnlock() + + return l.listLocked(since.UnixMilli()) +} + +func (l *TaskLedger) listLocked(minCreatedAtMS int64) []TaskLedgerEntry { + entries := make([]TaskLedgerEntry, 0, len(l.tasks)) + for _, entry := range l.tasks { + if minCreatedAtMS > 0 && entry.CreatedAtMS < minCreatedAtMS { + continue + } + entries = append(entries, cloneTaskEntry(*entry)) + } + sort.Slice(entries, func(i, j int) bool { + if entries[i].CreatedAtMS == entries[j].CreatedAtMS { + return entries[i].ID < entries[j].ID + } + return entries[i].CreatedAtMS < entries[j].CreatedAtMS + }) + return entries +} + +func (l *TaskLedger) load() error { + l.mu.Lock() + defer l.mu.Unlock() + + data, err := os.ReadFile(l.path) + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + + var store taskLedgerStore + if err := json.Unmarshal(data, &store); err != nil { + return err + } + + l.tasks = make(map[string]*TaskLedgerEntry, len(store.Tasks)) + for i := range store.Tasks { + entry := store.Tasks[i] + snapshot := cloneTaskEntry(entry) + l.tasks[entry.ID] = &snapshot + } + return nil +} + +func (l *TaskLedger) saveLocked() error { + store := taskLedgerStore{ + Version: 1, + Tasks: l.listLocked(0), + } + data, err := json.MarshalIndent(store, "", " ") + if err != nil { + return err + } + + if err := os.MkdirAll(filepath.Dir(l.path), 0o755); err != nil { + return err + } + + tmp := l.path + ".tmp" + if err := os.WriteFile(tmp, data, 0o600); err != nil { + return err + } + return os.Rename(tmp, l.path) +} + +func cloneTaskEntry(entry TaskLedgerEntry) TaskLedgerEntry { + snapshot := entry + if entry.Evidence != nil { + snapshot.Evidence = make([]TaskEvidence, len(entry.Evidence)) + for i := range entry.Evidence { + snapshot.Evidence[i] = entry.Evidence[i] + if entry.Evidence[i].Arguments != nil { + clonedArgs := make(map[string]any, len(entry.Evidence[i].Arguments)) + for k, v := range entry.Evidence[i].Arguments { + clonedArgs[k] = v + } + snapshot.Evidence[i].Arguments = clonedArgs + } + } + } + if entry.Remediations != nil { + snapshot.Remediations = make([]TaskRemediation, len(entry.Remediations)) + copy(snapshot.Remediations, entry.Remediations) + } + return snapshot +} diff --git a/pkg/tools/task_ledger_test.go b/pkg/tools/task_ledger_test.go new file mode 100644 index 000000000..50db194d7 --- /dev/null +++ b/pkg/tools/task_ledger_test.go @@ -0,0 +1,94 @@ +package tools + +import ( + "path/filepath" + "testing" + "time" +) + +func TestTaskLedger_CreateUpdateReload(t *testing.T) { + path := filepath.Join(t.TempDir(), "tasks", "ledger.json") + ledger := NewTaskLedger(path) + + err := ledger.CreateTask(TaskLedgerEntry{ + ID: "task-1", + AgentID: "main", + Source: "spawn", + Intent: "fetch weather", + OriginChannel: "telegram", + OriginChatID: "chat-1", + }) + if err != nil { + t.Fatalf("CreateTask failed: %v", err) + } + + if err := ledger.SetStatus("task-1", TaskStatusRunning, "", ""); err != nil { + t.Fatalf("SetStatus running failed: %v", err) + } + if err := ledger.AddEvidence("task-1", TaskEvidence{ + ToolName: "web_search", + Arguments: map[string]any{"query": "weather"}, + ResultPreview: "sunny", + }); err != nil { + t.Fatalf("AddEvidence failed: %v", err) + } + if err := ledger.SetStatus("task-1", TaskStatusCompleted, "done", ""); err != nil { + t.Fatalf("SetStatus completed failed: %v", err) + } + + entry, ok := ledger.Get("task-1") + if !ok { + t.Fatal("expected task entry") + } + if entry.Status != TaskStatusCompleted { + t.Fatalf("status = %q, want %q", entry.Status, TaskStatusCompleted) + } + if len(entry.Evidence) != 1 { + t.Fatalf("evidence count = %d, want 1", len(entry.Evidence)) + } + + // Reload from disk and verify persistence. + ledger2 := NewTaskLedger(path) + entry2, ok := ledger2.Get("task-1") + if !ok { + t.Fatal("expected reloaded task entry") + } + if entry2.Result != "done" { + t.Fatalf("result = %q, want %q", entry2.Result, "done") + } + if len(entry2.Evidence) != 1 { + t.Fatalf("reloaded evidence count = %d, want 1", len(entry2.Evidence)) + } +} + +func TestTaskLedger_ListSince(t *testing.T) { + path := filepath.Join(t.TempDir(), "tasks", "ledger.json") + ledger := NewTaskLedger(path) + + now := time.Now().UnixMilli() + old := now - int64((2 * time.Hour).Milliseconds()) + recent := now - int64((20 * time.Minute).Milliseconds()) + + if err := ledger.UpsertTask(TaskLedgerEntry{ + ID: "task-old", + Status: TaskStatusCompleted, + CreatedAtMS: old, + }); err != nil { + t.Fatalf("Upsert old task failed: %v", err) + } + if err := ledger.UpsertTask(TaskLedgerEntry{ + ID: "task-recent", + Status: TaskStatusCompleted, + CreatedAtMS: recent, + }); err != nil { + t.Fatalf("Upsert recent task failed: %v", err) + } + + items := ledger.ListSince(time.Now().Add(-1 * time.Hour)) + if len(items) != 1 { + t.Fatalf("ListSince returned %d items, want 1", len(items)) + } + if items[0].ID != "task-recent" { + t.Fatalf("ListSince first id = %q, want %q", items[0].ID, "task-recent") + } +} diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go index 1eed28de1..3cbf987fb 100644 --- a/pkg/tools/toolloop.go +++ b/pkg/tools/toolloop.go @@ -10,6 +10,7 @@ import ( "context" "encoding/json" "fmt" + "time" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" @@ -23,12 +24,25 @@ type ToolLoopConfig struct { Tools *ToolRegistry MaxIterations int LLMOptions map[string]any + SenderID string } // ToolLoopResult contains the result of running the tool loop. type ToolLoopResult struct { Content string Iterations int + Trace []ToolExecutionTrace +} + +// ToolExecutionTrace captures a single tool execution inside the loop. +type ToolExecutionTrace struct { + Iteration int + ToolName string + Arguments map[string]any + Result string + IsError bool + DurationMS int64 + ToolCallID string } // RunToolLoop executes the LLM + tool call iteration loop. @@ -41,6 +55,7 @@ func RunToolLoop( ) (*ToolLoopResult, error) { iteration := 0 var finalContent string + trace := make([]ToolExecutionTrace, 0) for iteration < config.MaxIterations { iteration++ @@ -123,6 +138,7 @@ func RunToolLoop( // 7. Execute tool calls for _, tc := range normalizedToolCalls { + start := time.Now() argsJSON, _ := json.Marshal(tc.Arguments) argsPreview := utils.Truncate(string(argsJSON), 200) logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), @@ -134,7 +150,15 @@ func RunToolLoop( // Execute tool (no async callback for subagents - they run independently) var toolResult *ToolResult if config.Tools != nil { - toolResult = config.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, channel, chatID, "", nil) + toolResult = config.Tools.ExecuteWithContext( + ctx, + tc.Name, + tc.Arguments, + channel, + chatID, + config.SenderID, + nil, + ) } else { toolResult = ErrorResult("No tools available") } @@ -145,6 +169,16 @@ func RunToolLoop( contentForLLM = toolResult.Err.Error() } + trace = append(trace, ToolExecutionTrace{ + Iteration: iteration, + ToolName: tc.Name, + Arguments: tc.Arguments, + Result: contentForLLM, + IsError: toolResult.IsError, + DurationMS: time.Since(start).Milliseconds(), + ToolCallID: tc.ID, + }) + // Add tool result message toolResultMsg := providers.Message{ Role: "tool", @@ -158,5 +192,6 @@ func RunToolLoop( return &ToolLoopResult{ Content: finalContent, Iterations: iteration, + Trace: trace, }, nil }