feat: add orchestration ledger and periodic audit supervisor
This commit is contained in:
root
parent
ed29e2bc1c
commit
49c1863a61
15 changed files with 2138 additions and 60 deletions
|
|
@ -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
|
||||
|
|
|
|||
112
docs/design/orchestration-audit.md
Normal file
112
docs/design/orchestration-audit.md
Normal file
|
|
@ -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:<task-id>`.
|
||||
- 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:
|
||||
|
||||
- `<workspace>/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`.
|
||||
|
||||
463
pkg/agent/audit.go
Normal file
463
pkg/agent/audit.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
89
pkg/agent/audit_test.go
Normal file
89
pkg/agent/audit_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
43
pkg/agent/orchestration_test.go
Normal file
43
pkg/agent/orchestration_test.go
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -57,6 +57,8 @@ type Config struct {
|
|||
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
|
||||
|
|
@ -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"`
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
// Check if it was cancelled
|
||||
if ctx.Err() != nil {
|
||||
task.Status = "canceled"
|
||||
task.Result = "Task canceled during execution"
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
341
pkg/tools/task_ledger.go
Normal file
341
pkg/tools/task_ledger.go
Normal file
|
|
@ -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
|
||||
}
|
||||
94
pkg/tools/task_ledger_test.go
Normal file
94
pkg/tools/task_ledger_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue