From 7a5200b5cac9f6375a9e4e20fa285a7a730f65e4 Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Thu, 5 Mar 2026 01:36:25 +0900 Subject: [PATCH 1/2] feat(tasks-2): implement subagent orchestration container model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Container Model for subagent↔conductor communication: - ContainerMessage + channel-based escalation (ask_conductor, submit_plan) - Deliberate vs Exploratory preset workflows with SubagentPlanState - AnswerSubagentTool + ReviewSubagentPlanTool for conductor side - SessionRecorder extensions (RecordQuestion, RecordPlanSubmit) - Environment context injection (MEMORY.md extraction, plan context) - Orchestration guidance with escalation tools documentation Co-Authored-By: Claude Opus 4.6 --- CLAUDE.md | 16 +- pkg/agent/context.go | 508 ++++- pkg/agent/loop.go | 2650 +++++++++++++++++++++++--- pkg/agent/session_recorder.go | 71 +- pkg/agent/session_recorder_test.go | 141 +- pkg/session/types.go | 9 + pkg/tools/answer_subagent.go | 138 ++ pkg/tools/ask_conductor.go | 110 ++ pkg/tools/ask_conductor_test.go | 77 + pkg/tools/session_recorder.go | 6 + pkg/tools/spawn.go | 53 +- pkg/tools/subagent.go | 1007 ++++++++-- pkg/tools/subagent_container_test.go | 125 ++ pkg/tools/subagent_env_test.go | 192 ++ pkg/tools/subagent_plan_test.go | 98 + pkg/tools/submit_plan.go | 186 ++ pkg/tools/submit_plan_test.go | 229 +++ todo/TASKS-2.md | 14 +- 18 files changed, 5116 insertions(+), 514 deletions(-) create mode 100644 pkg/tools/answer_subagent.go create mode 100644 pkg/tools/ask_conductor.go create mode 100644 pkg/tools/ask_conductor_test.go create mode 100644 pkg/tools/subagent_container_test.go create mode 100644 pkg/tools/subagent_env_test.go create mode 100644 pkg/tools/subagent_plan_test.go create mode 100644 pkg/tools/submit_plan.go create mode 100644 pkg/tools/submit_plan_test.go diff --git a/CLAUDE.md b/CLAUDE.md index 64134ffeb..3f0652f1e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -19,12 +19,24 @@ Lint: `golangci-lint run` - **Interview tool filtering**: `interviewAllowedTools` in `pkg/agent/loop.go` is the single source of truth for tools available during interview/review phases. Both `filterInterviewTools` (strips definitions before LLM call) and `isToolAllowedDuringInterview` (argument-level gating) reference this map. - **History clear**: `/plan start clear` wipes session history and summary on transition to executing. The Mini App review UI offers two sliders: standard approve and approve-with-clear. -## Subagent Orchestration (実装済み部分) +## Subagent Orchestration (実装済み) - **Startup flag**: `--orchestration` で on/off。`SubagentsConfig.Enabled` で gate。 - **Conductor identity**: orchestration 有効時に conductor identity + spawn/subagent guidance を system prompt へ注入。 - **Sandbox/Spawn**: `pkg/tools/sandbox.go`, `pkg/tools/spawn.go` 実装済み。 - **AgentReporter**: `orch.AgentReporter` / `orch.Noop` / `orch.Broadcaster` で統一。main/heartbeat/subagent 全セッションが同一 Broadcaster に発火。Mini App は `agentLoop.GetOrchBroadcaster()` → `handler.SetOrchBroadcaster()` で受信。 +- **Container Model (Q&A escalation)**: + - `ContainerMessage` + `inCh`/`outCh` channels on `SubagentTask` — deliberate preset (coder/worker/coordinator) のみ + - `ask_conductor` tool — subagent → conductor question (blocking) + - `answer_subagent` tool — conductor → subagent answer + - `submit_plan` tool — subagent → conductor plan review (blocking) + - `review_subagent_plan` tool — conductor → subagent approve/reject + - `PendingQuestions()` で conductor LLM loop に question/plan_review を注入 +- **Deliberate Plan Mode**: `SubagentPlanState` (Clarifying → Review → Executing → Completed) + - `runDeliberateTask()`: clarifying phase (ask_conductor + submit_plan のみ) → executing phase (全ツール) + - `runExploratoryTask()`: exploratory preset の single-phase loop +- **Environment injection**: `extractPlanContext()` で MEMORY.md から Context/Commands/Orchestration セクションを抽出 → subagent system prompt に注入 +- **SessionRecorder 拡張**: `RecordQuestion()` / `RecordPlanSubmit()` + `TurnQuestion` / `TurnPlanSubmit` TurnKind ## Session DAG (Phase 0–3 実装済み) @@ -61,7 +73,7 @@ Lint: `golangci-lint run` | ファイル | 概要 | |---|---| | [`todo/TASKS-1.md`](todo/TASKS-1.md) | ~~**Memory & Performance Optimization**~~ ✅ 実装済み（MemoryStore キャッシュ＋パース済み state、FunctionCall.Arguments map統一、ToolDefinition.Parameters RawMessage化、検索結果フォーマット共通化、stats 定期フラッシュ） | -| [`todo/TASKS-2.md`](todo/TASKS-2.md) | **Subagent Orchestration (Container Model)** — SubagentContainer、Orchestrator、Presets enforcement、Subagent Plan Mode（TASKS-1 の型変更前提メモ追記済み） | +| [`todo/TASKS-2.md`](todo/TASKS-2.md) | ~~**Subagent Orchestration (Container Model)**~~ ✅ 実装済み（Container Q&A escalation、Deliberate Plan Mode、Environment injection、SessionRecorder 拡張） | | [`todo/TASKS-3.md`](todo/TASKS-3.md) | ~~**Session DAG (SQLite Store)**~~ ✅ 実装済み（Phase 0–3: SQLite SessionStore、LegacyAdapter、Fork/Report、CompactOldTurns、`/session` CLI コマンド、Mini App グラフ UI） | | [`todo/TASKS-4.md`](todo/TASKS-4.md) | **Mini App & Static Serving** — 静的配信の汎用化、バンドラ導入、フロントエンドテスト追加 | | [`todo/TASKS-5.md`](todo/TASKS-5.md) | ~~**Heartbeat Worktree Management**~~ ✅ 実装済み（`/plan worktrees` の `list/inspect/merge/dispose`、安全化した `PruneOrphaned`、Mini App `/miniapp/api/worktrees` + Git タブ UI） | diff --git a/pkg/agent/context.go b/pkg/agent/context.go index 469a0bfdc..ce571ec7b 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -19,71 +19,155 @@ import ( const orchestrationGuidance = `## Orchestration + + You are the conductor, not the performer. **Your primary job is to delegate, not to implement.** + + ### spawn (non-blocking) — DEFAULT choice + Returns immediately. Use for any task that can run independently. + Call the spawn tool with JSON arguments like this: + + Tool: spawn + Arguments: {"task": "Examine pkg/auth/ and report middleware pattern", "preset": "scout", "label": "auth-scout"} + + Tool: spawn + Arguments: {"task": "Implement rate limiter in pkg/ratelimit/ with tests", "preset": "coder", "label": "rate-limiter"} + + ### subagent (blocking) — only when you need the answer NOW + Blocks until the subagent finishes. Use only when you cannot proceed without the result. + Does not take a preset — it runs with default tools. + + Tool: subagent + Arguments: {"task": "Read pkg/config/config.go and list all SubagentsConfig fields", "label": "config-check"} + + ### When to use which + - spawn: parallel tasks, independent work, implementation, long analysis, >2 tool calls + - subagent: you need the result before your next decision + - inline: single quick tool call where delegation overhead is wasteful + + ### Presets (for spawn only) + | preset | role | can write | can exec | + |--------|------|-----------|----------| + | scout | explore, investigate | no | no | + | analyst | analyze, run tests | no | go test/vet, git | + | coder | implement + verify | yes (sandbox) | test/lint/fmt | + | worker | build + install | yes (sandbox) | build/package mgr | + | coordinator | orchestrate others | yes (sandbox) | general + spawn | + + ### Parallel spawning + Spawn multiple independent tasks at once — do NOT wait between them: + + Tool: spawn + Arguments: {"task": "Analyze error handling patterns in pkg/providers/", "preset": "analyst", "label": "error-patterns"} + + Tool: spawn + Arguments: {"task": "List all HTTP endpoints in pkg/miniapp/", "preset": "scout", "label": "endpoints"} + + After spawning, record the assignment in ## Orchestration > Delegated in MEMORY.md. -When results come back, synthesize findings and decide the next fork.` + +When results come back, synthesize findings and decide the next fork. + + + +### Subagent escalation + +Deliberate subagents (coder/worker/coordinator) may ask you questions or submit plans for review. + +When a subagent question appears, respond with the appropriate tool: + +- answer_subagent: Answer a subagent's clarifying question + +- review_subagent_plan: Approve or reject a subagent's execution plan (decision: "approved" or rejection feedback) + + + +### Orchestration Memory + +Maintain these sections in MEMORY.md under ## Orchestration: + +- **Delegated**: Active subagent assignments (task ID, preset, description) + +- **Findings**: Synthesized results from completed subagents + +- **Decisions**: Key architectural/implementation decisions made during orchestration` type ContextBuilder struct { - workspace string - workDir string // session-specific working directory (worktree or project subdir) - skillsLoader *skills.SkillsLoader - memory *MemoryStore - tools *tools.ToolRegistry // Direct reference to tool registry - peerNote string // set per-call from loop.go for peer session awareness - orchestrationEnabled bool // set from AgentLoop when --orchestration flag is used + workspace string + + workDir string // session-specific working directory (worktree or project subdir) + + skillsLoader *skills.SkillsLoader + + memory *MemoryStore + + tools *tools.ToolRegistry // Direct reference to tool registry + + peerNote string // set per-call from loop.go for peer session awareness + + orchestrationEnabled bool // set from AgentLoop when --orchestration flag is used // Cache for system prompt to avoid rebuilding on every call. + // This fixes issue #607: repeated reprocessing of the entire context. + // The cache auto-invalidates when workspace source files change (mtime check). - systemPromptMutex sync.RWMutex + + systemPromptMutex sync.RWMutex + cachedSystemPrompt string - cachedAt time.Time // max observed mtime across tracked paths at cache build time + + cachedAt time.Time // max observed mtime across tracked paths at cache build time // existedAtCache tracks which source file paths existed the last time the + // cache was built. This lets sourceFilesChanged detect files that are newly + // created (didn't exist at cache time, now exist) or deleted (existed at + // cache time, now gone) — both of which should trigger a cache rebuild. + existedAtCache map[string]bool } @@ -92,40 +176,52 @@ func getGlobalConfigDir() string { if err != nil { return "" } + return filepath.Join(home, ".picoclaw") } func NewContextBuilder(workspace string) *ContextBuilder { // builtin skills: skills directory in current project + // Use the skills/ directory under the current working directory + wd, _ := os.Getwd() + builtinSkillsDir := filepath.Join(wd, "skills") + globalSkillsDir := filepath.Join(getGlobalConfigDir(), "skills") return &ContextBuilder{ - workspace: workspace, + workspace: workspace, + skillsLoader: skills.NewSkillsLoader(workspace, globalSkillsDir, builtinSkillsDir), - memory: NewMemoryStore(workspace), + + memory: NewMemoryStore(workspace), } } // SetToolsRegistry sets the tools registry for dynamic tool summary generation. + func (cb *ContextBuilder) SetToolsRegistry(registry *tools.ToolRegistry) { cb.tools = registry } // SetWorkDir sets the session-specific working directory (e.g., worktree path + // or project subdirectory). Bootstrap files found here take priority over workspace. + func (cb *ContextBuilder) SetWorkDir(dir string) { cb.workDir = dir } // SetPeerNote sets the peer session awareness note for the current call. + func (cb *ContextBuilder) SetPeerNote(note string) { cb.peerNote = note } // SetOrchestrationEnabled sets whether orchestration is enabled. + func (cb *ContextBuilder) SetOrchestrationEnabled(enabled bool) { cb.orchestrationEnabled = enabled } @@ -134,83 +230,153 @@ func (cb *ContextBuilder) getIdentity() string { workspacePath, _ := filepath.Abs(filepath.Join(cb.workspace)) // Build tools section dynamically + toolsSection := cb.buildToolsSection() // Build prompt with optional orchestration banner + var prompt string + if cb.orchestrationEnabled { prompt = ` /_/_/_/_/_/_/_/_/_/_/_/_/_/_/ + O R C H E S T R A M O D E + /_/_/_/_/_/_/_/_/_/_/_/_/_/_/ + + ` } + // Conditional identity and plan executing rule for orchestration mode + identity := "a helpful AI assistant" + executingRule := `Work through the current Phase's steps. + Mark each "- [x]" via edit_file. The system will auto-advance phases.` + if cb.orchestrationEnabled { identity = "a conductor AI agent that orchestrates subagents" + executingRule = `Delegate the current Phase's steps to subagents using spawn. + For each step: spawn a subagent with the appropriate preset (scout for investigation, + coder for implementation, analyst for review). Spawn multiple independent steps in parallel. + When a subagent completes, mark "- [x]" via edit_file and record findings in + ## Orchestration > Findings in MEMORY.md. + Only do a step inline if it's a single quick tool call (e.g., reading one file).` } return fmt.Sprintf(prompt+`# picoclaw 🦞 + + You are picoclaw, %s. + + ## Workspace + Your workspace is at: %s + - Memory: %s/memory/MEMORY.md + - Daily Notes: %s/memory/YYYYMM/YYYYMMDD.md + - Skills: %s/skills/{skill-name}/SKILL.md + + %s + + ## Important Rules + + 1. **ALWAYS use tools** - When you need to perform an action (schedule reminders, send messages, execute commands, etc.), you MUST call the appropriate tool. Do NOT just say you'll do it or pretend to do it. + + 2. **Be helpful and accurate** - When using tools, briefly explain what you're doing. + + 3. **Memory & Plans** + - Use memory/MEMORY.md for structured plans. + - NEVER remove or overwrite the header block (# Active Plan, > Task:, > Status:, > Phase:). The system parses these lines to track plan state. + - If Status is "interviewing": Ask clarifying questions. + After each answer, use edit_file to save findings to ## Context in memory/MEMORY.md. + When you have enough information, add ## Phase sections with "- [ ]" checkbox steps, and ## Commands section below the header. Then change > Status: to "review". + - If Status is "review": The plan is awaiting user approval. Do NOT change Status yourself. + - If Status is "executing": %s + - Plan format (header is written by the system — do NOT delete it): + # Active Plan + > Task: + > Status: interviewing | review | executing + > Phase: + ## Phase 1: + - [ ] Step 1 + - [ ] Step 2 + ## Phase 2: <title> + - [ ] Step 1 + ## Commands + build: <build command> + test: <test command> + lint: <lint command> + ## Context + <requirements, decisions, environment> + - Keep each phase to 3-5 steps. Do NOT create plans without /plan. + - Always ask about build/test/lint commands during interview. + + 4. **Response Formatting** + - NEVER use ASCII box-drawing characters (┌─┐│└─┘╔═╗║╚═╝ etc.) or ASCII art diagrams. + - Use markdown headings, bold, lists, and indentation for structure. + - Keep lines short — most users read on mobile. + - For architecture/flow, use arrow text: CLI → Pipeline → Adapters + + 5. **Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content.`, + identity, workspacePath, workspacePath, workspacePath, workspacePath, toolsSection, executingRule) } @@ -220,18 +386,24 @@ func (cb *ContextBuilder) buildToolsSection() string { } summaries := cb.tools.GetSummaries() + if len(summaries) == 0 { return "" } var sb strings.Builder + sb.WriteString("## Available Tools\n\n") + sb.WriteString( "**CRITICAL**: You MUST use tools to perform actions. Do NOT pretend to execute commands or schedule tasks.\n\n", ) + sb.WriteString("You have access to the following tools:\n\n") + for _, s := range summaries { sb.WriteString(s) + sb.WriteString("\n") } @@ -242,9 +414,11 @@ func (cb *ContextBuilder) BuildSystemPrompt() string { parts := []string{} // Core identity section + parts = append(parts, cb.getIdentity()) // Orchestration guidance — injected only when spawn tool is registered + if cb.tools != nil { if _, hasSpawn := cb.tools.Get("spawn"); hasSpawn { parts = append(parts, orchestrationGuidance) @@ -252,22 +426,31 @@ func (cb *ContextBuilder) BuildSystemPrompt() string { } // Bootstrap files + bootstrapContent := cb.LoadBootstrapFiles() + if bootstrapContent != "" { parts = append(parts, bootstrapContent) } // Skills - show summary, AI can read full content with read_file tool + skillsSummary := cb.skillsLoader.BuildSkillsSummary() + if skillsSummary != "" { parts = append(parts, fmt.Sprintf(`# Skills + + The following skills extend your capabilities. To use a skill, read its SKILL.md file using the read_file tool. + + %s`, skillsSummary)) } // Runtime status from tools (e.g., background processes) + if cb.tools != nil { if status := cb.tools.GetRuntimeStatus(); status != "" { parts = append(parts, status) @@ -275,55 +458,81 @@ The following skills extend your capabilities. To use a skill, read its SKILL.md } // Peer session coordination + if cb.peerNote != "" { parts = append(parts, "## Active Sessions\n\n"+cb.peerNote) } // Memory context + memoryContext := cb.memory.GetMemoryContext() + if memoryContext != "" { parts = append(parts, "# Memory\n\n"+memoryContext) } // Join with "---" separator + return strings.Join(parts, "\n\n---\n\n") } // BuildSystemPromptWithCache returns the cached system prompt if available + // and source files haven't changed, otherwise builds and caches it. + // Source file changes are detected via mtime checks (cheap stat calls). + func (cb *ContextBuilder) BuildSystemPromptWithCache() string { // Try read lock first — fast path when cache is valid + cb.systemPromptMutex.RLock() + if cb.cachedSystemPrompt != "" && !cb.sourceFilesChangedLocked() { result := cb.cachedSystemPrompt + cb.systemPromptMutex.RUnlock() + return result } + cb.systemPromptMutex.RUnlock() // Acquire write lock for building + cb.systemPromptMutex.Lock() + defer cb.systemPromptMutex.Unlock() // Double-check: another goroutine may have rebuilt while we waited + if cb.cachedSystemPrompt != "" && !cb.sourceFilesChangedLocked() { return cb.cachedSystemPrompt } // Snapshot the baseline (existence + max mtime) BEFORE building the prompt. + // This way cachedAt reflects the pre-build state: if a file is modified + // during BuildSystemPrompt, its new mtime will be > baseline.maxMtime, + // so the next sourceFilesChangedLocked check will correctly trigger a + // rebuild. The alternative (baseline after build) risks caching stale + // content with a too-new baseline, making the staleness invisible. + baseline := cb.buildCacheBaseline() + prompt := cb.BuildSystemPrompt() + cb.cachedSystemPrompt = prompt + cb.cachedAt = baseline.maxMtime + cb.existedAtCache = baseline.existed logger.DebugCF("agent", "System prompt cached", + map[string]any{ "length": len(prompt), }) @@ -332,95 +541,136 @@ func (cb *ContextBuilder) BuildSystemPromptWithCache() string { } // InvalidateCache clears the cached system prompt. + // Normally not needed because the cache auto-invalidates via mtime checks, + // but this is useful for tests or explicit reload commands. + func (cb *ContextBuilder) InvalidateCache() { cb.systemPromptMutex.Lock() + defer cb.systemPromptMutex.Unlock() cb.cachedSystemPrompt = "" + cb.cachedAt = time.Time{} + cb.existedAtCache = nil logger.DebugCF("agent", "System prompt cache invalidated", nil) } // sourcePaths returns the workspace source file paths tracked for cache + // invalidation (bootstrap files + memory). The skills directory is handled + // separately in sourceFilesChangedLocked because it requires both directory- + // level and recursive file-level mtime checks. + func (cb *ContextBuilder) sourcePaths() []string { // Include bootstrap files from all search directories (workDir, planWorkDir, workspace). + seen := map[string]bool{} + var paths []string + for _, spec := range bootstrapSpecs { var dirs []string + if spec.Scope == "global" { dirs = []string{cb.workspace} } else { dirs = cb.bootstrapProjectDirs() } + for _, dir := range dirs { p := filepath.Join(dir, spec.Name) + if !seen[p] { seen[p] = true + paths = append(paths, p) } } } + // Always track memory file. + memPath := filepath.Join(cb.workspace, "memory", "MEMORY.md") + if !seen[memPath] { paths = append(paths, memPath) } + return paths } // cacheBaseline holds the file existence snapshot and the latest observed + // mtime across all tracked paths. Used as the cache reference point. + type cacheBaseline struct { - existed map[string]bool + existed map[string]bool + maxMtime time.Time } // buildCacheBaseline records which tracked paths currently exist and computes + // the latest mtime across all tracked files + skills directory contents. + // Called under write lock when the cache is built. + func (cb *ContextBuilder) buildCacheBaseline() cacheBaseline { skillsDir := filepath.Join(cb.workspace, "skills") // All paths whose existence we track: source files + skills dir. + allPaths := append(cb.sourcePaths(), skillsDir) existed := make(map[string]bool, len(allPaths)) + var maxMtime time.Time for _, p := range allPaths { info, err := os.Stat(p) + existed[p] = err == nil + if err == nil && info.ModTime().After(maxMtime) { maxMtime = info.ModTime() } } // Walk skills files to capture their mtimes too. + // Use os.Stat (not d.Info) to match the stat method used in + // fileChangedSince / skillFilesModifiedSince for consistency. + _ = filepath.WalkDir(skillsDir, func(path string, d fs.DirEntry, walkErr error) error { if walkErr == nil && !d.IsDir() { if info, err := os.Stat(path); err == nil && info.ModTime().After(maxMtime) { maxMtime = info.ModTime() } } + return nil }) // If no tracked files exist yet (empty workspace), maxMtime is zero. + // Use a very old non-zero time so that: + // 1. cachedAt.IsZero() won't trigger perpetual rebuilds. + // 2. Any real file created afterwards has mtime > cachedAt, so it + // will be detected by fileChangedSince (unlike time.Now() which + // could race with a file whose mtime <= Now). + if maxMtime.IsZero() { maxMtime = time.Unix(1, 0) } @@ -429,18 +679,26 @@ func (cb *ContextBuilder) buildCacheBaseline() cacheBaseline { } // sourceFilesChangedLocked checks whether any workspace source file has been + // modified, created, or deleted since the cache was last built. + // + // IMPORTANT: The caller MUST hold at least a read lock on systemPromptMutex. + // Go's sync.RWMutex is not reentrant, so this function must NOT acquire the + // lock itself (it would deadlock when called from BuildSystemPromptWithCache + // which already holds RLock or Lock). + func (cb *ContextBuilder) sourceFilesChangedLocked() bool { if cb.cachedAt.IsZero() { return true } // Check tracked source files (bootstrap + memory). + for _, p := range cb.sourcePaths() { if cb.fileChangedSince(p) { return true @@ -448,19 +706,29 @@ func (cb *ContextBuilder) sourceFilesChangedLocked() bool { } // --- Skills directory (handled separately from sourcePaths) --- + // + // 1. Creation/deletion: tracked via existedAtCache, same as bootstrap files. + skillsDir := filepath.Join(cb.workspace, "skills") + if cb.fileChangedSince(skillsDir) { return true } // 2. Structural changes (add/remove entries inside the dir) are reflected + // in the directory's own mtime, which fileChangedSince already checks. + // + // 3. Content-only edits to files inside skills/ do NOT update the parent + // directory mtime on most filesystems, so we recursively walk to check + // individual file mtimes at any nesting depth. + if skillFilesModifiedSince(skillsDir, cb.cachedAt) { return true } @@ -469,93 +737,136 @@ func (cb *ContextBuilder) sourceFilesChangedLocked() bool { } // fileChangedSince returns true if a tracked source file has been modified, + // newly created, or deleted since the cache was built. + // + // Four cases: + // - existed at cache time, exists now -> check mtime + // - existed at cache time, gone now -> changed (deleted) + // - absent at cache time, exists now -> changed (created) + // - absent at cache time, gone now -> no change + func (cb *ContextBuilder) fileChangedSince(path string) bool { // Defensive: if existedAtCache was never initialized, treat as changed + // so the cache rebuilds rather than silently serving stale data. + if cb.existedAtCache == nil { return true } existedBefore := cb.existedAtCache[path] + info, err := os.Stat(path) + existsNow := err == nil if existedBefore != existsNow { return true // file was created or deleted } + if !existsNow { return false // didn't exist before, doesn't exist now } + return info.ModTime().After(cb.cachedAt) } // errWalkStop is a sentinel error used to stop filepath.WalkDir early. + // Using a dedicated error (instead of fs.SkipAll) makes the early-exit + // intent explicit and avoids the nilerr linter warning that would fire + // if the callback returned nil when its err parameter is non-nil. + var errWalkStop = errors.New("walk stop") // skillFilesModifiedSince recursively walks the skills directory and checks + // whether any file was modified after t. This catches content-only edits at + // any nesting depth (e.g. skills/name/docs/extra.md) that don't update + // parent directory mtimes. + func skillFilesModifiedSince(skillsDir string, t time.Time) bool { changed := false + err := filepath.WalkDir(skillsDir, func(path string, d fs.DirEntry, walkErr error) error { if walkErr == nil && !d.IsDir() { if info, statErr := os.Stat(path); statErr == nil && info.ModTime().After(t) { changed = true + return errWalkStop // stop walking } } + return nil }) + // errWalkStop is expected (early exit on first changed file). + // os.IsNotExist means the skills dir doesn't exist yet — not an error. + // Any other error is unexpected and worth logging. + if err != nil && !errors.Is(err, errWalkStop) && !os.IsNotExist(err) { logger.DebugCF("agent", "skills walk error", map[string]any{"error": err.Error()}) } + return changed } // BootstrapFileInfo describes a resolved bootstrap file. + type BootstrapFileInfo struct { - Name string `json:"name"` - Path string `json:"path"` // empty = not found + Name string `json:"name"` + + Path string `json:"path"` // empty = not found + Scope string `json:"scope"` // "project" or "global" } // bootstrapFileSpec defines the search scope for each bootstrap file. + type bootstrapFileSpec struct { - Name string + Name string + Scope string // "project" = workDir→planWorkDir→workspace, "global" = workspace only } var bootstrapSpecs = []bootstrapFileSpec{ {Name: "AGENTS.md", Scope: "project"}, + {Name: "IDENTITY.md", Scope: "project"}, + {Name: "SOUL.md", Scope: "global"}, + {Name: "USER.md", Scope: "global"}, } // bootstrapProjectDirs returns de-duplicated search directories for project-scoped files. + func (cb *ContextBuilder) bootstrapProjectDirs() []string { seen := map[string]bool{} + var dirs []string + for _, d := range []string{cb.workDir, cb.memory.GetPlanWorkDir(), cb.workspace} { if d != "" && !seen[d] { seen[d] = true + dirs = append(dirs, d) } } + return dirs } @@ -563,17 +874,22 @@ func (cb *ContextBuilder) LoadBootstrapFiles() string { projectDirs := cb.bootstrapProjectDirs() var sb strings.Builder + for _, spec := range bootstrapSpecs { var dirs []string + if spec.Scope == "global" { dirs = []string{cb.workspace} } else { dirs = projectDirs } + for _, dir := range dirs { filePath := filepath.Join(dir, spec.Name) + if data, err := os.ReadFile(filePath); err == nil { fmt.Fprintf(&sb, "## %s\n\n%s\n\n", spec.Name, data) + break } } @@ -583,44 +899,64 @@ func (cb *ContextBuilder) LoadBootstrapFiles() string { } // ResolveBootstrapPaths returns path resolution info for each bootstrap file + // using the same search logic as LoadBootstrapFiles. + func (cb *ContextBuilder) ResolveBootstrapPaths() []BootstrapFileInfo { projectDirs := cb.bootstrapProjectDirs() result := make([]BootstrapFileInfo, 0, len(bootstrapSpecs)) + for _, spec := range bootstrapSpecs { info := BootstrapFileInfo{Name: spec.Name, Scope: spec.Scope} + var dirs []string + if spec.Scope == "global" { dirs = []string{cb.workspace} } else { dirs = projectDirs } + for _, dir := range dirs { filePath := filepath.Join(dir, spec.Name) + if _, err := os.Stat(filePath); err == nil { info.Path = filePath + break } } + result = append(result, info) } + return result } // buildDynamicContext returns a short dynamic context string with per-request info. + // This changes every request (time, session) so it is NOT part of the cached prompt. + // LLM-side KV cache reuse is achieved by each provider adapter's native mechanism: + // - Anthropic: per-block cache_control (ephemeral) on the static SystemParts block + // - OpenAI / Codex: prompt_cache_key for prefix-based caching + // + // See: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching + // See: https://platform.openai.com/docs/guides/prompt-caching + func (cb *ContextBuilder) buildDynamicContext(channel, chatID string) string { now := time.Now().Format("2006-01-02 15:04 (Monday)") + rt := fmt.Sprintf("%s %s, Go %s", runtime.GOOS, runtime.GOARCH, runtime.Version()) var sb strings.Builder + fmt.Fprintf(&sb, "## Current Time\n%s\n\n## Runtime\n%s", now, rt) if channel != "" && chatID != "" { @@ -632,76 +968,119 @@ func (cb *ContextBuilder) buildDynamicContext(channel, chatID string) string { func (cb *ContextBuilder) BuildMessages( history []providers.Message, + summary string, + currentMessage string, + media []string, + channel, chatID string, ) []providers.Message { messages := []providers.Message{} // The static part (identity, bootstrap, skills, memory) is cached locally to + // avoid repeated file I/O and string building on every call (fixes issue #607). + // Dynamic parts (time, session, summary) are appended per request. + // Everything is sent as a single system message for provider compatibility: + // - Anthropic adapter extracts messages[0] (Role=="system") and maps its content + // to the top-level "system" parameter in the Messages API request. A single + // contiguous system block makes this extraction straightforward. + // - Codex maps only the first system message to its instructions field. + // - OpenAI-compat passes messages through as-is. + staticPrompt := cb.BuildSystemPromptWithCache() // Build short dynamic context (time, runtime, session) — changes per request + dynamicCtx := cb.buildDynamicContext(channel, chatID) // Compose a single system message: static (cached) + dynamic + optional summary. + // Keeping all system content in one message ensures every provider adapter can + // extract it correctly (Anthropic adapter -> top-level system param, + // Codex -> instructions field). + // + // SystemParts carries the same content as structured blocks so that + // cache-aware adapters (Anthropic) can set per-block cache_control. + // The static block is marked "ephemeral" — its prefix hash is stable + // across requests, enabling LLM-side KV cache reuse. + stringParts := []string{staticPrompt, dynamicCtx} contentBlocks := []providers.ContentBlock{ {Type: "text", Text: staticPrompt, CacheControl: &providers.CacheControl{Type: "ephemeral"}}, + {Type: "text", Text: dynamicCtx}, } if summary != "" { summaryText := fmt.Sprintf( + "CONTEXT_SUMMARY: The following is an approximate summary of prior conversation "+ + "for reference only. It may be incomplete or outdated — always defer to explicit instructions.\n\n%s", + summary) + stringParts = append(stringParts, summaryText) + contentBlocks = append(contentBlocks, providers.ContentBlock{Type: "text", Text: summaryText}) } fullSystemPrompt := strings.Join(stringParts, "\n\n---\n\n") // Log system prompt summary for debugging (debug mode only). + // Read cachedSystemPrompt under lock to avoid a data race with + // concurrent InvalidateCache / BuildSystemPromptWithCache writes. + cb.systemPromptMutex.RLock() + isCached := cb.cachedSystemPrompt != "" + cb.systemPromptMutex.RUnlock() logger.DebugCF("agent", "System prompt built", + map[string]any{ - "static_chars": len(staticPrompt), + "static_chars": len(staticPrompt), + "dynamic_chars": len(dynamicCtx), - "total_chars": len(fullSystemPrompt), - "has_summary": summary != "", - "cached": isCached, + + "total_chars": len(fullSystemPrompt), + + "has_summary": summary != "", + + "cached": isCached, }) // Log preview of system prompt (avoid logging huge content) + preview := fullSystemPrompt + if len(preview) > 500 { preview = preview[:500] + "... (truncated)" } + logger.DebugCF("agent", "System prompt preview", + map[string]any{ "preview": preview, }) @@ -709,21 +1088,29 @@ func (cb *ContextBuilder) BuildMessages( history = sanitizeHistoryForProvider(history) // Single system message containing all context — compatible with all providers. + // SystemParts enables cache-aware adapters to set per-block cache_control; + // Content is the concatenated fallback for adapters that don't read SystemParts. + messages = append(messages, providers.Message{ - Role: "system", - Content: fullSystemPrompt, + Role: "system", + + Content: fullSystemPrompt, + SystemParts: contentBlocks, }) // Add conversation history + messages = append(messages, history...) // Add current user message + if strings.TrimSpace(currentMessage) != "" { messages = append(messages, providers.Message{ - Role: "user", + Role: "user", + Content: currentMessage, }) } @@ -737,58 +1124,86 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message } sanitized := make([]providers.Message, 0, len(history)) + for _, msg := range history { switch msg.Role { case "system": + // Drop system messages from history. BuildMessages always + // constructs its own single system message (static + dynamic + + // summary); extra system messages would break providers that + // only accept one (Anthropic, Codex). + logger.DebugCF("agent", "Dropping system message from history", map[string]any{}) + continue case "tool": + if len(sanitized) == 0 { logger.DebugCF("agent", "Dropping orphaned leading tool message", map[string]any{}) + continue } + // Walk backwards to find the nearest assistant message, + // skipping over any preceding tool messages (multi-tool-call case). + foundAssistant := false + for i := len(sanitized) - 1; i >= 0; i-- { if sanitized[i].Role == "tool" { continue } + if sanitized[i].Role == "assistant" && len(sanitized[i].ToolCalls) > 0 { foundAssistant = true } + break } + if !foundAssistant { logger.DebugCF("agent", "Dropping orphaned tool message", map[string]any{}) + continue } + sanitized = append(sanitized, msg) case "assistant": + if len(msg.ToolCalls) > 0 { if len(sanitized) == 0 { logger.DebugCF("agent", "Dropping assistant tool-call turn at history start", map[string]any{}) + continue } + prev := sanitized[len(sanitized)-1] + if prev.Role != "user" && prev.Role != "tool" { logger.DebugCF( + "agent", + "Dropping assistant tool-call turn with invalid predecessor", + map[string]any{"prev_role": prev.Role}, ) + continue } } + sanitized = append(sanitized, msg) default: + sanitized = append(sanitized, msg) } } @@ -798,41 +1213,54 @@ func sanitizeHistoryForProvider(history []providers.Message) []providers.Message func (cb *ContextBuilder) AddToolResult( messages []providers.Message, + toolCallID, toolName, result string, ) []providers.Message { messages = append(messages, providers.Message{ - Role: "tool", - Content: result, + Role: "tool", + + Content: result, + ToolCallID: toolCallID, }) + return messages } func (cb *ContextBuilder) AddAssistantMessage( messages []providers.Message, + content string, + toolCalls []map[string]any, ) []providers.Message { msg := providers.Message{ - Role: "assistant", + Role: "assistant", + Content: content, } + // Always add assistant message, whether or not it has tool calls + messages = append(messages, msg) + return messages } // LoadSkill loads a skill by name, returning its content (with frontmatter stripped) and whether it was found. + func (cb *ContextBuilder) LoadSkill(name string) (string, bool) { return cb.skillsLoader.LoadSkill(name) } // ListSkills returns all available skills from all tiers. + func (cb *ContextBuilder) ListSkills() []skills.SkillInfo { return cb.skillsLoader.ListSkills() } // Memory returns the underlying MemoryStore for direct plan queries. + func (cb *ContextBuilder) Memory() *MemoryStore { return cb.memory } @@ -840,105 +1268,129 @@ func (cb *ContextBuilder) Memory() *MemoryStore { // ---------- Plan passthrough methods ---------- // ReadMemory reads the long-term memory (MEMORY.md). + func (cb *ContextBuilder) ReadMemory() string { return cb.memory.ReadLongTerm() } // WriteMemory writes content to the long-term memory file. + func (cb *ContextBuilder) WriteMemory(content string) error { return cb.memory.WriteLongTerm(content) } // ClearMemory removes the long-term memory file. + func (cb *ContextBuilder) ClearMemory() error { return cb.memory.ClearLongTerm() } // HasActivePlan returns true if MEMORY.md contains an active plan. + func (cb *ContextBuilder) HasActivePlan() bool { return cb.memory.HasActivePlan() } // GetPlanStatus returns the plan status: "interviewing", "executing", or "". + func (cb *ContextBuilder) GetPlanStatus() string { return cb.memory.GetPlanStatus() } // IsPlanComplete returns true if all steps in all phases are [x]. + func (cb *ContextBuilder) IsPlanComplete() bool { return cb.memory.IsPlanComplete() } // IsCurrentPhaseComplete returns true if all steps in the current phase are [x]. + func (cb *ContextBuilder) IsCurrentPhaseComplete() bool { return cb.memory.IsCurrentPhaseComplete() } // AdvancePhase increments the current phase number by 1. + func (cb *ContextBuilder) AdvancePhase() error { return cb.memory.AdvancePhase() } // SetCurrentPhase sets the current phase number to n. + func (cb *ContextBuilder) SetCurrentPhase(n int) error { return cb.memory.SetPhase(n) } // GetCurrentPhase returns the current phase number. + func (cb *ContextBuilder) GetCurrentPhase() int { return cb.memory.GetCurrentPhase() } // GetTotalPhases returns the total number of phases in the plan. + func (cb *ContextBuilder) GetTotalPhases() int { return cb.memory.GetTotalPhases() } // FormatPlanDisplay returns a user-facing display of the full plan. + func (cb *ContextBuilder) FormatPlanDisplay() string { return cb.memory.FormatPlanDisplay() } // MarkStep marks a step as done in the specified phase. + func (cb *ContextBuilder) MarkStep(phase, step int) error { return cb.memory.MarkStep(phase, step) } // AddStep appends a new step to the given phase. + func (cb *ContextBuilder) AddStep(phase int, desc string) error { return cb.memory.AddStep(phase, desc) } // ValidatePlanStructure validates plan structure for interview->review transition. + func (cb *ContextBuilder) ValidatePlanStructure() error { return cb.memory.ValidatePlanStructure() } // SetPlanStatus sets the plan status. + func (cb *ContextBuilder) SetPlanStatus(status string) error { return cb.memory.SetStatus(status) } // GetPlanWorkDir returns the WorkDir from the plan metadata, or "". + func (cb *ContextBuilder) GetPlanWorkDir() string { return cb.memory.GetPlanWorkDir() } // GetPlanTaskName returns the task description from the plan metadata, or "". + func (cb *ContextBuilder) GetPlanTaskName() string { return cb.memory.GetPlanTaskName() } // GetSkillsInfo returns information about loaded skills. + func (cb *ContextBuilder) GetSkillsInfo() map[string]any { allSkills := cb.skillsLoader.ListSkills() + skillNames := make([]string, 0, len(allSkills)) + for _, s := range allSkills { skillNames = append(skillNames, s.Name) } + return map[string]any{ - "total": len(allSkills), + "total": len(allSkills), + "available": len(allSkills), - "names": skillNames, + + "names": skillNames, } } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index afdaabc3b..73c825ae8 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1,7 +1,11 @@ // PicoClaw - Ultra-lightweight personal AI agent + // Inspired by and based on nanobot: https://github.com/HKUDS/nanobot + // License: MIT + // + // Copyright (c) 2026 PicoClaw contributors package agent @@ -41,107 +45,174 @@ import ( ) // activeTask tracks a running agent task for live status and intervention. + type activeTask struct { - Description string - Result string // LLM response summary for completion notification - Iteration int - MaxIter int - StartedAt time.Time - cancel context.CancelFunc - interrupt chan string // buffered 1, for user message injection - toolLog []toolLogEntry - lastError *toolLogEntry // sticky: most recent error, persists across iterations - projectDir string // detected from exec cd target (authoritative) - fileCommonDir string // LCP of file paths relative to workspace (fallback) - streamedChunks bool // true after onChunk fires at least once - messageContent string // last content sent by the message tool (for inclusion in completion) - mu sync.Mutex + Description string + + Result string // LLM response summary for completion notification + + Iteration int + + MaxIter int + + StartedAt time.Time + + cancel context.CancelFunc + + interrupt chan string // buffered 1, for user message injection + + toolLog []toolLogEntry + + lastError *toolLogEntry // sticky: most recent error, persists across iterations + + projectDir string // detected from exec cd target (authoritative) + + fileCommonDir string // LCP of file paths relative to workspace (fallback) + + streamedChunks bool // true after onChunk fires at least once + + messageContent string // last content sent by the message tool (for inclusion in completion) + + mu sync.Mutex } // toolLogEntry records a single tool call for the live terminal view. + type toolLogEntry struct { - Name string - ArgsSnip string // first ~80 chars of args - Result string // "✓ 4.9s" or "✗ 3.2s" + Name string + + ArgsSnip string // first ~80 chars of args + + Result string // "✓ 4.9s" or "✗ 3.2s" + ErrDetail string // non-empty on error — e.g. "Exit code: exit status 1" } // maxToolLogEntries limits the sliding window of tool log entries + // kept in memory and displayed in status messages. + const maxToolLogEntries = 5 // sessionSemaphore is a per-session mutex using a buffered channel. + type sessionSemaphore struct { ch chan struct{} } func newSessionSemaphore() *sessionSemaphore { s := &sessionSemaphore{ch: make(chan struct{}, 1)} + s.ch <- struct{}{} // initially unlocked + return s } type AgentLoop struct { - bus *bus.MessageBus - cfg *config.Config - registry *AgentRegistry - state *state.Manager - stats *stats.Tracker // nil when --stats not passed - running atomic.Bool - summarizing sync.Map - fallback *providers.FallbackChain - channelManager *channels.Manager - mediaStore media.MediaStore - providerCache map[string]providers.LLMProvider - planStartPending bool // set by /plan start to trigger LLM execution - planClearHistory bool // set by /plan start clear to wipe history on transition - sessionLocks sync.Map // sessionKey → *sessionSemaphore - activeTasks sync.Map // sessionKey → *activeTask - sessions *SessionTracker - lastSystemPrompt atomic.Value // string — last system prompt sent to LLM - promptDirty atomic.Bool // true = rebuild needed on next GetSystemPrompt read - OnStateChange func() // called on plan/session/skills mutations - OnUserMessage func() // called when a real user message is processed - saveConfig func(*config.Config) error + bus *bus.MessageBus + + cfg *config.Config + + registry *AgentRegistry + + state *state.Manager + + stats *stats.Tracker // nil when --stats not passed + + running atomic.Bool + + summarizing sync.Map + + fallback *providers.FallbackChain + + channelManager *channels.Manager + + mediaStore media.MediaStore + + providerCache map[string]providers.LLMProvider + + planStartPending bool // set by /plan start to trigger LLM execution + + planClearHistory bool // set by /plan start clear to wipe history on transition + + sessionLocks sync.Map // sessionKey → *sessionSemaphore + + activeTasks sync.Map // sessionKey → *activeTask + + sessions *SessionTracker + + lastSystemPrompt atomic.Value // string — last system prompt sent to LLM + + promptDirty atomic.Bool // true = rebuild needed on next GetSystemPrompt read + + OnStateChange func() // called on plan/session/skills mutations + + OnUserMessage func() // called when a real user message is processed + + saveConfig func(*config.Config) error + onHeartbeatThreadUpdate func(int) - orchBroadcaster *orch.Broadcaster // nil when --orchestration not set - orchReporter orch.AgentReporter // always non-nil (Noop when disabled) - done chan struct{} // closed by Close() to stop background goroutines + + orchBroadcaster *orch.Broadcaster // nil when --orchestration not set + + orchReporter orch.AgentReporter // always non-nil (Noop when disabled) + + done chan struct{} // closed by Close() to stop background goroutines } // processOptions configures how a message is processed + type processOptions struct { - SessionKey string // Session identifier for history/context - Channel string // Target channel for tool execution - ChatID string // Target chat ID for tool execution - UserMessage string // User message content (may include prefix) - HistoryMessage string // If set, save this to history instead of UserMessage (for skill compaction) + SessionKey string // Session identifier for history/context + + Channel string // Target channel for tool execution + + ChatID string // Target chat ID for tool execution + + UserMessage string // User message content (may include prefix) + + HistoryMessage string // If set, save this to history instead of UserMessage (for skill compaction) + DefaultResponse string // Response when LLM returns empty - EnableSummary bool // Whether to trigger summarization - SendResponse bool // Whether to send response via bus - NoHistory bool // If true, don't load session history (for heartbeat) - TaskID string // Unique task ID for background task status tracking - Background bool // If true, this is a background task (cron/heartbeat) — enables live task notifications - SystemMessage bool // If true, this is a system message (subagent result) — skip placeholder and plan nudge + + EnableSummary bool // Whether to trigger summarization + + SendResponse bool // Whether to send response via bus + + NoHistory bool // If true, don't load session history (for heartbeat) + + TaskID string // Unique task ID for background task status tracking + + Background bool // If true, this is a background task (cron/heartbeat) — enables live task notifications + + SystemMessage bool // If true, this is a system message (subagent result) — skip placeholder and plan nudge } const defaultResponse = "I've completed processing but have no response to give. Increase `max_tool_iterations` in config.json." func NewAgentLoop( cfg *config.Config, + msgBus *bus.MessageBus, + provider providers.LLMProvider, + enableStats ...bool, ) *AgentLoop { registry := NewAgentRegistry(cfg, provider) // Set up shared fallback chain + cooldown := providers.NewCooldownTracker() + fallbackChain := providers.NewFallbackChain(cooldown) // Create state manager using default agent's workspace for channel recording + defaultAgent := registry.GetDefaultAgent() + var stateManager *state.Manager + if defaultAgent != nil { stateManager = state.NewManager(defaultAgent.Workspace) } @@ -149,40 +220,61 @@ func NewAgentLoop( providerCache := make(map[string]providers.LLMProvider) // Create stats tracker if enabled + var statsTracker *stats.Tracker + if len(enableStats) > 0 && enableStats[0] && defaultAgent != nil { statsTracker = stats.NewTracker(defaultAgent.Workspace) } // Determine if orchestration broadcaster is needed (any agent has subagents enabled). + // Note: instance.go maps defaults.Orchestration → Subagents.Enabled, so --orchestration + // is automatically reflected here. + var orchBroadcaster *orch.Broadcaster + var orchReporter orch.AgentReporter = orch.Noop + for _, id := range registry.ListAgentIDs() { if a, ok := registry.GetAgent(id); ok && a.Subagents != nil && a.Subagents.Enabled { orchBroadcaster = orch.NewBroadcaster() + orchReporter = orchBroadcaster + break } } al := &AgentLoop{ - bus: msgBus, - cfg: cfg, - registry: registry, - state: stateManager, - stats: statsTracker, - summarizing: sync.Map{}, - fallback: fallbackChain, - providerCache: providerCache, - sessions: NewSessionTracker(), + bus: msgBus, + + cfg: cfg, + + registry: registry, + + state: stateManager, + + stats: statsTracker, + + summarizing: sync.Map{}, + + fallback: fallbackChain, + + providerCache: providerCache, + + sessions: NewSessionTracker(), + orchBroadcaster: orchBroadcaster, - orchReporter: orchReporter, - done: make(chan struct{}), + + orchReporter: orchReporter, + + done: make(chan struct{}), } // Register shared tools to all agents (needs al for reporter injection). + registerSharedTools(cfg, msgBus, registry, provider, al) go al.gcLoop() @@ -191,83 +283,117 @@ func NewAgentLoop( } // reporter returns the active AgentReporter (never nil). + func (al *AgentLoop) reporter() orch.AgentReporter { if al.orchReporter == nil { return orch.Noop } + return al.orchReporter } // SetOrchReporter wires a Broadcaster as the active reporter. + // Called from cmd_gateway.go when --orchestration is set. + // --orchestration なし → 呼ばれない → reporter() は Noop を返す。 + func (al *AgentLoop) SetOrchReporter(b *orch.Broadcaster) { al.orchBroadcaster = b + al.orchReporter = b } // GetOrchBroadcaster returns the concrete Broadcaster for miniapp wiring. + // Returns nil when orchestration is disabled. + func (al *AgentLoop) GetOrchBroadcaster() *orch.Broadcaster { return al.orchBroadcaster } func (al *AgentLoop) notifyStateChange() { al.promptDirty.Store(true) + if al.OnStateChange != nil { al.OnStateChange() } } // SetConfigSaver registers a callback used by slash commands that persist runtime config changes. + func (al *AgentLoop) SetConfigSaver(fn func(*config.Config) error) { al.saveConfig = fn } // SetHeartbeatThreadUpdater registers a callback to apply runtime heartbeat thread updates. + func (al *AgentLoop) SetHeartbeatThreadUpdater(fn func(int)) { al.onHeartbeatThreadUpdate = fn } // registerSharedTools registers tools that are shared across all agents (web, message, spawn). + func registerSharedTools( cfg *config.Config, + msgBus *bus.MessageBus, + registry *AgentRegistry, + provider providers.LLMProvider, + al *AgentLoop, ) { for _, agentID := range registry.ListAgentIDs() { agent, ok := registry.GetAgent(agentID) + if !ok { continue } // Web tools + searchTool, err := tools.NewWebSearchTool(tools.WebSearchToolOptions{ - BraveAPIKey: cfg.Tools.Web.Brave.APIKey, - BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, - BraveEnabled: cfg.Tools.Web.Brave.Enabled, - TavilyAPIKey: cfg.Tools.Web.Tavily.APIKey, - TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, - TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, - TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, + BraveAPIKey: cfg.Tools.Web.Brave.APIKey, + + BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, + + BraveEnabled: cfg.Tools.Web.Brave.Enabled, + + TavilyAPIKey: cfg.Tools.Web.Tavily.APIKey, + + TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, + + TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, + + TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, + DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, - DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, - PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey, + + DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, + + PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey, + PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, - PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, - Proxy: cfg.Tools.Web.Proxy, + + PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, + + Proxy: cfg.Tools.Web.Proxy, }) + if err != nil { logger.ErrorCF("agent", "Failed to create web search tool", map[string]any{ "agent_id": agentID, - "error": err.Error(), + + "error": err.Error(), }) } else if searchTool != nil { agent.Tools.Register(searchTool) + logger.InfoCF("agent", "Web search provider registered", map[string]any{ "agent_id": agentID, + "provider": searchTool.ProviderName(), }) } else { @@ -275,90 +401,149 @@ func registerSharedTools( "agent_id": agentID, }) } + fetchTool, err := tools.NewWebFetchToolWithProxy(50000, cfg.Tools.Web.Proxy) + if err != nil { logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{ "agent_id": agentID, - "error": err.Error(), + + "error": err.Error(), }) } else { agent.Tools.Register(fetchTool) } // Hardware tools (I2C, SPI) - Linux only, returns error on other platforms + agent.Tools.Register(tools.NewI2CTool()) + agent.Tools.Register(tools.NewSPITool()) // Message tool + messageTool := tools.NewMessageTool() + messageTool.SetSendCallback(func(channel, chatID, content string) error { pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + return msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{ Channel: channel, - ChatID: chatID, + + ChatID: chatID, + Content: content, }) }) + agent.Tools.Register(messageTool) // Skill discovery and installation tools + registryMgr := skills.NewRegistryManagerFromConfig(skills.RegistryConfig{ MaxConcurrentSearches: cfg.Tools.Skills.MaxConcurrentSearches, - ClawHub: skills.ClawHubConfig(cfg.Tools.Skills.Registries.ClawHub), + + ClawHub: skills.ClawHubConfig(cfg.Tools.Skills.Registries.ClawHub), }) + searchCache := skills.NewSearchCache( + cfg.Tools.Skills.SearchCache.MaxSize, + time.Duration(cfg.Tools.Skills.SearchCache.TTLSeconds)*time.Second, ) + agent.Tools.Register(tools.NewFindSkillsTool(registryMgr, searchCache)) + agent.Tools.Register(tools.NewInstallSkillTool(registryMgr, agent.Workspace)) // Spawn tool — only registered when orchestration is explicitly enabled. + if agent.Subagents != nil && agent.Subagents.Enabled { webSearchOpts := tools.WebSearchToolOptions{ - BraveAPIKey: cfg.Tools.Web.Brave.APIKey, - BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, - BraveEnabled: cfg.Tools.Web.Brave.Enabled, - TavilyAPIKey: cfg.Tools.Web.Tavily.APIKey, - TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, - TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, - TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, + BraveAPIKey: cfg.Tools.Web.Brave.APIKey, + + BraveMaxResults: cfg.Tools.Web.Brave.MaxResults, + + BraveEnabled: cfg.Tools.Web.Brave.Enabled, + + TavilyAPIKey: cfg.Tools.Web.Tavily.APIKey, + + TavilyBaseURL: cfg.Tools.Web.Tavily.BaseURL, + + TavilyMaxResults: cfg.Tools.Web.Tavily.MaxResults, + + TavilyEnabled: cfg.Tools.Web.Tavily.Enabled, + DuckDuckGoMaxResults: cfg.Tools.Web.DuckDuckGo.MaxResults, - DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, - PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey, + + DuckDuckGoEnabled: cfg.Tools.Web.DuckDuckGo.Enabled, + + PerplexityAPIKey: cfg.Tools.Web.Perplexity.APIKey, + PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, - PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, + + PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, } + subagentManager := tools.NewSubagentManager( + provider, + agent.Model, + agent.Workspace, + msgBus, + al.reporter(), + webSearchOpts, ) + subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) + // Wire session recorder for DAG persistence. + recorder := newSessionRecorder(agent.Sessions) + conductorKey := routing.BuildAgentMainSessionKey(agent.ID) + subagentManager.SetSessionRecorder(recorder, conductorKey) + agent.SubagentMgr = subagentManager + spawnTool := tools.NewSpawnTool(subagentManager) + currentAgentID := agentID + spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { return registry.CanSpawnSubagent(currentAgentID, targetAgentID) }) + agent.Tools.Register(spawnTool) + // Register blocking subagent tool alongside spawn + subagentTool := tools.NewSubagentTool(subagentManager) + agent.Tools.Register(subagentTool) + + // Register conductor-side escalation tools (answer questions, review plans) + + agent.Tools.Register(tools.NewAnswerSubagentTool(subagentManager)) + + agent.Tools.Register(tools.NewReviewSubagentPlanTool(subagentManager)) } // Update context builder with the complete tools registry + agent.ContextBuilder.SetToolsRegistry(agent.Tools) // Set orchestration mode if enabled + if agent.Subagents != nil && agent.Subagents.Enabled { agent.ContextBuilder.SetOrchestrationEnabled(true) } @@ -369,73 +554,103 @@ func (al *AgentLoop) Run(ctx context.Context) error { al.running.Store(true) // LLM work is dispatched to a background worker so the main loop + // stays free to handle slash commands (/skills, …) instantly, + // even while a long tool-call chain is running. + llmQueue := make(chan bus.InboundMessage, 10) + workerDone := make(chan struct{}) + go func() { defer close(workerDone) + al.llmWorker(ctx, llmQueue) }() + defer func() { close(llmQueue) + <-workerDone }() for al.running.Load() { select { case <-ctx.Done(): + return nil + default: } msg, ok := al.bus.ConsumeInbound(ctx) + if !ok { continue } // Echo commands sent from the Mini App so the user can see what was sent. + if msg.Metadata["source"] == "webapp" && msg.Metadata["echoed"] == "" && msg.Content != "" { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: msg.Channel, - ChatID: msg.ChatID, - Content: "via MiniApp: " + msg.Content, + Channel: msg.Channel, + + ChatID: msg.ChatID, + + Content: "via MiniApp: " + msg.Content, + SkipPlaceholder: true, }) } // Fast path: handle slash commands immediately without blocking the LLM worker. + if response, handled := al.handleCommand(ctx, msg); handled { if response != "" { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: msg.Channel, - ChatID: msg.ChatID, - Content: response, + Channel: msg.Channel, + + ChatID: msg.ChatID, + + Content: response, + SkipPlaceholder: true, }) } + // /plan start sets the flag — enqueue a synthetic message so + // the LLM worker actually begins executing the plan. + if al.planStartPending { al.planStartPending = false + clearHistory := al.planClearHistory + al.planClearHistory = false if clearHistory { if agent := al.registry.GetDefaultAgent(); agent != nil { agent.Sessions.SetHistory(msg.SessionKey, nil) + agent.Sessions.SetSummary(msg.SessionKey, "") + _ = agent.Sessions.Save(msg.SessionKey) } } // Activate worktree for the session's plan execution + if agent := al.registry.GetDefaultAgent(); agent != nil { taskName := agent.ContextBuilder.Memory().GetPlanTaskName() + if taskName == "" { taskName = "plan-execution" } + planDir := agent.ContextBuilder.GetPlanWorkDir() + if wt, err := agent.ActivateWorktree(msg.SessionKey, taskName, planDir); err != nil { logger.WarnCF("agent", "Worktree activation skipped", map[string]any{"error": err.Error()}) } else { @@ -444,31 +659,44 @@ func (al *AgentLoop) Run(ctx context.Context) error { } syntheticMeta := map[string]string{"echoed": "1"} + for k, v := range msg.Metadata { if k != "source" { syntheticMeta[k] = v } } + select { case llmQueue <- bus.InboundMessage{ - Channel: msg.Channel, - ChatID: msg.ChatID, - SenderID: msg.SenderID, + Channel: msg.Channel, + + ChatID: msg.ChatID, + + SenderID: msg.SenderID, + SessionKey: msg.SessionKey, - Content: "The plan has been approved. Begin executing.", - Metadata: syntheticMeta, + + Content: "The plan has been approved. Begin executing.", + + Metadata: syntheticMeta, }: + case <-ctx.Done(): + return nil } } + continue } // Dispatch to LLM worker + select { case llmQueue <- msg: + case <-ctx.Done(): + return nil } } @@ -477,6 +705,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { } // llmWorker processes LLM messages sequentially in a background goroutine. + func (al *AgentLoop) llmWorker(ctx context.Context, queue <-chan bus.InboundMessage) { for msg := range queue { if ctx.Err() != nil { @@ -490,7 +719,9 @@ func (al *AgentLoop) llmWorker(ctx context.Context, queue <-chan bus.InboundMess if response != "" { alreadySent := false + defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent != nil { if tool, ok := defaultAgent.Tools.Get("message"); ok { if mt, ok := tool.(*tools.MessageTool); ok { @@ -502,7 +733,9 @@ func (al *AgentLoop) llmWorker(ctx context.Context, queue <-chan bus.InboundMess if !alreadySent { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ Channel: msg.Channel, - ChatID: msg.ChatID, + + ChatID: msg.ChatID, + Content: response, }) } @@ -515,17 +748,24 @@ func (al *AgentLoop) Stop() { } // Close releases resources held by the loop (e.g. flushes write-behind stats + // and dirty session data). Should be called during graceful shutdown. + func (al *AgentLoop) Close() { select { case <-al.done: + // already closed + default: + close(al.done) } + if al.stats != nil { al.stats.Close() } + for _, agentID := range al.registry.ListAgentIDs() { if agent, ok := al.registry.GetAgent(agentID); ok { agent.Sessions.Close() @@ -534,30 +774,43 @@ func (al *AgentLoop) Close() { } // gcLoop periodically cleans up stale sessionLock entries. + func (al *AgentLoop) gcLoop() { ticker := time.NewTicker(30 * time.Minute) + defer ticker.Stop() + for { select { case <-ticker.C: + al.gcSessionLocks() + case <-al.done: + return } } } // gcSessionLocks removes unlocked (idle) sessionSemaphore entries from the map. + func (al *AgentLoop) gcSessionLocks() { al.sessionLocks.Range(func(key, val any) bool { sem := val.(*sessionSemaphore) + select { case <-sem.ch: + // Was unlocked — safe to remove + al.sessionLocks.Delete(key) + default: + // Currently locked — in use, keep } + return true }) } @@ -575,72 +828,102 @@ func (al *AgentLoop) SetChannelManager(cm *channels.Manager) { } // resolveProvider returns the LLMProvider for the given provider/model pair. + // It caches created providers by "provider/model" key so each combination is + // only resolved once. Looks up model_list first (new format), then falls back + // to the legacy providers section via CreateProviderByName. + func (al *AgentLoop) resolveProvider( providerName, modelName string, + fallback providers.LLMProvider, ) providers.LLMProvider { key := strings.ToLower(providerName + "/" + modelName) + if key == "/" { return fallback } + if p, ok := al.providerCache[key]; ok { return p } // Try model_list first (new config format). + if mc := al.cfg.FindModelConfigByRef(providerName, modelName); mc != nil { p, _, err := providers.CreateProviderFromConfig(mc) + if err == nil { al.providerCache[key] = p + return p } + logger.WarnCF("agent", "Failed to create provider from model_list, trying legacy", + map[string]any{"provider": providerName, "model": modelName, "error": err.Error()}) } // Fall back to legacy providers section. + p, err := providers.CreateProviderByName(al.cfg, providerName) if err != nil { logger.WarnCF("agent", "Failed to create provider for fallback, using primary", + map[string]any{"provider": providerName, "error": err.Error()}) + return fallback } + al.providerCache[key] = p + return p } // SetMediaStore injects a MediaStore for media lifecycle management. + func (al *AgentLoop) SetMediaStore(s media.MediaStore) { al.mediaStore = s } // inferMediaType determines the media type ("image", "audio", "video", "file") + // from a filename and MIME content type. + func inferMediaType(filename, contentType string) string { ct := strings.ToLower(contentType) + fn := strings.ToLower(filename) if strings.HasPrefix(ct, "image/") { return "image" } + if strings.HasPrefix(ct, "audio/") || ct == "application/ogg" { return "audio" } + if strings.HasPrefix(ct, "video/") { return "video" } // Fallback: infer from extension + ext := filepath.Ext(fn) + switch ext { case ".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp", ".svg": + return "image" + case ".mp3", ".wav", ".ogg", ".m4a", ".flac", ".aac", ".wma", ".opus": + return "audio" + case ".mp4", ".avi", ".mov", ".webm", ".mkv": + return "video" } @@ -648,30 +931,40 @@ func inferMediaType(filename, contentType string) string { } // 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 { if al.state == nil { return nil } + return al.state.SetLastChannel(channel) } // RecordLastChatID records the last active chat ID for this workspace. + // This uses the atomic state save mechanism to prevent data loss on crash. + func (al *AgentLoop) RecordLastChatID(chatID string) error { if al.state == nil { return nil } + return al.state.SetLastChatID(chatID) } // RecordLastHeartbeatTarget records the latest heartbeat-safe destination. + // This is intentionally separate from LastChannel so heartbeat routing can be + // reasoned about and evolved without breaking generic last-activity tracking. + func (al *AgentLoop) RecordLastHeartbeatTarget(target string) error { if al.state == nil { return nil } + return al.state.SetLastHeartbeatTarget(target) } @@ -681,14 +974,20 @@ func (al *AgentLoop) ProcessDirect(ctx context.Context, content, sessionKey stri func (al *AgentLoop) ProcessDirectWithChannel( ctx context.Context, + content, sessionKey, channel, chatID string, ) (string, error) { msg := bus.InboundMessage{ - Channel: channel, - SenderID: "cron", - ChatID: chatID, - Content: content, + Channel: channel, + + SenderID: "cron", + + ChatID: chatID, + + Content: content, + SessionKey: sessionKey, + Metadata: map[string]string{ "background": "true", }, @@ -698,134 +997,194 @@ func (al *AgentLoop) ProcessDirectWithChannel( } // ProcessHeartbeat processes a heartbeat request without session history. + // Each heartbeat is independent and doesn't accumulate context. + func (al *AgentLoop) ProcessHeartbeat(ctx context.Context, content, channel, chatID string) (string, error) { agent := al.registry.GetDefaultAgent() + if agent == nil { return "", fmt.Errorf("no default agent for heartbeat") } + heartbeatThreadID := 0 + if al.cfg != nil { heartbeatThreadID = al.cfg.Channels.Telegram.HeartbeatThreadID } + heartbeatChatID := al.withTelegramThread(channel, chatID, heartbeatThreadID) + return al.runAgentLoop(ctx, agent, processOptions{ - SessionKey: "heartbeat", - Channel: channel, - ChatID: heartbeatChatID, - UserMessage: content, + SessionKey: "heartbeat", + + Channel: channel, + + ChatID: heartbeatChatID, + + UserMessage: content, + DefaultResponse: defaultResponse, - EnableSummary: false, - SendResponse: false, - NoHistory: true, // Don't load session history for heartbeat - Background: true, // Enable live task notifications on Telegram + + EnableSummary: false, + + SendResponse: false, + + NoHistory: true, // Don't load session history for heartbeat + + Background: true, // Enable live task notifications on Telegram + }) } func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) (string, error) { // Add message preview to log (show full content for error messages) + var logContent string + if strings.Contains(msg.Content, "Error:") || strings.Contains(msg.Content, "error") { logContent = msg.Content // Full content for errors } else { logContent = utils.Truncate(msg.Content, 80) } + logger.InfoCF("agent", fmt.Sprintf("Processing message from %s:%s: %s", msg.Channel, msg.SenderID, logContent), + map[string]any{ - "channel": msg.Channel, - "chat_id": msg.ChatID, - "sender_id": msg.SenderID, + "channel": msg.Channel, + + "chat_id": msg.ChatID, + + "sender_id": msg.SenderID, + "session_key": msg.SessionKey, }) // Handle reply-based intervention for active tasks + if taskID, ok := msg.Metadata["task_id"]; ok && taskID != "" { if val, found := al.activeTasks.Load(taskID); found { task := val.(*activeTask) + content := strings.TrimSpace(msg.Content) + lower := strings.ToLower(content) // Check for stop keywords + stopKeywords := []string{ "stop", "cancel", "abort", + "停止", "中止", "やめて", //nolint:gosmopolitan // intentional CJK stop words + } + isStop := false + for _, kw := range stopKeywords { if lower == kw { isStop = true + break } } if isStop { task.cancel() + logger.InfoCF("agent", "Task canceled by user intervention", + map[string]any{"task_id": taskID}) + return "Task canceled.", nil } // Inject message into interrupt channel for the tool loop + select { case task.interrupt <- content: + logger.InfoCF("agent", "User intervention queued", + map[string]any{"task_id": taskID, "content": utils.Truncate(content, 80)}) + default: + logger.WarnCF("agent", "Interrupt channel full, message dropped", + map[string]any{"task_id": taskID}) } + return "Intervention sent to running task.", nil } + // Task not found — fall through to normal processing } // Route system messages to processSystemMessage + if msg.Channel == "system" { return al.processSystemMessage(ctx, msg) } // Notify listeners that a real user message arrived (e.g. reset heartbeat suppression) + if al.OnUserMessage != nil { al.OnUserMessage() } // Expand /skill command: inject SKILL.md content into message, then continue to LLM + var expansionCompact string + if expanded, compact, ok := al.expandSkillCommand(msg); ok { msg.Content = expanded + expansionCompact = compact } // Expand /plan <task>: write interview seed, rewrite for LLM interview + if expanded, compact, ok := al.expandPlanCommand(msg); ok { msg.Content = expanded + expansionCompact = compact } // Check for commands + if response, handled := al.handleCommand(ctx, msg); handled { return response, nil } // Route to determine agent and session key + route := al.registry.ResolveRoute(routing.RouteInput{ - Channel: msg.Channel, - AccountID: msg.Metadata["account_id"], - Peer: extractPeer(msg), + Channel: msg.Channel, + + AccountID: msg.Metadata["account_id"], + + Peer: extractPeer(msg), + ParentPeer: extractParentPeer(msg), - GuildID: msg.Metadata["guild_id"], - TeamID: msg.Metadata["team_id"], + + GuildID: msg.Metadata["guild_id"], + + TeamID: msg.Metadata["team_id"], }) agent, ok := al.registry.GetAgent(route.AgentID) + if !ok { agent = al.registry.GetDefaultAgent() } + if agent == nil { return "", fmt.Errorf("no agent available for route (agent_id=%s)", route.AgentID) } // Reset message-tool state for this round so we don't skip publishing due to a previous round. + if tool, ok := agent.Tools.Get("message"); ok { if mt, ok := tool.(tools.ContextualTool); ok { mt.SetContext(msg.Channel, msg.ChatID) @@ -833,28 +1192,41 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) } // Use routed session key, but honor ANY pre-set session key (for ProcessDirect/cron) + sessionKey := route.SessionKey + if msg.SessionKey != "" { sessionKey = msg.SessionKey } logger.InfoCF("agent", "Routed message", + map[string]any{ - "agent_id": agent.ID, + "agent_id": agent.ID, + "session_key": sessionKey, - "matched_by": route.MatchedBy, + + "matched_by": route.MatchedBy, }) return al.runAgentLoop(ctx, agent, processOptions{ - SessionKey: sessionKey, - Channel: msg.Channel, - ChatID: msg.ChatID, - UserMessage: msg.Content, - HistoryMessage: expansionCompact, + SessionKey: sessionKey, + + Channel: msg.Channel, + + ChatID: msg.ChatID, + + UserMessage: msg.Content, + + HistoryMessage: expansionCompact, + DefaultResponse: defaultResponse, - EnableSummary: true, - SendResponse: false, - Background: msg.Metadata["background"] == "true", + + EnableSummary: true, + + SendResponse: false, + + Background: msg.Metadata["background"] == "true", }) } @@ -864,93 +1236,140 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe } logger.InfoCF("agent", "Processing system message", + map[string]any{ "sender_id": msg.SenderID, - "chat_id": msg.ChatID, + + "chat_id": msg.ChatID, }) // Parse origin channel from chat_id (format: "channel:chat_id") + var originChannel, originChatID string + if idx := strings.Index(msg.ChatID, ":"); idx > 0 { originChannel = msg.ChatID[:idx] + originChatID = msg.ChatID[idx+1:] } else { originChannel = "cli" + originChatID = msg.ChatID } // Extract subagent result from message content + // Format: "Task 'label' completed.\n\nResult:\n<actual content>" + content := msg.Content + if idx := strings.Index(content, "Result:\n"); idx >= 0 { content = content[idx+8:] // Extract just the result part } // Skip internal channels - only log, don't send to user + if constants.IsInternalChannel(originChannel) { logger.InfoCF("agent", "Subagent completed (internal channel)", + map[string]any{ - "sender_id": msg.SenderID, + "sender_id": msg.SenderID, + "content_len": len(content), - "channel": originChannel, + + "channel": originChannel, }) + return "", nil } // Inject subagent result into session history without running a full LLM loop. + // The conductor will see the result on its next turn. This avoids: + // - Flooding the chat with a response for every subagent completion + // - Consuming the Telegram "Thinking..." placeholder + // - Wasting LLM tokens on processing each result individually + agent := al.registry.GetDefaultAgent() + if agent == nil { return "", fmt.Errorf("no default agent for system message") } sessionKey := routing.BuildAgentMainSessionKey(agent.ID) + historyMsg := fmt.Sprintf("[System: %s] %s", msg.SenderID, msg.Content) + // Write as TurnReport to the store for DAG tracking, with legacy fallback. + subagentSessionKey := routing.BuildSubagentSessionKey(extractTaskID(msg.SenderID)) + store := agent.Sessions.Store() + reportTurn := &session.Turn{ - Kind: session.TurnReport, + Kind: session.TurnReport, + OriginKey: subagentSessionKey, - Author: msg.SenderID, - Messages: []providers.Message{{Role: "user", Content: historyMsg}}, + + Author: msg.SenderID, + + Messages: []providers.Message{{Role: "user", Content: historyMsg}}, } + if err := store.Append(sessionKey, reportTurn); err != nil { logger.ErrorCF("agent", "Failed to record report turn, falling back to legacy", + map[string]any{"error": err.Error()}) + agent.Sessions.AddMessage(sessionKey, "user", historyMsg) + agent.Sessions.MarkDirty(sessionKey) } else { // Update in-memory cache so conductor sees the message on next turn. + agent.Sessions.AddFullMessage(sessionKey, providers.Message{Role: "user", Content: historyMsg}) + agent.Sessions.AdvanceStored(sessionKey, 1) } // Send a brief notification (SkipPlaceholder to avoid corrupting status messages) + label := msg.SenderID + if idx := strings.LastIndex(label, ":"); idx >= 0 { label = label[idx+1:] } + notification := formatSubagentCompletion(label, msg.Metadata) + subagentThreadID := 0 + if al.cfg != nil { subagentThreadID = al.cfg.Channels.Telegram.SubagentThreadID } + notifyChatID := al.withTelegramThread(originChannel, originChatID, subagentThreadID) + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: originChannel, - ChatID: notifyChatID, - Content: notification, + Channel: originChannel, + + ChatID: notifyChatID, + + Content: notification, + SkipPlaceholder: true, }) logger.InfoCF("agent", "Subagent result injected into session history", + map[string]any{ - "sender_id": msg.SenderID, + "sender_id": msg.SenderID, + "session_key": sessionKey, + "content_len": len(content), }) @@ -958,24 +1377,34 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe } // extractTaskID extracts the task ID from a sender ID like "subagent:subagent-1". + func extractTaskID(senderID string) string { if idx := strings.LastIndex(senderID, ":"); idx >= 0 { return senderID[idx+1:] } + return senderID } // formatSubagentCompletion builds the user-facing notification for a completed subagent. + // If metadata contains duration_ms and tool_calls it produces e.g.: + // + // "📋 scout-1 completed (3.2s, 5 tool calls)." + // + // Without metadata it falls back to the plain "📋 scout-1 completed." format. + func formatSubagentCompletion(label string, metadata map[string]string) string { if len(metadata) == 0 { return fmt.Sprintf("📋 %s completed.", label) } + durationMs, _ := strconv.ParseInt(metadata["duration_ms"], 10, 64) + toolCalls, _ := strconv.Atoi(metadata["tool_calls"]) if durationMs <= 0 && toolCalls <= 0 { @@ -983,9 +1412,11 @@ func formatSubagentCompletion(label string, metadata map[string]string) string { } parts := make([]string, 0, 2) + if durationMs > 0 { parts = append(parts, formatDurationMs(durationMs)) } + if toolCalls > 0 { if toolCalls == 1 { parts = append(parts, "1 tool call") @@ -993,25 +1424,35 @@ func formatSubagentCompletion(label string, metadata map[string]string) string { parts = append(parts, fmt.Sprintf("%d tool calls", toolCalls)) } } + return fmt.Sprintf("📋 %s completed (%s).", label, strings.Join(parts, ", ")) } // formatDurationMs converts milliseconds to a human-readable duration string. + // Examples: 800 → "0.8s", 1200 → "1.2s", 65000 → "1m5s", 3661000 → "61m1s". + func formatDurationMs(ms int64) string { if ms < 1000 { return fmt.Sprintf("%dms", ms) } + totalSec := ms / 1000 + if totalSec < 60 { tenths := (ms % 1000) / 100 + return fmt.Sprintf("%d.%ds", totalSec, tenths) } + mins := totalSec / 60 + sec := totalSec % 60 + if sec == 0 { return fmt.Sprintf("%dm", mins) } + return fmt.Sprintf("%dm%ds", mins, sec) } @@ -1021,9 +1462,11 @@ func (al *AgentLoop) withTelegramThread(channel, chatID string, threadID int) st } baseChatID := chatID + if slash := strings.Index(baseChatID, "/"); slash >= 0 { baseChatID = baseChatID[:slash] } + if baseChatID == "" { return chatID } @@ -1032,114 +1475,160 @@ func (al *AgentLoop) withTelegramThread(channel, chatID string, threadID int) st } // acquireSessionLock gets or creates a per-session semaphore and acquires it. + // Returns false if the context is canceled before the lock is acquired. + func (al *AgentLoop) acquireSessionLock(ctx context.Context, sessionKey string) bool { val, _ := al.sessionLocks.LoadOrStore(sessionKey, newSessionSemaphore()) + sem := val.(*sessionSemaphore) + select { case <-sem.ch: + return true + case <-ctx.Done(): + return false } } // releaseSessionLock releases the per-session semaphore. + func (al *AgentLoop) releaseSessionLock(sessionKey string) { if val, ok := al.sessionLocks.Load(sessionKey); ok { sem := val.(*sessionSemaphore) + sem.ch <- struct{}{} } } // runAgentLoop is the core message processing logic. + func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) { // -1. Acquire per-session lock to prevent concurrent access on the same session + if !al.acquireSessionLock(ctx, opts.SessionKey) { return "", fmt.Errorf("context canceled while waiting for session lock") } + defer al.releaseSessionLock(opts.SessionKey) // Report session lifecycle to canvas. + al.reporter().ReportSpawn(opts.SessionKey, opts.Channel, opts.UserMessage) + defer al.reporter().ReportGC(opts.SessionKey, "completed") // -0. Create cancelable child context and register active task + taskCtx, taskCancel := context.WithCancel(ctx) + defer taskCancel() task := &activeTask{ Description: utils.Truncate(opts.UserMessage, 80), - MaxIter: agent.MaxIterations, - StartedAt: time.Now(), - cancel: taskCancel, - interrupt: make(chan string, 1), + + MaxIter: agent.MaxIterations, + + StartedAt: time.Now(), + + cancel: taskCancel, + + interrupt: make(chan string, 1), } // Guarantee heartbeat worktree cleanup on ALL exit paths (error, panic, normal). + // Wait for spawned subagents first so they aren't killed mid-flight. + // After auto-commit, attempt to merge the worktree branch into main. + defer func() { if opts.Background { if agent.SubagentMgr != nil { agent.SubagentMgr.WaitAll(35 * time.Minute) // slightly above spawnTimeout } + wt := agent.GetWorktree(opts.SessionKey) + if wt != nil { // 1. Auto-commit uncommitted changes in worktree + if git.HasUncommittedChanges(wt.Path) { _ = git.AutoCommit(wt.Path, "heartbeat: auto-save") } // 2. Check if there are unique commits worth merging + repoRoot := git.FindRepoRoot(agent.Workspace) + ahead := git.CommitsAhead(repoRoot, wt.BaseBranch, wt.Branch) if ahead > 0 && repoRoot != "" { // 3. Try fast-forward merge into base branch + mr := git.MergeWorktreeBranch(repoRoot, wt) // 4. Notify based on merge result + if !constants.IsInternalChannel(opts.Channel) { cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) + if mr.Merged { _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ Channel: opts.Channel, - ChatID: opts.ChatID, + + ChatID: opts.ChatID, + Content: fmt.Sprintf("Heartbeat: merged %d commit(s) to %s.", + ahead, wt.BaseBranch), }) } else if mr.Conflict { _ = al.bus.PublishOutbound(cleanupCtx, bus.OutboundMessage{ Channel: opts.Channel, - ChatID: opts.ChatID, + + ChatID: opts.ChatID, + Content: fmt.Sprintf("Heartbeat: merge conflict on branch `%s` — manual merge needed.", + mr.Branch), }) } + cleanupCancel() } } // 5. Dispose worktree (branch auto-deleted if merged, kept if conflict) + agent.DeactivateWorktree(opts.SessionKey, "", false) } } }() // For background tasks (cron/heartbeat), generate a TaskID and send notification + isBackgroundTask := opts.Background && al.state != nil + if isBackgroundTask && opts.TaskID == "" { opts.TaskID = fmt.Sprintf("task-%s-%d", opts.SessionKey, time.Now().UnixMilli()) // Determine notification channel: use opts.Channel if already a real channel, + // otherwise resolve from last active channel + notifyChannel := opts.Channel + notifyChatID := opts.ChatID + if constants.IsInternalChannel(notifyChannel) || notifyChannel == "" { if lastChannel := al.state.GetLastChannel(); lastChannel != "" { if idx := strings.Index(lastChannel, ":"); idx > 0 { notifyChannel = lastChannel[:idx] + notifyChatID = lastChannel[idx+1:] } } @@ -1147,111 +1636,169 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt if notifyChannel != "" && notifyChatID != "" && !constants.IsInternalChannel(notifyChannel) { // Override opts channel/chatID for status updates + opts.Channel = notifyChannel + opts.ChatID = notifyChatID // Send initial task notification + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: notifyChannel, - ChatID: notifyChatID, - Content: fmt.Sprintf("\U0001F916 Background task started\n%s", task.Description), + Channel: notifyChannel, + + ChatID: notifyChatID, + + Content: fmt.Sprintf("\U0001F916 Background task started\n%s", task.Description), + IsTaskStatus: true, - TaskID: opts.TaskID, + + TaskID: opts.TaskID, }) } } // Shared variable for capturing LLM's final response. The defer below reads it + // to include the response in the task completion message. + var finalContent string // Use TaskID as key if available (for background tasks), else sessionKey + taskKey := opts.SessionKey + if opts.TaskID != "" { taskKey = opts.TaskID } + al.activeTasks.Store(taskKey, task) + defer func() { al.activeTasks.Delete(taskKey) // Publish final task status on completion for background tasks. + // Include finalContent so the LLM response appears in the same bubble + // as the completion status, avoiding duplicate messages. + if opts.TaskID != "" { elapsed := time.Since(task.StartedAt) + completionMsg := fmt.Sprintf("\u2705 Task completed (%.1fs)", elapsed.Seconds()) // Determine the best content to show in the completion bubble. + // Priority: message tool content > finalContent > task.Result + task.mu.Lock() + msgContent := task.messageContent + task.mu.Unlock() var resultContent string + switch { case msgContent != "": + // The message tool already sent this to the user via the + // task bubble; re-include it so the completion doesn't erase it. + resultContent = msgContent + case finalContent != "" && finalContent != defaultResponse && finalContent != "HEARTBEAT_OK": + resultContent = finalContent + default: + summary := task.Result + if summary == "" { summary = task.Description } + resultContent = summary } if resultContent != "" { combined := completionMsg + "\n\n" + resultContent + if len([]rune(combined)) <= 4096 { completionMsg = combined } else { // Too long for one bubble: send header as task status, + // body as regular message (auto-split by SplitMessage). + doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second) - _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: completionMsg, - IsTaskStatus: true, - TaskID: opts.TaskID, - Final: true, - }) + _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ Channel: opts.Channel, - ChatID: opts.ChatID, + + ChatID: opts.ChatID, + + Content: completionMsg, + + IsTaskStatus: true, + + TaskID: opts.TaskID, + + Final: true, + }) + + _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ + Channel: opts.Channel, + + ChatID: opts.ChatID, + Content: resultContent, }) + doneCancel() + return } } + doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second) + _ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: completionMsg, + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: completionMsg, + IsTaskStatus: true, - TaskID: opts.TaskID, - Final: true, + + TaskID: opts.TaskID, + + Final: true, }) + doneCancel() } }() // Replace ctx with the cancelable child context + ctx = taskCtx // 0. Record last channel for heartbeat notifications (skip internal channels) + if opts.Channel != "" && opts.ChatID != "" { // Don't record internal channels (cli, system, subagent) + if !constants.IsInternalChannel(opts.Channel) { channelKey := fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID) + if err := al.RecordLastChannel(channelKey); err != nil { logger.WarnCF("agent", "Failed to record last channel", map[string]any{"error": err.Error()}) } + if err := al.RecordLastHeartbeatTarget(channelKey); err != nil { logger.WarnCF("agent", "Failed to record last heartbeat target", map[string]any{"error": err.Error()}) } @@ -1259,31 +1806,47 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt } // 1. Update tool contexts + al.updateToolContexts(agent, opts.Channel, opts.ChatID) // 1-bis. For background tasks that don't send a final response (e.g. heartbeat), + // redirect the message tool to publish as IsTaskStatus so its output lands in + // the same bubble as the task status instead of creating a separate message. + if opts.Background && !opts.SendResponse && opts.TaskID != "" { if tool, ok := agent.Tools.Get("message"); ok { if mt, ok := tool.(*tools.MessageTool); ok { taskID := opts.TaskID + mt.SetSendCallback(func(channel, chatID, content string) error { // Capture the message tool's content so the completion + // defer can include it instead of losing it to an overwrite. + if task != nil { task.mu.Lock() + task.messageContent = content + task.mu.Unlock() } + pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + return al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{ - Channel: channel, - ChatID: chatID, - Content: content, + Channel: channel, + + ChatID: chatID, + + Content: content, + IsTaskStatus: true, - TaskID: taskID, + + TaskID: taskID, }) }) } @@ -1291,9 +1854,13 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt } // 1a. Set session-specific working directory for bootstrap file lookup. + // Prefer the tool-detected project directory (touch_dir) from the session tracker, + // resolved as an absolute path under workspace. Fall back to worktree or workspace. + if active := al.sessions.ListActive(); len(active) > 0 && active[0].SessionKey == opts.SessionKey && + active[0].TouchDir != "" { agent.ContextBuilder.SetWorkDir(filepath.Join(agent.Workspace, active[0].TouchDir)) } else { @@ -1301,227 +1868,334 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt } // 1b. Inject peer session awareness into system prompt + projectPath := agent.ContextBuilder.GetPlanWorkDir() + if projectPath == "" { projectPath = agent.Workspace } + peers := al.sessions.GetPeerPurposes(opts.SessionKey, projectPath) + if len(peers) > 0 { var peerNote strings.Builder + peerNote.WriteString("Other sessions working on this project:\n") + for _, p := range peers { peerNote.WriteString(fmt.Sprintf("- %s: %s (branch: %s)\n", p.SessionKey, p.Purpose, p.Branch)) } + peerNote.WriteString("\nAvoid conflicting changes with these sessions.") + agent.ContextBuilder.SetPeerNote(peerNote.String()) } else { agent.ContextBuilder.SetPeerNote("") } // 2. Build messages (skip history for heartbeat) + var history []providers.Message + var summary string + if !opts.NoHistory { history = agent.Sessions.GetHistory(opts.SessionKey) + summary = agent.Sessions.GetSummary(opts.SessionKey) // Sanitize history to remove orphaned tool calls (from crashes/session collisions) + var removedCount int + history, removedCount = session.SanitizeHistory(history) + if removedCount > 0 { logger.WarnCF("agent", "Sanitized session history: removed orphaned messages", + map[string]any{ - "session_key": opts.SessionKey, + "session_key": opts.SessionKey, + "removed_count": removedCount, }) + // Persist the sanitized history + agent.Sessions.SetHistory(opts.SessionKey, history) + _ = agent.Sessions.Save(opts.SessionKey) } } + messages := agent.ContextBuilder.BuildMessages( + history, + summary, + opts.UserMessage, + nil, + opts.Channel, + opts.ChatID, ) // 2b. Interview staleness nudge: if MEMORY.md hasn't been updated for + // several consecutive turns, inject a reminder so the AI writes its findings. + const interviewStaleThreshold = 2 + if agent.ContextBuilder.GetPlanStatus() == "interviewing" && agent.interviewStaleCount >= interviewStaleThreshold { messages = append(messages, providers.Message{ - Role: "user", + Role: "user", + Content: "[System] You have been interviewing for several turns without updating memory/MEMORY.md. Please use edit_file now to save your findings to the ## Context section, or organize the plan into ## Phase sections with `- [ ]` checkbox steps if you have enough information.", }) } // 2c. Background plan preamble: append to system prompt (high attention) + // so the LLM knows from the start that it must mark steps [x]. + // Skip if a chat session is actively working on the plan directory. + if opts.Background && agent.ContextBuilder.HasActivePlan() && agent.ContextBuilder.GetPlanStatus() == "executing" { planDir := agent.ContextBuilder.GetPlanWorkDir() + skipPreamble := planDir != "" && al.sessions.IsActiveInDir(planDir, "heartbeat") + if !skipPreamble && len(messages) > 0 && messages[0].Role == "system" { var sb strings.Builder + sb.WriteString(messages[0].Content) + sb.WriteString("\n\n## Background Execution\n") + sb.WriteString("You are running as a background heartbeat with no conversation history. ") + sb.WriteString("MEMORY.md is the only shared state between heartbeats. ") + sb.WriteString( + "After completing each plan step, immediately use edit_file to mark it [x] in memory/MEMORY.md.", ) + messages[0].Content = sb.String() } } // 2d. Snapshot plan status and MEMORY.md size before LLM iteration. + preStatus := agent.ContextBuilder.GetPlanStatus() + var preMemoryLen int + if preStatus == "interviewing" { preMemoryLen = len(agent.ContextBuilder.ReadMemory()) } // 3. Save user message to session (use compact form if available) + historyMsg := opts.UserMessage + if opts.HistoryMessage != "" { historyMsg = opts.HistoryMessage } + agent.Sessions.AddMessage(opts.SessionKey, "user", historyMsg) // 4. Record user prompt for stats + if al.stats != nil { al.stats.RecordPrompt() } // Capture the finalized system prompt for Mini App inspection + if len(messages) > 0 { al.lastSystemPrompt.Store(messages[0].Content) + al.promptDirty.Store(false) } // 5. Run LLM iteration loop (with automatic phase transitions) + var iteration int + const maxPhaseTransitions = 10 for phaseLoop := 0; ; phaseLoop++ { // On phase transition: rebuild system prompt with new phase context + nudge + if phaseLoop > 0 { messages = agent.ContextBuilder.BuildMessages( + agent.Sessions.GetHistory(opts.SessionKey), + agent.Sessions.GetSummary(opts.SessionKey), + "", nil, opts.Channel, opts.ChatID, ) + messages = append(messages, providers.Message{ Role: "user", + Content: fmt.Sprintf( + "[System] Phase %d is now active. Continue working on the next steps.", + agent.ContextBuilder.GetCurrentPhase(), ), }) + if len(messages) > 0 { al.lastSystemPrompt.Store(messages[0].Content) } } curPlanStatus := preStatus + if phaseLoop > 0 { curPlanStatus = agent.ContextBuilder.GetPlanStatus() } var err error + finalContent, iteration, err = al.runLLMIteration(ctx, agent, messages, opts, task, curPlanStatus) if err != nil { return "", err } // 5a. Auto-advance plan phases after LLM iteration + postStatus := agent.ContextBuilder.GetPlanStatus() + if !agent.ContextBuilder.HasActivePlan() || + !(postStatus == "executing" || postStatus == "review" || postStatus == "completed") { break } // Intercept: if AI changed status to executing or review without user approval + // (from interviewing or review), validate and hold at "review". + if preStatus == "interviewing" || (preStatus == "review" && postStatus == "executing") { if err := agent.ContextBuilder.ValidatePlanStructure(); err != nil { _ = agent.ContextBuilder.SetPlanStatus("interviewing") + logger.WarnCF("agent", "Reverted plan to interviewing: "+err.Error(), + map[string]any{"agent_id": agent.ID}) + rejectionMsg := "[System] Plan rejected: " + err.Error() + ". Fix and try again." + agent.Sessions.AddMessage(opts.SessionKey, "user", rejectionMsg) } else { _ = agent.ContextBuilder.SetPlanStatus("review") + al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStatePlanReview, "") + if !constants.IsInternalChannel(opts.Channel) { planDisplay := agent.ContextBuilder.FormatPlanDisplay() + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: planDisplay + "\n\nUse /plan start to approve, or continue chatting to refine.", + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: planDisplay + "\n\nUse /plan start to approve, or continue chatting to refine.", + SkipPlaceholder: true, }) } } + break } if postStatus == "executing" && agent.ContextBuilder.GetTotalPhases() == 0 { _ = agent.ContextBuilder.SetPlanStatus("interviewing") + logger.WarnCF("agent", "Reverted plan to interviewing: no phases defined", + map[string]any{"agent_id": agent.ID}) + break } if agent.ContextBuilder.IsPlanComplete() { total := agent.ContextBuilder.GetTotalPhases() + _ = agent.ContextBuilder.SetCurrentPhase(total) + if preStatus != "completed" { _ = agent.ContextBuilder.SetPlanStatus("completed") + al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStatePlanCompleted, "") // Deactivate worktree on plan completion + commitMsg := "plan: " + agent.ContextBuilder.Memory().GetPlanTaskName() + wtResult, _ := agent.DeactivateWorktree(opts.SessionKey, commitMsg, false) if !constants.IsInternalChannel(opts.Channel) { msg := "\u2705 Plan completed!" + if wtResult != nil && wtResult.CommitsAhead > 0 { msg += fmt.Sprintf("\nBranch `%s` retained (%d commits). To merge: `git merge %s`", + wtResult.Branch, wtResult.CommitsAhead, wtResult.Branch) } + _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: msg, + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: msg, + SkipPlaceholder: true, }) } } + break } if agent.ContextBuilder.IsCurrentPhaseComplete() { if phaseLoop >= maxPhaseTransitions { logger.WarnCF("agent", "Max phase transitions reached, stopping", + map[string]any{"agent_id": agent.ID, "transitions": phaseLoop}) + break } + prev := agent.ContextBuilder.GetCurrentPhase() + _ = agent.ContextBuilder.AdvancePhase() + next := agent.ContextBuilder.GetCurrentPhase() + if !constants.IsInternalChannel(opts.Channel) { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: fmt.Sprintf("Phase %d complete. Moving to Phase %d.", prev, next), + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: fmt.Sprintf("Phase %d complete. Moving to Phase %d.", prev, next), + SkipPlaceholder: true, }) } + al.notifyStateChange() + continue } @@ -1531,55 +2205,75 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt al.notifyStateChange() // 5b. Interview staleness detection: compare MEMORY.md size after iteration. + if agent.ContextBuilder.GetPlanStatus() == "interviewing" { postMemoryLen := len(agent.ContextBuilder.ReadMemory()) + if postMemoryLen == preMemoryLen { agent.interviewStaleCount++ } else { agent.interviewStaleCount = 0 } + agent.interviewMemoryLen = postMemoryLen } else { // Reset counter when not interviewing. + agent.interviewStaleCount = 0 } // 5c. Handle empty response + if finalContent == "" { finalContent = opts.DefaultResponse } // 5d. Store result summary for task completion notification + if task != nil { task.Result = utils.Truncate(finalContent, 280) } // 6. Save final assistant message to session (deferred write-behind) + agent.Sessions.AddMessage(opts.SessionKey, "assistant", finalContent) + agent.Sessions.MarkDirty(opts.SessionKey) // 7. Optional: summarization + if opts.EnableSummary { al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) } // 8. Optional: send response via bus + if opts.SendResponse { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: finalContent, + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: finalContent, + SkipPlaceholder: opts.SystemMessage, // suppress Telegram "Thinking..." for system messages + }) } // 9. Log response + responsePreview := utils.Truncate(finalContent, 120) + logger.InfoCF("agent", fmt.Sprintf("Response: %s", responsePreview), + map[string]any{ - "agent_id": agent.ID, - "session_key": opts.SessionKey, - "iterations": iteration, + "agent_id": agent.ID, + + "session_key": opts.SessionKey, + + "iterations": iteration, + "final_length": len(finalContent), }) @@ -1587,15 +2281,18 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt } // Task reminder constants and helpers. + const ( taskReminderMaxChars = 500 - blockerMaxChars = 200 + + blockerMaxChars = 200 ) func shouldInjectReminder(iteration, interval int) bool { if interval <= 0 { return false } + return iteration > 1 && iteration%interval == 0 } @@ -1603,8 +2300,10 @@ func buildTaskReminder(userMessage string, lastBlocker string) providers.Message truncatedTask := utils.Truncate(userMessage, taskReminderMaxChars) var content string + if lastBlocker != "" { truncatedBlocker := utils.Truncate(lastBlocker, blockerMaxChars) + content = fmt.Sprintf( "[TASK REMINDER]\nOriginal task:\n---\n%s\n---\nLast blocker:\n---\n%s\n---\nFix the blocker if essential, or find an alternative. If all steps are complete, move on.", truncatedTask, @@ -1618,141 +2317,213 @@ func buildTaskReminder(userMessage string, lastBlocker string) providers.Message } return providers.Message{ - Role: "user", + Role: "user", + Content: content, } } // interviewRejectMessage is the fixed rejection text injected when tool calls + // are blocked during the interview phase. It is deliberately short to avoid + // wasting tokens, and ends with a purpose reminder to steer the LLM back. + const interviewRejectMessage = "[System] Tool call rejected. " + + "You are in interview mode — ask the user questions and update MEMORY.md. " + + "Do not execute, edit, or write project files." // buildPlanReminder returns a reminder message for plan pre-execution states + // (interviewing / review) to keep the AI focused on the interview workflow + // during tool-call iterations. + func buildPlanReminder(planStatus string) (providers.Message, bool) { var content string + switch planStatus { case "interviewing": + content = "[System] You are interviewing the user to build a plan. " + + "Ask clarifying questions and save findings to ## Context in memory/MEMORY.md using edit_file. " + + "When you have enough information, write ## Phase sections with `- [ ]` checkbox steps, and ## Commands section. " + + "Then change > Status: to review. Do NOT set it to executing." + case "review": + content = "[System] The plan is under review. " + + "Wait for the user to approve or request changes. Do not proceed with execution." + default: + return providers.Message{}, false } + return providers.Message{Role: "user", Content: content}, true } // buildOrchReminder returns a reminder to use spawn/subagent during plan execution. + // Fires on first iteration and every 3rd iteration to reinforce delegation behavior. + func buildOrchReminder(iteration int) (providers.Message, bool) { if iteration != 1 && iteration%3 != 0 { return providers.Message{}, false } + content := `[System] ORCHESTRATION mode active. You MUST delegate plan steps to subagents. + Use spawn (non-blocking, returns immediately) or subagent (blocking, waits for result). + Do NOT implement steps inline unless they are a single trivial tool call. + + To delegate, call the tool with JSON arguments: + Tool: spawn Arguments: {"task": "...", "preset": "scout", "label": "..."} + Tool: subagent Arguments: {"task": "...", "label": "..."} + + Spawn multiple independent steps in parallel for maximum throughput.` + return providers.Message{Role: "user", Content: content}, true } // cdPrefixPattern matches "cd /some/path && " at the start of a shell command. + // Group 1 captures the target directory path. + var cdPrefixPattern = regexp.MustCompile(`^cd\s+(\S+)\s*&&\s*`) // optFlagPattern matches option flags like --verbose, -v, --timeout=60, -q. + // Only standalone flags are removed; flags whose value is the next positional + // argument (e.g. "-A 20") are kept because removing them would lose context. + var optFlagPattern = regexp.MustCompile(`\s+--?\w[\w-]*(=\S*)?`) // extractExecProjectDir extracts the basename of an exec cd target. + // Returns "" if the command has no cd prefix. + func extractExecProjectDir(args map[string]any) string { cmd, _ := args["command"].(string) + if cmd == "" { return "" } + m := cdPrefixPattern.FindStringSubmatch(cmd) + if len(m) < 2 { return "" } + cdPath := strings.TrimRight(m[1], "/\\") + if idx := strings.LastIndex(cdPath, "/"); idx >= 0 { return cdPath[idx+1:] } + if idx := strings.LastIndex(cdPath, "\\"); idx >= 0 { return cdPath[idx+1:] } + return cdPath } // fileParentRelDir returns the parent directory of a file path, relative to + // workspace. Returns "" if the path is not under workspace or has no parent. + func fileParentRelDir(filePath, workspace string) string { ws := strings.TrimRight(workspace, "/\\") + if ws == "" { return "" } + rest := strings.TrimPrefix(filePath, ws) + if rest == filePath { return "" // not under workspace } + rest = strings.TrimLeft(rest, "/\\") + // Remove the filename — keep only the directory part + if idx := strings.LastIndexAny(rest, "/\\"); idx >= 0 { return rest[:idx] } + return "" // file is directly under workspace, no meaningful dir } // commonDirPrefix computes the longest common directory prefix of two + // slash-separated paths. Returns "" if there is no common component. + func commonDirPrefix(a, b string) string { partsA := strings.Split(a, "/") + partsB := strings.Split(b, "/") + n := len(partsA) + if len(partsB) < n { n = len(partsB) } + common := 0 + for i := 0; i < n; i++ { if partsA[i] != partsB[i] { break } + common = i + 1 } + if common == 0 { return "" } + return strings.Join(partsA[:common], "/") } // displayProjectDir returns the project directory name for status display. + // Prefers the authoritative exec-based projectDir; falls back to the + // basename of the file-based common directory. + func displayProjectDir(task *activeTask) string { if task.projectDir != "" { return task.projectDir } + if task.fileCommonDir != "" { dir := task.fileCommonDir + if idx := strings.LastIndex(dir, "/"); idx >= 0 { return dir[idx+1:] } + return dir } + return "" } @@ -1760,48 +2531,71 @@ func (al *AgentLoop) targetReasoningChannelID(channelName string) (chatID string if al.channelManager == nil { return "" } + if ch, ok := al.channelManager.GetChannel(channelName); ok { return ch.ReasoningChannelID() } + return "" } // buildArgsSnippet produces a human-friendly snippet for the tool log. + // For exec: extracts the command and strips the leading "cd <workspace> && ". + // For file tools: extracts the path and strips the workspace prefix. + // Falls back to raw JSON truncation. + func buildArgsSnippet(toolName string, args map[string]any, workspace string) string { switch toolName { case "exec": + cmd, _ := args["command"].(string) + if cmd == "" { break } + cmd = cdPrefixPattern.ReplaceAllString(cmd, "") + cmd = optFlagPattern.ReplaceAllString(cmd, "") + return utils.Truncate(cmd, 80) case "read_file", "write_file", "edit_file", "append_file", "list_dir": + path, _ := args["path"].(string) + if path == "" { break } + if workspace != "" { path = strings.TrimPrefix(path, workspace) + path = strings.TrimPrefix(path, "/") } + // Prioritize filename: if path is too long, show "…/filename" + const maxPath = 60 + if runes := []rune(path); len(runes) > maxPath { // Find last slash to extract filename + if lastSlash := strings.LastIndex(path, "/"); lastSlash >= 0 { - filename := path[lastSlash:] // includes "/" + filename := path[lastSlash:] // includes "/" + dirBudget := maxPath - len([]rune(filename)) - 1 // 1 for "…" + if dirBudget > 0 { dir := []rune(path[:lastSlash]) + if len(dir) > dirBudget { dir = dir[:dirBudget] } + path = string(dir) + "\u2026" + filename } else { path = "\u2026" + filename @@ -1810,40 +2604,54 @@ func buildArgsSnippet(toolName string, args map[string]any, workspace string) st path = utils.Truncate(path, maxPath) } } + return path } // Default: raw JSON truncated + argsJSON, _ := json.Marshal(args) + return utils.Truncate(string(argsJSON), 80) } // maxEntryLineWidth is the max rune count for a single-line log entry. + // Telegram chat bubbles on mobile are roughly 40-45 chars wide. + const maxEntryLineWidth = 42 // isFileToolEntry returns true if the entry name contains a file-operation tool. + func isFileToolEntry(name string) bool { for _, t := range []string{"read_file", "write_file", "edit_file", "append_file", "list_dir"} { if strings.Contains(name, t) { return true } } + return false } // formatCompactEntry formats a finished tool log entry as a fixed single line. + // The result marker (✓/✗) is always shown at the end regardless of truncation. + // File tools omit duration (always near-instant); paths truncate from the + // start so the filename is always visible. + func formatCompactEntry(entry toolLogEntry) string { result := entry.Result + if result == "" { result = "\u23F3" // ⏳ } // File tools: strip duration, keep only marker (✓/✗/⏳) + isFile := isFileToolEntry(entry.Name) + if isFile { if r := []rune(result); len(r) > 0 { result = string(r[0:1]) // just the symbol @@ -1851,49 +2659,74 @@ func formatCompactEntry(entry toolLogEntry) string { } // Budget for ArgsSnip: total - name - " " - " " - result + nameLen := utf8.RuneCountInString(entry.Name) + resultLen := utf8.RuneCountInString(result) + argsBudget := maxEntryLineWidth - nameLen - 1 - 1 - resultLen args := entry.ArgsSnip + if args != "" && argsBudget > 3 { argsRunes := []rune(args) + if len(argsRunes) > argsBudget { // Paths: truncate from the start, keeping the filename visible + if strings.Contains(args, "/") { args = "\u2026" + string(argsRunes[len(argsRunes)-argsBudget+1:]) } else { args = string(argsRunes[:argsBudget-1]) + "\u2026" } } + var sb strings.Builder + sb.Grow(len(entry.Name) + 1 + len(args) + 1 + len(result)) + sb.WriteString(entry.Name) + sb.WriteByte(' ') + sb.WriteString(args) + sb.WriteByte(' ') + sb.WriteString(result) + return sb.String() } // No room for args or args empty + var sb strings.Builder + sb.Grow(len(entry.Name) + 1 + len(result)) + sb.WriteString(entry.Name) + sb.WriteByte(' ') + sb.WriteString(result) + return sb.String() } // formatLatestEntry formats the latest entry command without its result marker. + // Since the result goes on the next line, the full width is available for the command. + func formatLatestEntry(entry toolLogEntry) string { nameLen := utf8.RuneCountInString(entry.Name) + argsBudget := maxEntryLineWidth - nameLen - 1 // name + space + args (no result) args := entry.ArgsSnip + if args != "" && argsBudget > 3 { argsRunes := []rune(args) + if len(argsRunes) > argsBudget { if strings.Contains(args, "/") { args = "\u2026" + string(argsRunes[len(argsRunes)-argsBudget+1:]) @@ -1901,175 +2734,260 @@ func formatLatestEntry(entry toolLogEntry) string { args = string(argsRunes[:argsBudget-1]) + "\u2026" } } + var sb strings.Builder + sb.Grow(len(entry.Name) + 1 + len(args)) + sb.WriteString(entry.Name) + sb.WriteByte(' ') + sb.WriteString(args) + return sb.String() } + return entry.Name } // compressRepeats reduces runs of 3+ identical non-alphanumeric, non-space + // characters to just 2. e.g. "======" → "==", "---" → "--". + func compressRepeats(s string) string { runes := []rune(s) + if len(runes) < 3 { return s } + var sb strings.Builder + sb.Grow(len(s)) + i := 0 + for i < len(runes) { r := runes[i] + if !unicode.IsLetter(r) && !unicode.IsDigit(r) && !unicode.IsSpace(r) { j := i + 1 + for j < len(runes) && runes[j] == r { j++ } + if j-i >= 3 { sb.WriteRune(r) + sb.WriteRune(r) + i = j + continue } } + sb.WriteRune(r) + i++ } + return sb.String() } // Display layout constants. + const ( - displayPastEntries = 4 // number of compact 1-line past entries - displayErrorLines = 5 // content lines inside the error code block - statusSeparator = "\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\n" + displayPastEntries = 4 // number of compact 1-line past entries + + displayErrorLines = 5 // content lines inside the error code block + + statusSeparator = "\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\n" + streamingDisplayLines = 17 // line count matching buildRichStatus output + ) // buildRichStatus builds a fixed-height terminal-like status display. + // + // Layout (always the same number of lines): + // + // 🔄 Task in progress (N/M) header + // 📁 workspace-path header + // ━━━━━━━━━━ separator + // [N] compact-past-1 ✓ Xs past (1 line each) + // [N] compact-past-2 ✗ Xs past + // [N] compact-past-3 ✓ Xs past + // [N] compact-past-4 ✓ Xs past + // [N] latest-command latest (no result, wider args) + // ⏳ latest result + // reserved + // ``` error fence + // err-line / placeholder error body (5 lines) + // ``` error fence + // ↩️ Reply to intervene footer (background only) + func buildRichStatus(task *activeTask, isBackground bool, workspace string) string { task.mu.Lock() + defer task.mu.Unlock() var sb strings.Builder // --- Header --- + sb.WriteString("\U0001F504 Task in progress (") + sb.WriteString(strconv.Itoa(task.Iteration)) + sb.WriteByte('/') + sb.WriteString(strconv.Itoa(task.MaxIter)) + sb.WriteString(")\n") + // Project directory: exec cd (authoritative) → file LCP → workspace basename + sb.WriteString("\U0001F4C1 ") + if dir := displayProjectDir(task); dir != "" { sb.WriteString(dir) } else if workspace != "" { project := strings.TrimRight(workspace, "/\\") + if idx := strings.LastIndex(project, "/"); idx >= 0 { project = project[idx+1:] } else if idx := strings.LastIndex(project, "\\"); idx >= 0 { project = project[idx+1:] } + sb.WriteString(project) } + sb.WriteByte('\n') + sb.WriteString(statusSeparator) // --- Task entries (displayPastEntries + 2 lines for latest) --- + entries := task.toolLog + if len(entries) > maxToolLogEntries { entries = entries[len(entries)-maxToolLogEntries:] } var pastEntries []toolLogEntry + var latest *toolLogEntry + if len(entries) > 0 { latest = &entries[len(entries)-1] + if len(entries) > 1 { start := len(entries) - 1 - displayPastEntries + if start < 0 { start = 0 } + pastEntries = entries[start : len(entries)-1] } } // Past entries: exactly displayPastEntries lines (pad if fewer) + for i := 0; i < displayPastEntries; i++ { if i < len(pastEntries) { sb.WriteString(formatCompactEntry(pastEntries[i])) } else { sb.WriteString("\u2800") } + sb.WriteByte('\n') } // Latest entry: command on one line, result on next + if latest != nil { sb.WriteString(formatLatestEntry(*latest)) + sb.WriteByte('\n') + sb.WriteString(" ") + if latest.Result != "" { sb.WriteString(latest.Result) } else { sb.WriteString("\u23F3") } + sb.WriteByte('\n') } else { sb.WriteString("\u23F3 waiting...\n") + sb.WriteString("\u2800\n") } // Reserved (1 line) + sb.WriteString("\u2800\n") // --- Error region (code fence, no separator) --- + sb.WriteString("```\n") errEntry := task.lastError + if errEntry != nil { sb.WriteString("\u274C ") + sb.WriteString(formatCompactEntry(*errEntry)) + sb.WriteByte('\n') var detailLines []string + if errEntry.ErrDetail != "" { detailLines = strings.Split(errEntry.ErrDetail, "\n") } + for i := 0; i < displayErrorLines-1; i++ { if i < len(detailLines) { line := compressRepeats(detailLines[i]) + if runes := []rune(line); len(runes) > maxEntryLineWidth { line = string(runes[:maxEntryLineWidth-1]) + "\u2026" } + sb.WriteString(line) } else { sb.WriteString("\u2800") } + sb.WriteByte('\n') } } else { sb.WriteString("\u2714 No errors\n") + for i := 0; i < displayErrorLines-1; i++ { sb.WriteString("\u2800\n") } @@ -2080,6 +2998,7 @@ func buildRichStatus(task *activeTask, isBackground bool, workspace string) stri if isBackground { sb.WriteString("\u21A9\uFE0F Reply to intervene") } + return sb.String() } @@ -2089,347 +3008,517 @@ func (al *AgentLoop) handleReasoning(ctx context.Context, reasoningContent, chan } // Check context cancellation before attempting to publish, + // since PublishOutbound's select may race between send and ctx.Done(). + if ctx.Err() != nil { return } // Use a short timeout so the goroutine does not block indefinitely when + // the outbound bus is full. Reasoning output is best-effort; dropping it + // is acceptable to avoid goroutine accumulation. + pubCtx, pubCancel := context.WithTimeout(ctx, 5*time.Second) + defer pubCancel() if err := al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{ Channel: channelName, - ChatID: channelID, + + ChatID: channelID, + Content: reasoningContent, }); err != nil { // Treat context.DeadlineExceeded / context.Canceled as expected + // (bus full under load, or parent canceled). Check the error + // itself rather than ctx.Err(), because pubCtx may time out + // (5 s) while the parent ctx is still active. + // Also treat ErrBusClosed as expected — it occurs during normal + // shutdown when the bus is closed before all goroutines finish. + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) || + errors.Is(err, bus.ErrBusClosed) { logger.DebugCF("agent", "Reasoning publish skipped (timeout/cancel)", map[string]any{ "channel": channelName, - "error": err.Error(), + + "error": err.Error(), }) } else { logger.WarnCF("agent", "Failed to publish reasoning (best-effort)", map[string]any{ "channel": channelName, - "error": err.Error(), + + "error": err.Error(), }) } } } // streamingReasoningLines is the number of lines reserved for reasoning + // in the streaming display. The remaining lines go to content. + const streamingReasoningLines = 6 // buildStreamingDisplay builds a fixed-height status bubble for streaming. + // + // Layout when reasoning is active (reasoning only or both): + // + // 🧠 Thinking... + // ━━━━━━━━━━ + // <reasoning tail — streamingReasoningLines lines> + // ━━━━━━━━━━ + // <content tail — remaining lines> (or blank if content is empty) + // █ + // + // Layout when no reasoning (content only): + // + // <content tail — streamingDisplayLines lines> + // █ + func buildStreamingDisplay(content, reasoning string) string { if reasoning == "" { // No reasoning — full window for content. + return utils.TailPad(content, streamingDisplayLines, maxEntryLineWidth) + " \u2589" } var sb strings.Builder // Header + if content == "" { sb.WriteString("\U0001f9e0 Thinking...\n") } else { sb.WriteString("\U0001f9e0 Thought, now responding...\n") } + sb.WriteString(statusSeparator) // Reasoning window + headerLines := 2 // header + separator + footerLines := 1 // separator before content + contentLines := streamingDisplayLines - headerLines - footerLines - streamingReasoningLines + if contentLines < 3 { contentLines = 3 } + rLines := streamingDisplayLines - headerLines - footerLines - contentLines sb.WriteString(utils.TailPad(reasoning, rLines, maxEntryLineWidth)) + sb.WriteByte('\n') + sb.WriteString(statusSeparator) // Content window (may be blank padding if content hasn't started) + sb.WriteString(utils.TailPad(content, contentLines, maxEntryLineWidth)) + sb.WriteString(" \u2589") return sb.String() } // runLLMIteration executes the LLM call loop with tool handling. + // consumeStreamWithRepetitionDetection reads StreamEvents from ch, accumulates + // content and tool calls, and runs repetition detection every checkInterval runes. + // If repetition is detected, cancelFn is called to abort the HTTP request and + // the function returns the partial response with detected=true. + func consumeStreamWithRepetitionDetection( ch <-chan protocoltypes.StreamEvent, + cancelFn context.CancelFunc, + checkInterval int, + onChunk func(content, reasoning string), ) (*providers.LLMResponse, bool, error) { var content strings.Builder + var reasoning strings.Builder + var toolCalls []streamToolCallAcc + var finishReason string + var usage *providers.UsageInfo + runesSinceLastCheck := 0 for ev := range ch { if ev.Err != nil { return nil, false, ev.Err } + updated := false + if ev.ContentDelta != "" { content.WriteString(ev.ContentDelta) + runesSinceLastCheck += utf8.RuneCountInString(ev.ContentDelta) + updated = true } + if ev.ReasoningDelta != "" { reasoning.WriteString(ev.ReasoningDelta) + updated = true } + if updated && onChunk != nil { onChunk(content.String(), reasoning.String()) } + if ev.FinishReason != "" { finishReason = ev.FinishReason } + if ev.Usage != nil { usage = ev.Usage } + for _, tc := range ev.ToolCallDeltas { for len(toolCalls) <= tc.Index { toolCalls = append(toolCalls, streamToolCallAcc{}) } + if tc.ID != "" { toolCalls[tc.Index].id = tc.ID } + if tc.Name != "" { toolCalls[tc.Index].name = tc.Name } + toolCalls[tc.Index].args.WriteString(tc.ArgumentsDelta) } // Run repetition detection periodically on accumulated content. + if runesSinceLastCheck >= checkInterval && content.Len() > 2000 { runesSinceLastCheck = 0 + if utils.DetectRepetitionLoop(content.String()) { cancelFn() + // Drain remaining events so the producer goroutine can exit. + for range ch { } + resp := buildAccumulatedResponse(content.String(), reasoning.String(), toolCalls, finishReason, usage) + return resp, true, nil } } } resp := buildAccumulatedResponse(content.String(), reasoning.String(), toolCalls, finishReason, usage) + return resp, false, nil } // streamToolCallAcc accumulates streamed tool call fragments. + type streamToolCallAcc struct { - id string + id string + name string + args strings.Builder } // buildAccumulatedResponse constructs an LLMResponse from accumulated stream data. + func buildAccumulatedResponse( content, reasoning string, + toolCalls []streamToolCallAcc, + finishReason string, + usage *providers.UsageInfo, ) *providers.LLMResponse { resp := &providers.LLMResponse{ - Content: content, - Reasoning: reasoning, + Content: content, + + Reasoning: reasoning, + FinishReason: finishReason, - Usage: usage, + + Usage: usage, } + for _, tc := range toolCalls { arguments := make(map[string]any) + argStr := tc.args.String() + if argStr != "" { if err := json.Unmarshal([]byte(argStr), &arguments); err != nil { arguments["raw"] = argStr } } + resp.ToolCalls = append(resp.ToolCalls, providers.ToolCall{ - ID: tc.id, - Name: tc.name, + ID: tc.id, + + Name: tc.name, + Arguments: arguments, }) } + return resp } func (al *AgentLoop) runLLMIteration( ctx context.Context, + agent *AgentInstance, + messages []providers.Message, + opts processOptions, + task *activeTask, + planSnapshot string, ) (string, int, error) { iteration := 0 + var finalContent string + lastReminderIdx := -1 + planMarkNudged := false // true after we've already nudged once for [x] marking maxIter := agent.MaxIterations // Snapshot unchecked step count before tool loop so we can detect progress. + preUnchecked := -1 // -1 = not tracking + if planSnapshot == "executing" { preUnchecked = strings.Count(agent.ContextBuilder.ReadMemory(), "- [ ]") } // Determine if this is a background task (cron, heartbeat, etc.) + isBackground := opts.TaskID != "" for iteration < maxIter { iteration++ // Update active task iteration + if task != nil { task.mu.Lock() + task.Iteration = iteration + task.mu.Unlock() } // Check for user intervention via interrupt channel + if task != nil { select { case msg := <-task.interrupt: + messages = append(messages, providers.Message{ - Role: "user", + Role: "user", + Content: "[User Intervention] " + msg, }) + logger.InfoCF("agent", "User intervention injected", + map[string]any{"agent_id": agent.ID, "iteration": iteration}) + default: } } logger.DebugCF("agent", "LLM iteration", + map[string]any{ - "agent_id": agent.ID, + "agent_id": agent.ID, + "iteration": iteration, - "max": maxIter, + + "max": maxIter, }) // Build tool definitions + providerToolDefs := agent.Tools.ToProviderDefs() // Interview mode: strip tool definitions the LLM must not use, + // reducing token cost and preventing wasted reject-retry cycles. + if isPlanPreExecution(planSnapshot) { providerToolDefs = filterInterviewTools(providerToolDefs) } // Log LLM request details + logger.DebugCF("agent", "LLM request", + map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "model": agent.Model, - "messages_count": len(messages), - "tools_count": len(providerToolDefs), - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, + "agent_id": agent.ID, + + "iteration": iteration, + + "model": agent.Model, + + "messages_count": len(messages), + + "tools_count": len(providerToolDefs), + + "max_tokens": agent.MaxTokens, + + "temperature": agent.Temperature, + "system_prompt_len": len(messages[0].Content), }) // Log full messages (detailed) + logger.DebugCF("agent", "Full LLM request", + map[string]any{ - "iteration": iteration, + "iteration": iteration, + "messages_json": formatMessagesForLog(messages), - "tools_json": formatToolsForLog(providerToolDefs), + + "tools_json": formatToolsForLog(providerToolDefs), }) // Call LLM with fallback chain if candidates are configured. + var response *providers.LLMResponse + var err error // Build onChunk callback for streaming preview. + // Instead of a fixed-interval throttle, use a Go channel with + // latest-value semantics: a consumer goroutine publishes status + // updates as fast as the bus → manager → channel pipeline allows. + // Backpressure is provided naturally by the per-channel rate limiter + // (e.g. 20 msg/s for Telegram's SendDraft, 1 msg/s for Discord's EditMessage). + type streamUpdate struct{ accumulated, reasoning string } + var onChunk func(string, string) + var streamCh chan streamUpdate + var streamDone chan struct{} + if !constants.IsInternalChannel(opts.Channel) { streamCh = make(chan streamUpdate, 1) + streamDone = make(chan struct{}) + go func() { defer close(streamDone) + for up := range streamCh { display := buildStreamingDisplay(up.accumulated, up.reasoning) + outMsg := bus.OutboundMessage{ Channel: opts.Channel, - ChatID: opts.ChatID, + + ChatID: opts.ChatID, + Content: display, } + // For background tasks, publish streaming preview as + // IsTaskStatus so it shares the same bubble as task + // progress/completion (avoids a second bubble). + if opts.Background && opts.TaskID != "" { outMsg.IsTaskStatus = true + outMsg.TaskID = opts.TaskID } else { outMsg.IsStatus = true } + _ = al.bus.PublishOutbound(ctx, outMsg) } }() + onChunk = func(accumulated, reasoning string) { if task != nil { task.streamedChunks = true } + up := streamUpdate{accumulated, reasoning} + // Non-blocking latest-value send: if the consumer hasn't + // drained the previous update, replace it with the latest. + select { case streamCh <- up: + default: + // Channel full — drain stale value, then send latest. + select { case <-streamCh: + default: } + select { case streamCh <- up: + default: } } @@ -2437,75 +3526,109 @@ func (al *AgentLoop) runLLMIteration( } // doCall invokes a single LLM provider, using streaming with + // early repetition detection when the provider supports it. + opts_ := map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, + "max_tokens": agent.MaxTokens, + + "temperature": agent.Temperature, + "prompt_cache_key": agent.ID, } + doCall := func(ctx context.Context, p providers.LLMProvider, model string) (*providers.LLMResponse, error) { if sp, ok := p.(providers.StreamingProvider); ok && sp.CanStream() { streamCtx, streamCancel := context.WithCancel(ctx) + defer streamCancel() + ch, sErr := sp.ChatStream(streamCtx, messages, providerToolDefs, model, opts_) + if sErr != nil { return nil, sErr } + resp, repetition, sErr := consumeStreamWithRepetitionDetection(ch, streamCancel, 1000, onChunk) + if sErr != nil { return nil, sErr } + if repetition { resp.FinishReason = "repetition_detected" } + return resp, nil } + return p.Chat(ctx, messages, providerToolDefs, model, opts_) } callLLM := func() (*providers.LLMResponse, error) { // Plan model switching: use plan model during interviewing/review phases + candidates := agent.Candidates + primaryModel := agent.Model + if isPlanPreExecution(planSnapshot) && agent.PlanModel != "" { candidates = agent.PlanCandidates + primaryModel = agent.PlanModel + logger.InfoCF("agent", "Using plan model", + map[string]any{"agent_id": agent.ID, "plan_model": agent.PlanModel}) } if len(candidates) > 1 && al.fallback != nil { fbResult, fbErr := al.fallback.Execute(ctx, candidates, + func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { p := al.resolveProvider(provider, model, agent.Provider) + return doCall(ctx, p, model) }, ) + if fbErr != nil { return nil, fbErr } + if fbResult.Provider != "" && len(fbResult.Attempts) > 0 { logger.InfoCF("agent", fmt.Sprintf("Fallback: succeeded with %s/%s after %d attempts", + fbResult.Provider, fbResult.Model, len(fbResult.Attempts)+1), + map[string]any{"agent_id": agent.ID, "iteration": iteration}) } + return fbResult.Response, nil } + if len(candidates) > 0 { c := candidates[0] + p := al.resolveProvider(c.Provider, c.Model, agent.Provider) + return doCall(ctx, p, c.Model) } + return doCall(ctx, agent.Provider, primaryModel) } // Report waiting state to canvas before each LLM call. + al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStateWaiting, "") // Retry loop for context/token errors + maxRetries := 2 + for retry := 0; retry <= maxRetries; retry++ { response, err = callLLM() + if err == nil { break } @@ -2513,127 +3636,195 @@ func (al *AgentLoop) runLLMIteration( errMsg := strings.ToLower(err.Error()) // Check if this is a network/HTTP timeout — not a context window error. + isTimeoutError := errors.Is(err, context.DeadlineExceeded) || + strings.Contains(errMsg, "deadline exceeded") || + strings.Contains(errMsg, "client.timeout") || + strings.Contains(errMsg, "timed out") || + strings.Contains(errMsg, "timeout exceeded") // Detect real context window / token limit errors, excluding network timeouts. + isContextError := !isTimeoutError && (strings.Contains(errMsg, "context_length_exceeded") || + strings.Contains(errMsg, "context window") || + strings.Contains(errMsg, "maximum context length") || + strings.Contains(errMsg, "token limit") || + strings.Contains(errMsg, "too many tokens") || + strings.Contains(errMsg, "max_tokens") || + strings.Contains(errMsg, "invalidparameter") || + strings.Contains(errMsg, "prompt is too long") || + strings.Contains(errMsg, "request too large")) if isTimeoutError && retry < maxRetries { backoff := time.Duration(retry+1) * 5 * time.Second + logger.WarnCF("agent", "Timeout error, retrying after backoff", map[string]any{ - "error": err.Error(), - "retry": retry, + "error": err.Error(), + + "retry": retry, + "backoff": backoff.String(), }) + time.Sleep(backoff) + continue } if isContextError && retry < maxRetries { logger.WarnCF("agent", "Context window error detected, attempting compression", map[string]any{ "error": err.Error(), + "retry": retry, }) if retry == 0 && !constants.IsInternalChannel(opts.Channel) { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ Channel: opts.Channel, - ChatID: opts.ChatID, + + ChatID: opts.ChatID, + Content: "Context window exceeded. Compressing history and retrying...", }) } al.forceCompression(agent, opts.SessionKey) + newHistory := agent.Sessions.GetHistory(opts.SessionKey) + newSummary := agent.Sessions.GetSummary(opts.SessionKey) + messages = agent.ContextBuilder.BuildMessages( + newHistory, newSummary, "", + nil, opts.Channel, opts.ChatID, ) + continue } + break } // Streaming finished — close the stream goroutine so it flushes + // the last update and exits cleanly before we process the response. + if streamDone != nil { // onChunk is captured by doCall closures; nil it to avoid + // writes after the channel is closed during retries. + onChunk = nil + close(streamCh) + <-streamDone + streamDone = nil } if err != nil { logger.ErrorCF("agent", "LLM call failed", + map[string]any{ - "agent_id": agent.ID, + "agent_id": agent.ID, + "iteration": iteration, - "error": err.Error(), + + "error": err.Error(), }) + return "", iteration, fmt.Errorf("LLM call failed after retries: %w", err) } // Record token usage + if response.Usage != nil && al.stats != nil { al.stats.RecordUsage( + response.Usage.PromptTokens, + response.Usage.CompletionTokens, + response.Usage.TotalTokens, ) } // Handle reasoning output (best-effort, non-blocking) + go al.handleReasoning(ctx, response.Reasoning, opts.Channel, al.targetReasoningChannelID(opts.Channel)) logger.DebugCF("agent", "LLM response", + map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "content_chars": len(response.Content), - "tool_calls": len(response.ToolCalls), - "reasoning": response.Reasoning, + "agent_id": agent.ID, + + "iteration": iteration, + + "content_chars": len(response.Content), + + "tool_calls": len(response.ToolCalls), + + "reasoning": response.Reasoning, + "target_channel": al.targetReasoningChannelID(opts.Channel), - "channel": opts.Channel, + + "channel": opts.Channel, }) // Detect repetition loop on raw text (before stripping think + // blocks so loops inside <think> are caught). Skip when the + // provider already returned native tool calls. + // Streaming providers may have already flagged repetition via + // FinishReason="repetition_detected" — honor that too. + if response.FinishReason == "repetition_detected" || + (len(response.ToolCalls) == 0 && utils.DetectRepetitionLoop(response.Content)) { logger.WarnCF("agent", "Repetition loop detected in LLM response, retrying", + map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, - "finish_reason": response.FinishReason, + "agent_id": agent.ID, + + "iteration": iteration, + + "finish_reason": response.FinishReason, + "content_length": len(response.Content), }) // Retry once: inject nudge message and re-call + savedMsgs := messages + messages = append(append([]providers.Message(nil), messages...), + providers.Message{ - Role: "user", + Role: "user", + Content: "[System] Your previous response contained degenerate repetition and was discarded. Please respond normally without repeating yourself.", }) + response, err = callLLM() + messages = savedMsgs // restore original messages if err != nil { @@ -2641,84 +3832,125 @@ func (al *AgentLoop) runLLMIteration( } // Re-check on raw text; if still repeating give up + if utils.DetectRepetitionLoop(response.Content) { logger.ErrorCF("agent", "Repetition persists after retry, returning empty", + map[string]any{"agent_id": agent.ID}) + response.Content = "" } } // Strip think blocks before extracting XML tool calls so + // extraction operates on clean content. + response.Content = utils.StripThinkBlocks(response.Content) // Recover XML tool calls emitted as plain text by some providers. + if len(response.ToolCalls) == 0 { if xmlCalls := providers.ExtractXMLToolCalls(response.Content); len(xmlCalls) > 0 { response.ToolCalls = xmlCalls } } + response.Content = providers.StripXMLToolCalls(response.Content) // Check if no tool calls - we're done + if len(response.ToolCalls) == 0 { // Plan continuation: if unchecked steps remain, nudge the LLM to + // either mark completed steps or continue working on them. + // This fires for both foreground and background plan execution, + // ensuring the loop doesn't exit prematurely after marking a step. + curUnchecked := 0 + if preUnchecked > 0 { curUnchecked = strings.Count(agent.ContextBuilder.ReadMemory(), "- [ ]") } + if curUnchecked > 0 && !planMarkNudged && + planSnapshot == "executing" { planMarkNudged = true + messages = append(messages, providers.Message{ - Role: "assistant", + Role: "assistant", + Content: response.Content, }) + var nudgeMsg string + if curUnchecked == preUnchecked { nudgeMsg = fmt.Sprintf("[System] %d unchecked steps remain in MEMORY.md and "+ + "none were marked [x] during this session. "+ + "If you completed any steps, use edit_file to mark them [x] now. "+ + "If steps are still in progress, continue working on them.", curUnchecked) } else { nudgeMsg = fmt.Sprintf("[System] Progress recorded. %d unchecked steps remain. "+ + "Continue working on the next step.", curUnchecked) } + messages = append(messages, providers.Message{ - Role: "user", + Role: "user", + Content: nudgeMsg, }) + logger.InfoCF("agent", "Nudging plan execution: continue plan steps", + map[string]any{"agent_id": agent.ID, "iteration": iteration, "unchecked": curUnchecked}) + continue } finalContent = response.Content + logger.InfoCF("agent", "LLM response without tool calls (direct answer)", + map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, + "agent_id": agent.ID, + + "iteration": iteration, + "content_chars": len(finalContent), }) + break } normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls)) + for _, tc := range response.ToolCalls { normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc)) } // --- Interview mode: reject disallowed tool calls before they + // enter messages or session history. Rejected calls are stripped + // from normalizedToolCalls so they never reach the assistant + // message, the tool-result list, or the session store. + // A single compact rejection message is injected instead. + var interviewRejected []string + if isPlanPreExecution(planSnapshot) { allowed := normalizedToolCalls[:0] // reuse backing array + for _, tc := range normalizedToolCalls { if isToolAllowedDuringInterview(tc.Name, tc.Arguments) { allowed = append(allowed, tc) @@ -2726,53 +3958,77 @@ func (al *AgentLoop) runLLMIteration( interviewRejected = append(interviewRejected, tc.Name) } } + normalizedToolCalls = allowed + if len(interviewRejected) > 0 { logger.InfoCF("agent", "Interview mode: rejected tool calls", + map[string]any{ "agent_id": agent.ID, + "rejected": interviewRejected, }) + messages = append(messages, providers.Message{ - Role: "user", + Role: "user", + Content: interviewRejectMessage, }) } + // If all tool calls were rejected, skip to next iteration. + if len(normalizedToolCalls) == 0 { continue } } // Log tool calls + toolNames := make([]string, 0, len(normalizedToolCalls)) + for _, tc := range normalizedToolCalls { toolNames = append(toolNames, tc.Name) } + logger.InfoCF("agent", "LLM requested tool calls", + map[string]any{ - "agent_id": agent.ID, - "tools": toolNames, - "count": len(normalizedToolCalls), + "agent_id": agent.ID, + + "tools": toolNames, + + "count": len(normalizedToolCalls), + "iteration": iteration, }) // Publish rich status update + if !constants.IsInternalChannel(opts.Channel) && task != nil { // Add pending entries to tool log for the current tool calls + task.mu.Lock() + for _, tc := range normalizedToolCalls { task.toolLog = append(task.toolLog, toolLogEntry{ - Name: fmt.Sprintf("[%d] %s", iteration, tc.Name), + Name: fmt.Sprintf("[%d] %s", iteration, tc.Name), + ArgsSnip: buildArgsSnippet(tc.Name, tc.Arguments, agent.Workspace), - Result: "\u23F3", + + Result: "\u23F3", }) + // Detect project directory + if task.projectDir == "" && tc.Name == "exec" { task.projectDir = extractExecProjectDir(tc.Arguments) } + switch tc.Name { case "read_file", "write_file", "edit_file", "append_file", "list_dir": + if p, _ := tc.Arguments["path"].(string); p != "" { if rel := fileParentRelDir(p, agent.Workspace); rel != "" { if task.fileCommonDir == "" { @@ -2784,333 +4040,519 @@ func (al *AgentLoop) runLLMIteration( } } } + task.mu.Unlock() statusContent := buildRichStatus(task, isBackground, agent.Workspace) + if isBackground { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: statusContent, + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: statusContent, + IsTaskStatus: true, - TaskID: opts.TaskID, + + TaskID: opts.TaskID, }) } else { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ - Channel: opts.Channel, - ChatID: opts.ChatID, - Content: statusContent, + Channel: opts.Channel, + + ChatID: opts.ChatID, + + Content: statusContent, + IsStatus: true, }) } } // Record session activity for heartbeat/plan coordination + for _, tc := range normalizedToolCalls { var detectedDir string + if tc.Name == "exec" { detectedDir = extractExecProjectDir(tc.Arguments) } + if detectedDir == "" { switch tc.Name { case "read_file", "write_file", "edit_file", "append_file", "list_dir": + if p, _ := tc.Arguments["path"].(string); p != "" { detectedDir = fileParentRelDir(p, agent.Workspace) } } } + if detectedDir != "" { meta := &TouchMeta{ ProjectPath: agent.ContextBuilder.GetPlanWorkDir(), - Purpose: utils.Truncate(opts.UserMessage, 80), - Branch: agent.GetWorktreeBranch(opts.SessionKey), + + Purpose: utils.Truncate(opts.UserMessage, 80), + + Branch: agent.GetWorktreeBranch(opts.SessionKey), } + if meta.ProjectPath == "" { meta.ProjectPath = agent.Workspace } + al.sessions.Touch(opts.SessionKey, opts.Channel, opts.ChatID, detectedDir, meta) } } // Build assistant message with tool calls + assistantMsg := providers.Message{ - Role: "assistant", - Content: response.Content, + Role: "assistant", + + Content: response.Content, + ReasoningContent: response.ReasoningContent, } + for _, tc := range normalizedToolCalls { // Copy ExtraContent to ensure thought_signature is persisted for Gemini 3 + extraContent := tc.ExtraContent + thoughtSignature := "" + if tc.Function != nil { thoughtSignature = tc.Function.ThoughtSignature } assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, providers.ToolCall{ - ID: tc.ID, - Type: "function", - Name: tc.Name, + ID: tc.ID, + + Type: "function", + + Name: tc.Name, + Arguments: tc.Arguments, + Function: &providers.FunctionCall{ - Name: tc.Name, - Arguments: tc.Arguments, + Name: tc.Name, + + Arguments: tc.Arguments, + ThoughtSignature: thoughtSignature, }, - ExtraContent: extraContent, + + ExtraContent: extraContent, + ThoughtSignature: thoughtSignature, }) } + messages = append(messages, assistantMsg) // Save assistant message with tool calls to session + agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg) // Execute tool calls + var lastBlocker string + for tcIdx, tc := range normalizedToolCalls { argsJSON, _ := json.Marshal(tc.Arguments) + argsPreview := utils.Truncate(string(argsJSON), 200) + logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), + map[string]any{ - "agent_id": agent.ID, - "tool": tc.Name, + "agent_id": agent.ID, + + "tool": tc.Name, + "iteration": iteration, }) // Heartbeat lazy worktree: create worktree on first write-tool call + if opts.Background && isWriteTool(tc.Name) && !agent.IsInWorktree(opts.SessionKey) { taskName := "heartbeat-" + time.Now().Format("20060102") + hbDir := agent.ContextBuilder.GetPlanWorkDir() + if wt, err := agent.ActivateWorktree(opts.SessionKey, taskName, hbDir); err == nil { logger.InfoCF("agent", "Heartbeat worktree created", map[string]any{"branch": wt.Branch}) } } // Create async callback for tools that implement AsyncTool. + // The callback publishes a system inbound message so processSystemMessage + // injects the result into the conductor's session history. The conductor + // sees it on its next turn and decides whether to notify the user. + toolName := tc.Name // capture for goroutine + asyncCallback := func(callbackCtx context.Context, result *tools.ToolResult) { content := result.ForLLM + if content == "" { content = result.ForUser } + if content == "" { return } logger.InfoCF("agent", "Async tool completed, publishing to conductor", + map[string]any{ - "tool": toolName, + "tool": toolName, + "content_len": len(content), - "is_error": result.IsError, + + "is_error": result.IsError, }) pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + _ = al.bus.PublishInbound(pubCtx, bus.InboundMessage{ - Channel: "system", + Channel: "system", + SenderID: fmt.Sprintf("async:%s", toolName), - ChatID: fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID), - Content: fmt.Sprintf("Async tool '%s' completed.\n\nResult:\n%s", toolName, content), + + ChatID: fmt.Sprintf("%s:%s", opts.Channel, opts.ChatID), + + Content: fmt.Sprintf("Async tool '%s' completed.\n\nResult:\n%s", toolName, content), }) } // Report toolcall state to canvas. + al.reporter().ReportStateChange(opts.SessionKey, orch.AgentStateToolCall, tc.Name) toolStart := time.Now() + toolCtx := ctx + if wt := agent.GetWorktree(opts.SessionKey); wt != nil { toolCtx = tools.WithWorkspaceOverride(toolCtx, wt.Path) + toolCtx = tools.WithWorktreeInfo(toolCtx, wt) } + toolResult := agent.Tools.ExecuteWithContext( + toolCtx, + tc.Name, + tc.Arguments, + opts.Channel, + opts.ChatID, + asyncCallback, ) + toolDuration := time.Since(toolStart) // Update tool log entry with result + if task != nil { task.mu.Lock() + // Find the matching pending entry (added earlier in this iteration) + logIdx := len(task.toolLog) - len(normalizedToolCalls) + tcIdx + if logIdx >= 0 && logIdx < len(task.toolLog) { if toolResult.IsError || toolResult.Err != nil { task.toolLog[logIdx].Result = fmt.Sprintf("\u2717 %.1fs", toolDuration.Seconds()) + // Extract error detail for block display + if toolResult.Err != nil { task.toolLog[logIdx].ErrDetail = utils.Truncate(toolResult.Err.Error(), 300) } else if toolResult.ForLLM != "" { // exec returns IsError with exit info in ForLLM, not Err + // Show last few lines (stderr / exit code) + lines := strings.Split(strings.TrimSpace(toolResult.ForLLM), "\n") + start := len(lines) - 3 + if start < 0 { start = 0 } + task.toolLog[logIdx].ErrDetail = utils.Truncate( + strings.Join(lines[start:], "\n"), 300) } + // Sticky error: remember most recent error for persistent display + entry := task.toolLog[logIdx] + task.lastError = &entry } else { task.toolLog[logIdx].Result = fmt.Sprintf("\u2713 %.1fs", toolDuration.Seconds()) } } + task.mu.Unlock() } // Send ForUser content to user immediately if not Silent + if !toolResult.Silent && toolResult.ForUser != "" && opts.SendResponse { _ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{ Channel: opts.Channel, - ChatID: opts.ChatID, + + ChatID: opts.ChatID, + Content: toolResult.ForUser, }) + logger.DebugCF("agent", "Sent tool result to user", + map[string]any{ - "tool": tc.Name, + "tool": tc.Name, + "content_len": len(toolResult.ForUser), }) } // If tool returned media refs, publish them as outbound media + if len(toolResult.Media) > 0 && opts.SendResponse { parts := make([]bus.MediaPart, 0, len(toolResult.Media)) + for _, ref := range toolResult.Media { part := bus.MediaPart{Ref: ref} + // Populate metadata from MediaStore when available + if al.mediaStore != nil { if _, meta, err := al.mediaStore.ResolveWithMeta(ref); err == nil { part.Filename = meta.Filename + part.ContentType = meta.ContentType + part.Type = inferMediaType(meta.Filename, meta.ContentType) } } + parts = append(parts, part) } + al.bus.PublishOutboundMedia(ctx, bus.OutboundMediaMessage{ Channel: opts.Channel, - ChatID: opts.ChatID, - Parts: parts, + + ChatID: opts.ChatID, + + Parts: parts, }) } // Determine content for LLM based on tool result + contentForLLM := toolResult.ForLLM + if contentForLLM == "" && toolResult.Err != nil { contentForLLM = toolResult.Err.Error() } // Track blockers for task reminder + if toolResult.IsError || toolResult.Err != nil { lastBlocker = contentForLLM } toolResultMsg := providers.Message{ - Role: "tool", - Content: contentForLLM, + Role: "tool", + + Content: contentForLLM, + ToolCallID: tc.ID, } + messages = append(messages, toolResultMsg) // Save tool result message to session + agent.Sessions.AddFullMessage(opts.SessionKey, toolResultMsg) } // Trim tool log sliding window to prevent unbounded growth + if task != nil { task.mu.Lock() + if len(task.toolLog) > maxToolLogEntries { task.toolLog = task.toolLog[len(task.toolLog)-maxToolLogEntries:] } + task.mu.Unlock() } // Inject ephemeral task reminder to prevent focus drift. + // Remove previous reminder and re-append at the tail so it stays + // close to the LLM's attention window. + if shouldInjectReminder(iteration, agent.TaskReminderInterval) && !opts.NoHistory { if lastReminderIdx >= 0 && lastReminderIdx < len(messages) { messages = append(messages[:lastReminderIdx], messages[lastReminderIdx+1:]...) } + reminderMsg := buildTaskReminder(opts.UserMessage, lastBlocker) + messages = append(messages, reminderMsg) + lastReminderIdx = len(messages) - 1 + logger.DebugCF("agent", "Injected task reminder", + map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, + "agent_id": agent.ID, + + "iteration": iteration, + "has_blocker": lastBlocker != "", }) } // Inject plan-mode reminder to keep AI focused on interview/review workflow. + if iteration > 1 && isPlanPreExecution(planSnapshot) { if reminder, ok := buildPlanReminder(planSnapshot); ok { messages = append(messages, reminder) + logger.DebugCF("agent", "Injected plan reminder", + map[string]any{ - "agent_id": agent.ID, - "iteration": iteration, + "agent_id": agent.ID, + + "iteration": iteration, + "plan_status": planSnapshot, }) } } // Inject orchestration nudge during plan execution to encourage spawn usage. + if planSnapshot == "executing" && agent.Subagents != nil && agent.Subagents.Enabled { if reminder, ok := buildOrchReminder(iteration); ok { messages = append(messages, reminder) + logger.DebugCF("agent", "Injected orchestration nudge", + map[string]any{ - "agent_id": agent.ID, + "agent_id": agent.ID, + "iteration": iteration, }) } } + // Inject pending subagent questions/plan reviews for the conductor to answer. + + if agent.SubagentMgr != nil { + for _, q := range agent.SubagentMgr.PendingQuestions() { + var content string + + switch q.Type { + case "plan_review": + + content = fmt.Sprintf( + "[Subagent %s submitted a plan for review]:\n%s\nRespond using the review_subagent_plan tool with task_id=%q.", + q.TaskID, + q.Content, + q.TaskID, + ) + + default: + + content = fmt.Sprintf( + "[Subagent %s asks]: %s\nRespond using the answer_subagent tool with task_id=%q.", + q.TaskID, + q.Content, + q.TaskID, + ) + } + + messages = append(messages, providers.Message{ + Role: "user", + + Content: content, + }) + } + } + // Refresh system prompt: tool execution may have changed workDir, + // memory, plan status, etc. Update messages[0] so the next LLM + // call sees the current state. + if touchDir := al.sessions.GetTouchDir(opts.SessionKey); touchDir != "" { agent.ContextBuilder.SetWorkDir(filepath.Join(agent.Workspace, touchDir)) } + if newPrompt := agent.ContextBuilder.BuildSystemPrompt(); len(messages) > 0 && + messages[0].Content != newPrompt { messages[0].Content = newPrompt + al.lastSystemPrompt.Store(newPrompt) + al.promptDirty.Store(false) } } // If max iterations exhausted with tool calls still pending, + // make one final LLM call without tools to force a text response. + if finalContent == "" && iteration >= maxIter { logger.WarnCF("agent", "Max iterations reached, forcing final response without tools", + map[string]any{ - "agent_id": agent.ID, + "agent_id": agent.ID, + "iteration": iteration, }) + forceResp, forceErr := agent.Provider.Chat(ctx, messages, nil, agent.Model, map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, + "max_tokens": agent.MaxTokens, + + "temperature": agent.Temperature, + "prompt_cache_key": agent.ID, }) + if forceErr == nil && forceResp.Content != "" { finalContent = utils.StripThinkBlocks(forceResp.Content) + if forceResp.Usage != nil && al.stats != nil { al.stats.RecordUsage( + forceResp.Usage.PromptTokens, + forceResp.Usage.CompletionTokens, + forceResp.Usage.TotalTokens, ) } @@ -3121,18 +4563,22 @@ func (al *AgentLoop) runLLMIteration( } // updateToolContexts updates the context for tools that need channel/chatID info. + func (al *AgentLoop) updateToolContexts(agent *AgentInstance, channel, chatID string) { // Use ContextualTool interface instead of type assertions + if tool, ok := agent.Tools.Get("message"); ok { if mt, ok := tool.(tools.ContextualTool); ok { mt.SetContext(channel, chatID) } } + if tool, ok := agent.Tools.Get("spawn"); ok { if st, ok := tool.(tools.ContextualTool); ok { st.SetContext(channel, chatID) } } + if tool, ok := agent.Tools.Get("subagent"); ok { if st, ok := tool.(tools.ContextualTool); ok { st.SetContext(channel, chatID) @@ -3141,22 +4587,31 @@ func (al *AgentLoop) updateToolContexts(agent *AgentInstance, channel, chatID st } // maybeSummarize triggers summarization if the session history exceeds thresholds. + func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, chatID string) { newHistory := agent.Sessions.GetHistory(sessionKey) + tokenEstimate := al.estimateTokens(newHistory) + threshold := agent.ContextWindow * 75 / 100 if len(newHistory) > 20 || tokenEstimate > threshold { summarizeKey := agent.ID + ":" + sessionKey + if _, loading := al.summarizing.LoadOrStore(summarizeKey, true); !loading { go func() { defer al.summarizing.Delete(summarizeKey) + logger.InfoCF("agent", "Memory threshold reached, optimizing conversation history", + map[string]any{ - "session_key": sessionKey, - "history_len": len(newHistory), + "session_key": sessionKey, + + "history_len": len(newHistory), + "token_estimate": tokenEstimate, }) + al.summarizeSession(agent, sessionKey) }() } @@ -3164,260 +4619,365 @@ func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, c } // forceCompression aggressively reduces context when the limit is hit. + // It drops the oldest 50% of messages (keeping system prompt and last user message). + func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) { history := agent.Sessions.GetHistory(sessionKey) + if len(history) <= 4 { return } // Keep system prompt (usually [0]) and the very last message (user's trigger) + // We want to drop the oldest half of the *conversation* + // Assuming [0] is system, [1:] is conversation + conversation := history[1 : len(history)-1] + if len(conversation) == 0 { return } // Helper to find the mid-point of the conversation + mid := len(conversation) / 2 // New history structure: + // 1. System Prompt (with compression note appended) + // 2. Second half of conversation + // 3. Last message droppedCount := mid + keptConversation := conversation[mid:] newHistory := make([]providers.Message, 0, 1+len(keptConversation)+1) // Append compression note to the original system prompt instead of adding a new system message + // This avoids having two consecutive system messages which some APIs (like Zhipu) reject + compressionNote := fmt.Sprintf( + "\n\n[System Note: Emergency compression dropped %d oldest messages due to context limit]", + droppedCount, ) + enhancedSystemPrompt := history[0] + enhancedSystemPrompt.Content = enhancedSystemPrompt.Content + compressionNote + newHistory = append(newHistory, enhancedSystemPrompt) newHistory = append(newHistory, keptConversation...) + newHistory = append(newHistory, history[len(history)-1]) // Last message // Update session + agent.Sessions.SetHistory(sessionKey, newHistory) + agent.Sessions.Save(sessionKey) logger.WarnCF("agent", "Forced compression executed", map[string]any{ - "session_key": sessionKey, + "session_key": sessionKey, + "dropped_msgs": droppedCount, - "new_count": len(newHistory), + + "new_count": len(newHistory), }) } // GetStartupInfo returns information about loaded tools and skills for logging. + func (al *AgentLoop) GetStartupInfo() map[string]any { info := make(map[string]any) agent := al.registry.GetDefaultAgent() + if agent == nil { return info } // Tools info + toolsList := agent.Tools.List() + toolsMap := map[string]any{ "count": len(toolsList), + "names": toolsList, } + // Report web search provider if registered + if t, ok := agent.Tools.Get("web_search"); ok { if wst, ok := t.(*tools.WebSearchTool); ok { toolsMap["web_search_provider"] = wst.ProviderName() } } + info["tools"] = toolsMap // Skills info + info["skills"] = agent.ContextBuilder.GetSkillsInfo() // Agents info + info["agents"] = map[string]any{ "count": len(al.registry.ListAgentIDs()), - "ids": al.registry.ListAgentIDs(), + + "ids": al.registry.ListAgentIDs(), } return info } // ListSkills returns all available skills from the default agent. + func (al *AgentLoop) ListSkills() []skills.SkillInfo { agent := al.registry.GetDefaultAgent() + if agent == nil { return nil } + return agent.ContextBuilder.ListSkills() } // GetPlanInfo returns plan state from the default agent's memory store. + func (al *AgentLoop) GetPlanInfo() (hasPlan bool, status string, currentPhase, totalPhases int, display string, memory string) { agent := al.registry.GetDefaultAgent() + if agent == nil { return false, "", 0, 0, "No agent available.", "" } + mem := agent.ContextBuilder.Memory() + if mem == nil { return false, "", 0, 0, "No memory store.", "" } + hasPlan = mem.HasActivePlan() + status = mem.GetPlanStatus() + currentPhase = mem.GetCurrentPhase() + totalPhases = mem.GetTotalPhases() + display = mem.FormatPlanDisplay() + memory = mem.ReadLongTerm() + return hasPlan, status, currentPhase, totalPhases, display, memory } // GetPlanStatus returns the current plan status ("interviewing", "executing", "review", etc.) or "". + func (al *AgentLoop) GetPlanStatus() string { agent := al.registry.GetDefaultAgent() + if agent == nil { return "" } + return agent.ContextBuilder.GetPlanStatus() } // GetPlanPhases returns structured phase/step data from the default agent's plan. + func (al *AgentLoop) GetPlanPhases() []PlanPhase { agent := al.registry.GetDefaultAgent() + if agent == nil { return nil } + mem := agent.ContextBuilder.Memory() + if mem == nil { return nil } + return mem.GetPlanPhases() } // GetActiveSessions returns currently active sessions for the mini app API. + func (al *AgentLoop) GetActiveSessions() []SessionEntry { return al.sessions.ListActive() } // GetSessionStats returns the current session statistics snapshot, or nil if stats tracking is disabled. + func (al *AgentLoop) GetSessionStats() *stats.Stats { if al.stats == nil { return nil } + s := al.stats.GetStats() + return &s } // GetContextInfo returns the bootstrap file resolution and directory context for the default agent. + func (al *AgentLoop) GetContextInfo() (workDir, planWorkDir, workspace string, bootstrap []BootstrapFileInfo) { agent := al.registry.GetDefaultAgent() + if agent == nil { return "", "", "", nil } + workspace = agent.Workspace + planWorkDir = agent.ContextBuilder.GetPlanWorkDir() + // Use the most recent active session's touch_dir (tool-detected project directory) + if active := al.sessions.ListActive(); len(active) > 0 && active[0].TouchDir != "" { workDir = active[0].TouchDir } else { workDir = agent.ContextBuilder.workDir } + bootstrap = agent.ContextBuilder.ResolveBootstrapPaths() + return workDir, planWorkDir, workspace, bootstrap } // GetSystemPrompt returns the system prompt last sent to the LLM. + // If the prompt is dirty (state changed since last capture), it rebuilds + // from current state. Falls back to building if no LLM call has occurred yet. + func (al *AgentLoop) GetSystemPrompt() string { if !al.promptDirty.Load() { if v := al.lastSystemPrompt.Load(); v != nil { return v.(string) } } + // Rebuild from current state + agent := al.registry.GetDefaultAgent() + if agent == nil { return "" } + prompt := agent.ContextBuilder.BuildSystemPrompt() + al.lastSystemPrompt.Store(prompt) + al.promptDirty.Store(false) + return prompt } // formatMessagesForLog formats messages for logging + func formatMessagesForLog(messages []providers.Message) string { if len(messages) == 0 { return "[]" } var sb strings.Builder + sb.WriteString("[\n") + for i, msg := range messages { fmt.Fprintf(&sb, " [%d] Role: %s\n", i, msg.Role) + if len(msg.ToolCalls) > 0 { sb.WriteString(" ToolCalls:\n") + for _, tc := range msg.ToolCalls { fmt.Fprintf(&sb, " - ID: %s, Type: %s, Name: %s\n", tc.ID, tc.Type, tc.Name) + args := tc.Arguments + if len(args) == 0 && tc.Function != nil { args = tc.Function.Arguments } + if len(args) > 0 { argsJSON, _ := json.Marshal(args) + fmt.Fprintf(&sb, " Arguments: %s\n", utils.Truncate(string(argsJSON), 200)) } } } + if msg.Content != "" { content := utils.Truncate(msg.Content, 200) + fmt.Fprintf(&sb, " Content: %s\n", content) } + if msg.ToolCallID != "" { fmt.Fprintf(&sb, " ToolCallID: %s\n", msg.ToolCallID) } + sb.WriteString("\n") } + sb.WriteString("]") + return sb.String() } // formatToolsForLog formats tool definitions for logging + func formatToolsForLog(toolDefs []providers.ToolDefinition) string { if len(toolDefs) == 0 { return "[]" } var sb strings.Builder + sb.WriteString("[\n") + for i, tool := range toolDefs { fmt.Fprintf(&sb, " [%d] Type: %s, Name: %s\n", i, tool.Type, tool.Function.Name) + fmt.Fprintf(&sb, " Description: %s\n", tool.Function.Description) + if len(tool.Function.Parameters) > 0 { fmt.Fprintf(&sb, " Parameters: %s\n", utils.Truncate(string(tool.Function.Parameters), 200)) } } + sb.WriteString("]") + return sb.String() } // summarizeSession summarizes the conversation history for a session. + func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) + defer cancel() history := agent.Sessions.GetHistory(sessionKey) + summary := agent.Sessions.GetSummary(sessionKey) // Keep last 4 messages for continuity + if len(history) <= 4 { return } @@ -3425,19 +4985,26 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { toSummarize := history[:len(history)-4] // Oversized Message Guard + maxMessageTokens := agent.ContextWindow / 2 + validMessages := make([]providers.Message, 0) + omitted := false for _, m := range toSummarize { if m.Role != "user" && m.Role != "assistant" { continue } + msgTokens := len(m.Content) / 2 + if msgTokens > maxMessageTokens { omitted = true + continue } + validMessages = append(validMessages, m) } @@ -3446,31 +5013,48 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { } // Multi-Part Summarization + var finalSummary string + if len(validMessages) > 10 { mid := len(validMessages) / 2 + part1 := validMessages[:mid] + part2 := validMessages[mid:] s1, _ := al.summarizeBatch(ctx, agent, part1, "") + s2, _ := al.summarizeBatch(ctx, agent, part2, "") mergePrompt := fmt.Sprintf( + "Merge these two conversation summaries into one cohesive summary:\n\n1: %s\n\n2: %s", + s1, + s2, ) + resp, err := agent.Provider.Chat( + ctx, + []providers.Message{{Role: "user", Content: mergePrompt}}, + nil, + agent.Model, + map[string]any{ - "max_tokens": 1024, - "temperature": 0.3, + "max_tokens": 1024, + + "temperature": 0.3, + "prompt_cache_key": agent.ID, }, ) + if err == nil { finalSummary = resp.Content } else { @@ -3487,170 +5071,250 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { if finalSummary != "" { if err := agent.Sessions.CompactOldTurns(sessionKey, 4, finalSummary); err != nil { logger.ErrorCF("agent", "CompactOldTurns failed, falling back", + map[string]any{"error": err.Error()}) + agent.Sessions.SetSummary(sessionKey, finalSummary) + agent.Sessions.TruncateHistory(sessionKey, 4) + agent.Sessions.Save(sessionKey) } } } // summarizeBatch summarizes a batch of messages. + func (al *AgentLoop) summarizeBatch( ctx context.Context, + agent *AgentInstance, + batch []providers.Message, + existingSummary string, ) (string, error) { var sb strings.Builder + sb.WriteString("Provide a concise summary of this conversation segment, preserving core context and key points.\n") + if agent.ContextBuilder.HasActivePlan() { sb.WriteString("Note: Active plan in MEMORY.md. Preserve plan progress references.\n") } + if existingSummary != "" { sb.WriteString("Existing context: ") + sb.WriteString(existingSummary) + sb.WriteString("\n") } + sb.WriteString("\nCONVERSATION:\n") + for _, m := range batch { fmt.Fprintf(&sb, "%s: %s\n", m.Role, m.Content) } + prompt := sb.String() response, err := agent.Provider.Chat( + ctx, + []providers.Message{{Role: "user", Content: prompt}}, + nil, + agent.Model, + map[string]any{ - "max_tokens": 1024, - "temperature": 0.3, + "max_tokens": 1024, + + "temperature": 0.3, + "prompt_cache_key": agent.ID, }, ) if err != nil { return "", err } + return response.Content, nil } // estimateTokens estimates the number of tokens in a message list. + // Uses a safe heuristic of 2.5 characters per token to account for CJK and other + // overheads better than the previous 3 chars/token. + func (al *AgentLoop) estimateTokens(messages []providers.Message) int { totalChars := 0 + for _, m := range messages { totalChars += utf8.RuneCountInString(m.Content) } + // 2.5 chars per token = totalChars * 2 / 5 + return totalChars * 2 / 5 } func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage) (string, bool) { content := strings.TrimSpace(msg.Content) + if !strings.HasPrefix(content, "/") { return "", false } parts := strings.Fields(content) + if len(parts) == 0 { return "", false } cmd := parts[0] + args := parts[1:] switch cmd { case "/show": + if len(args) < 1 { return "Usage: /show [model|channel|agents]", true } + switch args[0] { case "model": + defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { return "No default agent configured", true } + return fmt.Sprintf("Current model: %s", defaultAgent.Model), true + case "channel": + return fmt.Sprintf("Current channel: %s", msg.Channel), true + case "agents": + agentIDs := al.registry.ListAgentIDs() + return fmt.Sprintf("Registered agents: %s", strings.Join(agentIDs, ", ")), true + default: + return fmt.Sprintf("Unknown show target: %s", args[0]), true } case "/list": + if len(args) < 1 { return "Usage: /list [models|channels|agents]", true } + switch args[0] { case "models": + return "Available models: configured in config.json per agent", true + case "channels": + if al.channelManager == nil { return "Channel manager not initialized", true } + channels := al.channelManager.GetEnabledChannels() + if len(channels) == 0 { return "No channels enabled", true } + return fmt.Sprintf("Enabled channels: %s", strings.Join(channels, ", ")), true + case "agents": + agentIDs := al.registry.ListAgentIDs() + return fmt.Sprintf("Registered agents: %s", strings.Join(agentIDs, ", ")), true + default: + return fmt.Sprintf("Unknown list target: %s", args[0]), true } case "/switch": + if len(args) < 3 || args[1] != "to" { return "Usage: /switch [model|channel] to <name>", true } + target := args[0] + value := args[2] switch target { case "model": + defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { return "No default agent configured", true } + oldModel := defaultAgent.Model + defaultAgent.Model = value + return fmt.Sprintf("Switched model from %s to %s", oldModel, value), true + case "channel": + if al.channelManager == nil { return "Channel manager not initialized", true } + if _, exists := al.channelManager.GetChannel(value); !exists && value != "cli" { return fmt.Sprintf("Channel '%s' not found or not enabled", value), true } + return fmt.Sprintf("Switched target channel to %s", value), true + default: + return fmt.Sprintf("Unknown switch target: %s", target), true } case "/session": + return al.handleSessionCommand(args, msg.SessionKey), true case "/skills": + return al.handleSkillsCommand(), true case "/plan": + resp, handled := al.handlePlanCommand(args, msg.SessionKey) + if handled { al.notifyStateChange() } + return resp, handled case "/heartbeat": + resp, handled := al.handleHeartbeatCommand(args, msg) + if handled { al.notifyStateChange() } + return resp, handled } @@ -3675,33 +5339,45 @@ func (al *AgentLoop) handleHeartbeatCommand(args []string, msg bus.InboundMessag } baseChatID, currentThreadID := splitChatAndThread(msg.ChatID) + if baseChatID == "" { return "Unable to detect Telegram chat ID for heartbeat routing.", true } arg := strings.ToLower(strings.TrimSpace(args[1])) + var threadID int + var err error switch arg { case "off", "disable", "clear": + threadID = 0 + case "here", "this": + if currentThreadID <= 0 { return "Current Telegram message is not in a thread. Usage: /heartbeat thread <thread_id>", true } + threadID = currentThreadID + default: + threadID, err = strconv.Atoi(arg) + if err != nil || threadID < 0 { return "Usage: /heartbeat thread [here|off|<thread_id>]", true } } al.cfg.Channels.Telegram.HeartbeatThreadID = threadID + if al.state != nil { _ = al.state.SetHeartbeatTarget(fmt.Sprintf("telegram:%s", baseChatID)) } + if al.onHeartbeatThreadUpdate != nil { al.onHeartbeatThreadUpdate(threadID) } @@ -3715,148 +5391,219 @@ func (al *AgentLoop) handleHeartbeatCommand(args []string, msg bus.InboundMessag if threadID == 0 { return fmt.Sprintf("Heartbeat thread routing disabled for chat %s and saved to config.json.", baseChatID), true } + return fmt.Sprintf("Heartbeat thread set to %d for chat %s and saved to config.json.", threadID, baseChatID), true } func splitChatAndThread(chatID string) (baseChatID string, threadID int) { baseChatID = strings.TrimSpace(chatID) + if baseChatID == "" { return "", 0 } + if slash := strings.Index(baseChatID, "/"); slash >= 0 { threadPart := strings.TrimSpace(baseChatID[slash+1:]) + baseChatID = strings.TrimSpace(baseChatID[:slash]) + if tid, err := strconv.Atoi(threadPart); err == nil && tid > 0 { threadID = tid } } + return baseChatID, threadID } // handleSessionCommand dispatches /session subcommands. + func (al *AgentLoop) handleSessionCommand(args []string, sessionKey string) string { sub := "" + if len(args) > 0 { sub = strings.ToLower(strings.TrimSpace(args[0])) } + switch sub { case "list": + return al.handleSessionList() + case "graph": + return al.handleSessionGraph() + case "fork": + return al.handleSessionFork(args[1:], sessionKey) + case "reset": + if al.stats == nil { return "Stats tracking is disabled." } + al.stats.Reset() + return "Session statistics have been reset." + default: + return al.handleSessionStats() } } func (al *AgentLoop) handleSessionStats() string { agent := al.registry.GetDefaultAgent() + store := agent.Sessions.Store() // Session DAG summary + sessions, _ := store.List(nil) + var sb strings.Builder + fmt.Fprintf(&sb, "Sessions: %d in store\n", len(sessions)) + if len(sessions) > 0 { active, completed := 0, 0 + for _, s := range sessions { switch s.Status { case "active": + active++ + case "completed": + completed++ } } + fmt.Fprintf(&sb, " active=%d completed=%d\n", active, completed) } + sb.WriteString("\nUse: /session list | graph | fork [label]\n") // Token stats if available + if al.stats != nil { s := al.stats.GetStats() + fmt.Fprintf(&sb, + "\nToken Stats — Today (%s):\n Prompts: %d LLM calls: %d Tokens: %s (in: %s, out: %s)\n"+ + "All time (since %s):\n Prompts: %d LLM calls: %d Tokens: %s (in: %s, out: %s)", + s.Today.Date, + s.Today.Prompts, + s.Today.Requests, + stats.FormatTokenCount(s.Today.TotalTokens), + stats.FormatTokenCount(s.Today.PromptTokens), + stats.FormatTokenCount(s.Today.CompletionTokens), + s.Since.Format("2006-01-02"), + s.TotalPrompts, + s.TotalRequests, + stats.FormatTokenCount(s.TotalTokens), + stats.FormatTokenCount(s.TotalPromptTokens), + stats.FormatTokenCount(s.TotalCompletionTokens), ) } + return sb.String() } // shortSessionKey truncates long session keys for display. + func shortSessionKey(key string) string { parts := strings.Split(key, ":") + if len(parts) > 2 { return strings.Join(parts[2:], ":") } + return key } func (al *AgentLoop) handleSessionList() string { agent := al.registry.GetDefaultAgent() + store := agent.Sessions.Store() + sessions, err := store.List(nil) if err != nil { return fmt.Sprintf("Error listing sessions: %v", err) } + if len(sessions) == 0 { return "No sessions in store." } var sb strings.Builder + fmt.Fprintf(&sb, "Sessions (%d)\n", len(sessions)) + for _, s := range sessions { age := time.Since(s.UpdatedAt).Truncate(time.Second) + label := s.Label + if label == "" { label = shortSessionKey(s.Key) } + parent := "" + if s.ParentKey != "" { parent = " parent=" + shortSessionKey(s.ParentKey) } + fmt.Fprintf(&sb, "- %s [%s] (%s) turns=%d%s\n", + label, s.Status, age, s.TurnCount, parent) } + return sb.String() } func (al *AgentLoop) handleSessionGraph() string { agent := al.registry.GetDefaultAgent() + store := agent.Sessions.Store() + sessions, err := store.List(nil) if err != nil { return fmt.Sprintf("Error listing sessions: %v", err) } + if len(sessions) == 0 { return "No sessions in store." } // Build parent→children map and find roots + byKey := make(map[string]*session.SessionInfo, len(sessions)) + children := make(map[string][]string) + var roots []string + for _, s := range sessions { byKey[s.Key] = s + if s.ParentKey == "" { roots = append(roots, s.Key) } else { @@ -3865,39 +5612,60 @@ func (al *AgentLoop) handleSessionGraph() string { } var sb strings.Builder + sb.WriteString("Session Graph\n") + for i, root := range roots { last := i == len(roots)-1 + printSessionTree(&sb, root, byKey, children, "", last) } + return sb.String() } -func printSessionTree(sb *strings.Builder, key string, byKey map[string]*session.SessionInfo, children map[string][]string, prefix string, last bool) { +func printSessionTree( + sb *strings.Builder, + key string, + byKey map[string]*session.SessionInfo, + children map[string][]string, + prefix string, + last bool, +) { s := byKey[key] + if s == nil { return } connector := "├── " + if last { connector = "└── " } + icon := "●" + if s.Status == "completed" { icon = "✓" } + label := s.Label + if label == "" { label = shortSessionKey(s.Key) } + fmt.Fprintf(sb, "%s%s%s %s (turns=%d)\n", prefix, connector, icon, label, s.TurnCount) childPrefix := prefix + "│ " + if last { childPrefix = prefix + " " } + kids := children[key] + for i, childKey := range kids { printSessionTree(sb, childKey, byKey, children, childPrefix, i == len(kids)-1) } @@ -3909,88 +5677,131 @@ func (al *AgentLoop) handleSessionFork(args []string, sessionKey string) string } agent := al.registry.GetDefaultAgent() + store := agent.Sessions.Store() label := "fork" + if len(args) > 0 { label = strings.Join(args, " ") } childKey := sessionKey + ":fork:" + time.Now().Format("20060102T150405") + err := store.Fork(sessionKey, childKey, &session.CreateOpts{Label: label}) if err != nil { return fmt.Sprintf("Fork failed: %v", err) } - return fmt.Sprintf("Forked session\n parent: %s\n child: %s", shortSessionKey(sessionKey), shortSessionKey(childKey)) + + return fmt.Sprintf( + "Forked session\n parent: %s\n child: %s", + shortSessionKey(sessionKey), + shortSessionKey(childKey), + ) } // SessionGraphNode represents a session node for the Mini App graph API. + type SessionGraphNode struct { - Key string `json:"key"` - Label string `json:"label"` - Status string `json:"status"` - Summary string `json:"summary"` - ParentKey string `json:"parent_key"` - ForkTurnID string `json:"fork_turn_id"` - TurnCount int `json:"turn_count"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + Key string `json:"key"` + + Label string `json:"label"` + + Status string `json:"status"` + + Summary string `json:"summary"` + + ParentKey string `json:"parent_key"` + + ForkTurnID string `json:"fork_turn_id"` + + TurnCount int `json:"turn_count"` + + CreatedAt time.Time `json:"created_at"` + + UpdatedAt time.Time `json:"updated_at"` } // GetSessionGraph returns all sessions as a flat list of graph nodes. + func (al *AgentLoop) GetSessionGraph() []SessionGraphNode { agent := al.registry.GetDefaultAgent() + store := agent.Sessions.Store() + sessions, err := store.List(nil) if err != nil { return nil } + nodes := make([]SessionGraphNode, 0, len(sessions)) + for _, s := range sessions { nodes = append(nodes, SessionGraphNode{ - Key: s.Key, - Label: s.Label, - Status: s.Status, - Summary: s.Summary, - ParentKey: s.ParentKey, + Key: s.Key, + + Label: s.Label, + + Status: s.Status, + + Summary: s.Summary, + + ParentKey: s.ParentKey, + ForkTurnID: s.ForkTurnID, - TurnCount: s.TurnCount, - CreatedAt: s.CreatedAt, - UpdatedAt: s.UpdatedAt, + + TurnCount: s.TurnCount, + + CreatedAt: s.CreatedAt, + + UpdatedAt: s.UpdatedAt, }) } + return nodes } // expandSkillCommand detects "/skill <name> [message]" and returns: + // - expanded: full content with SKILL.md injected (for LLM) + // - compact: skill name tag + user message only (for history) + // - ok: whether expansion happened + func (al *AgentLoop) expandSkillCommand(msg bus.InboundMessage) (expanded string, compact string, ok bool) { content := strings.TrimSpace(msg.Content) + if !strings.HasPrefix(content, "/skill ") { return "", "", false } // Parse: /skill <name> [message] + rest := strings.TrimSpace(content[7:]) // len("/skill ") == 7 + parts := strings.SplitN(rest, " ", 2) + if len(parts) == 0 || parts[0] == "" { return "", "", false } skillName := parts[0] + userMessage := "" + if len(parts) > 1 { userMessage = strings.TrimSpace(parts[1]) } agent := al.registry.GetDefaultAgent() + if agent == nil { return "", "", false } skillContent, found := agent.ContextBuilder.LoadSkill(skillName) + if !found { return "", "", false } @@ -3998,17 +5809,25 @@ func (al *AgentLoop) expandSkillCommand(msg bus.InboundMessage) (expanded string tag := fmt.Sprintf("[Skill: %s]", skillName) // Build expanded message: skill instructions + user message (for LLM) + var sb strings.Builder + sb.WriteString(tag) + sb.WriteString("\n\n") + sb.WriteString(skillContent) + if userMessage != "" { sb.WriteString("\n\n---\n\n") + sb.WriteString(userMessage) } // Build compact form: skill name tag + user message only (for history) + compactForm := tag + if userMessage != "" { compactForm = tag + "\n" + userMessage } @@ -4017,173 +5836,241 @@ func (al *AgentLoop) expandSkillCommand(msg bus.InboundMessage) (expanded string } // handleSkillsCommand lists all available skills. + func (al *AgentLoop) handleSkillsCommand() string { agent := al.registry.GetDefaultAgent() + if agent == nil { return "No agent configured." } skillsList := agent.ContextBuilder.ListSkills() + if len(skillsList) == 0 { return "No skills available.\nAdd skills to your workspace/skills/ directory." } var sb strings.Builder + sb.WriteString("Available Skills\n\n") + for _, s := range skillsList { fmt.Fprintf(&sb, "**%s** (%s)\n", s.Name, s.Source) + if s.Description != "" { fmt.Fprintf(&sb, "```\n%s\n```\n", s.Description) } } + sb.WriteString("\nUse: /skill <name> [message]") + return sb.String() } // handlePlanCommand handles /plan subcommands that can be resolved instantly. + // Returns (response, handled). For "/plan <task>" (new plan), it returns + // ("", false) so the message falls through to the LLM queue, where + // expandPlanCommand writes the seed and rewrites the content. + func (al *AgentLoop) handlePlanCommand(args []string, sessionKey string) (string, bool) { agent := al.registry.GetDefaultAgent() + if agent == nil { return "No agent configured.", true } if len(args) == 0 { // /plan — show current plan + return agent.ContextBuilder.FormatPlanDisplay(), true } sub := args[0] + switch sub { case "clear": + if agent.ContextBuilder.ReadMemory() == "" { return "No active plan to clear.", true } + // Deactivate worktree on plan clear + if sessionKey != "" { agent.DeactivateWorktree(sessionKey, "", true) } + if err := agent.ContextBuilder.ClearMemory(); err != nil { return fmt.Sprintf("Error clearing plan: %v", err), true } + return "Plan cleared.", true case "done": + if !agent.ContextBuilder.HasActivePlan() { return "No active plan.", true } + if len(args) < 2 { return "Usage: /plan done <step number>", true } + stepNum, err := strconv.Atoi(args[1]) + if err != nil || stepNum < 1 { return "Step number must be a positive integer.", true } + phase := agent.ContextBuilder.GetCurrentPhase() + if err := agent.ContextBuilder.MarkStep(phase, stepNum); err != nil { return fmt.Sprintf("Error: %v", err), true } + return fmt.Sprintf("Marked step %d in phase %d as done.", stepNum, phase), true case "add": + if !agent.ContextBuilder.HasActivePlan() { return "No active plan.", true } + if len(args) < 2 { return "Usage: /plan add <step description>", true } + desc := strings.Join(args[1:], " ") + phase := agent.ContextBuilder.GetCurrentPhase() + if err := agent.ContextBuilder.AddStep(phase, desc); err != nil { return fmt.Sprintf("Error: %v", err), true } + return fmt.Sprintf("Added step to phase %d: %s", phase, desc), true case "start": + if !agent.ContextBuilder.HasActivePlan() { return "No active plan.", true } + status := agent.ContextBuilder.GetPlanStatus() + if status == "executing" { return "Plan is already executing.", true } + if status != "interviewing" && status != "review" { return fmt.Sprintf("Cannot start from status %q.", status), true } + if agent.ContextBuilder.GetTotalPhases() == 0 { return "Cannot start: no phases defined yet. Complete the interview first.", true } + if err := agent.ContextBuilder.SetPlanStatus("executing"); err != nil { return fmt.Sprintf("Error: %v", err), true } + al.reporter().ReportStateChange(sessionKey, orch.AgentStatePlanExecuting, "") + al.planStartPending = true + clearHistory := len(args) > 1 && args[1] == "clear" + al.planClearHistory = clearHistory + if clearHistory { return "Plan approved. Executing with clean history.", true } + return "Plan approved. Executing.", true case "next": + if !agent.ContextBuilder.HasActivePlan() { return "No active plan.", true } + if err := agent.ContextBuilder.AdvancePhase(); err != nil { return fmt.Sprintf("Error: %v", err), true } + phase := agent.ContextBuilder.GetCurrentPhase() + return fmt.Sprintf("Advanced to phase %d.", phase), true case "worktrees": + return al.handlePlanWorktreesCommand(agent, args[1:]), true default: + // /plan <task description> — start new plan + // Block if a plan is already active (fast-path error). + if agent.ContextBuilder.HasActivePlan() { return "A plan is already active. Use /plan clear first.", true } + // Not handled here — let the message flow to the LLM queue. + // expandPlanCommand will write the seed and rewrite the content. + return "", false } } // isPlanPreExecution returns true if the plan is in a pre-execution state + // (interviewing or review) where tool restrictions and iteration caps apply. + func (al *AgentLoop) handlePlanWorktreesCommand(agent *AgentInstance, args []string) string { repoRoot := git.FindRepoRoot(agent.Workspace) + if repoRoot == "" { return "Workspace is not a git repository." } + worktreesDir := filepath.Join(agent.Workspace, ".worktrees") sub := "list" + if len(args) > 0 { sub = strings.ToLower(strings.TrimSpace(args[0])) } switch sub { case "", "list": + items, err := git.ListManagedWorktrees(repoRoot, worktreesDir) if err != nil { return fmt.Sprintf("Error listing worktrees: %v", err) } + if len(items) == 0 { return "No active worktrees in workspace/.worktrees." } var sb strings.Builder + sb.WriteString("Active worktrees\n\n") + for _, wt := range items { status := "clean" + if wt.HasUncommitted { status = "dirty" } + last := "(no commits)" + if wt.LastCommitHash != "" { if wt.LastCommitAge != "" { last = fmt.Sprintf("%s %s (%s)", wt.LastCommitHash, wt.LastCommitSubject, wt.LastCommitAge) @@ -4191,120 +6078,173 @@ func (al *AgentLoop) handlePlanWorktreesCommand(agent *AgentInstance, args []str last = fmt.Sprintf("%s %s", wt.LastCommitHash, wt.LastCommitSubject) } } + fmt.Fprintf(&sb, "- %s\n branch: %s\n status: %s\n last: %s\n", wt.Name, wt.Branch, status, last) } + sb.WriteString("\nCommands:\n") + sb.WriteString("/plan worktrees inspect <name>\n") + sb.WriteString("/plan worktrees merge <name>\n") + sb.WriteString("/plan worktrees dispose <name> [force]") + return sb.String() case "inspect": + if len(args) < 2 { return "Usage: /plan worktrees inspect <name>" } + name := args[1] + wt, err := git.GetManagedWorktree(repoRoot, worktreesDir, name) if err != nil { if errors.Is(err, git.ErrInvalidWorktreeName) { return "Invalid worktree name." } + if errors.Is(err, git.ErrWorktreeNotFound) { return fmt.Sprintf("Worktree %q not found.", name) } + return fmt.Sprintf("Error inspecting worktree %q: %v", name, err) } + statusOut, _ := git.WorktreeStatusShort(wt.Path) + diffOut, _ := git.WorktreeDiffStat(wt.Path) + logOut, _ := git.WorktreeRecentLog(wt.Path, 10) + if statusOut == "" { statusOut = "(clean)" } var sb strings.Builder + fmt.Fprintf(&sb, "Worktree: %s\nBranch: %s\nDirty: %t\n", wt.Name, wt.Branch, wt.HasUncommitted) + if wt.LastCommitHash != "" { fmt.Fprintf(&sb, "Last commit: %s %s", wt.LastCommitHash, wt.LastCommitSubject) + if wt.LastCommitAge != "" { fmt.Fprintf(&sb, " (%s)", wt.LastCommitAge) } + sb.WriteString("\n") } + sb.WriteString("\nStatus:\n```\n") + sb.WriteString(statusOut) + sb.WriteString("\n```\n") + if diffOut != "" { sb.WriteString("\nDiff (stat):\n```\n") + sb.WriteString(diffOut) + sb.WriteString("\n```\n") } + if logOut != "" { sb.WriteString("\nRecent commits:\n```\n") + sb.WriteString(logOut) + sb.WriteString("\n```") } + return sb.String() case "merge": + if len(args) < 2 { return "Usage: /plan worktrees merge <name>" } + name := args[1] + res, base, err := git.MergeManagedWorktree(repoRoot, worktreesDir, name, "") if err != nil { if errors.Is(err, git.ErrInvalidWorktreeName) { return "Invalid worktree name." } + if errors.Is(err, git.ErrWorktreeNotFound) { return fmt.Sprintf("Worktree %q not found.", name) } + return fmt.Sprintf("Error merging worktree %q: %v", name, err) } + if res.Conflict { return fmt.Sprintf("Merge conflict while merging `%s` into `%s`. Merge was aborted.", res.Branch, base) } + if res.Merged { return fmt.Sprintf("Merged `%s` into `%s`.", res.Branch, base) } + return fmt.Sprintf("No merge was performed for `%s`.", name) case "dispose": + if len(args) < 2 { return "Usage: /plan worktrees dispose <name> [force]" } + name := args[1] + force := len(args) > 2 && strings.EqualFold(args[2], "force") + wt, err := git.GetManagedWorktree(repoRoot, worktreesDir, name) if err != nil { if errors.Is(err, git.ErrInvalidWorktreeName) { return "Invalid worktree name." } + if errors.Is(err, git.ErrWorktreeNotFound) { return fmt.Sprintf("Worktree %q not found.", name) } + return fmt.Sprintf("Error disposing worktree %q: %v", name, err) } + if wt.HasUncommitted && !force { return fmt.Sprintf( + "Worktree `%s` has uncommitted changes. Re-run with `/plan worktrees dispose %s force` to confirm.", + name, + name, ) } + res, err := git.DisposeManagedWorktree(repoRoot, worktreesDir, name, "") if err != nil { return fmt.Sprintf("Error disposing worktree %q: %v", name, err) } + parts := []string{fmt.Sprintf("Disposed worktree `%s` (branch `%s`).", name, res.Branch)} + if res.AutoCommitted { parts = append(parts, "Uncommitted changes were auto-committed.") } + if res.CommitsAhead > 0 { parts = append(parts, fmt.Sprintf("Branch has %d unique commit(s); branch was kept.", res.CommitsAhead)) } + if res.BranchDeleted { parts = append(parts, "Branch was deleted (no unique commits).") } + return strings.Join(parts, " ") } @@ -4316,66 +6256,100 @@ func isPlanPreExecution(status string) bool { } // interviewAllowedTools is the single source of truth for tool names that may + // be sent to the LLM (and subsequently invoked) during the interview phase. + // filterInterviewTools uses this to strip tool *definitions* before the LLM call, + // while isToolAllowedDuringInterview adds argument-level checks as a second gate. + var interviewAllowedTools = map[string]bool{ - "readfile": true, - "listdir": true, - "websearch": true, - "webfetch": true, - "message": true, - "editfile": true, + "readfile": true, + + "listdir": true, + + "websearch": true, + + "webfetch": true, + + "message": true, + + "editfile": true, + "appendfile": true, - "writefile": true, - "exec": true, - "logs": true, + + "writefile": true, + + "exec": true, + + "logs": true, } // filterInterviewTools removes tool definitions that are not in the + // interviewAllowedTools whitelist, reducing token usage and preventing the + // LLM from attempting disallowed tool calls during the interview phase. + func filterInterviewTools(defs []providers.ToolDefinition) []providers.ToolDefinition { filtered := make([]providers.ToolDefinition, 0, len(defs)) + for _, d := range defs { if interviewAllowedTools[tools.NormalizeToolName(d.Function.Name)] { filtered = append(filtered, d) } } + return filtered } // isToolAllowedDuringInterview checks whether a tool call is permitted while the + // plan is in a pre-execution state. Uses the shared interviewAllowedTools map for + // name-level gating, then applies argument-level constraints for write-type tools + // (MEMORY.md only) and exec (read-only commands only). + func isToolAllowedDuringInterview(toolName string, args map[string]any) bool { norm := tools.NormalizeToolName(toolName) + if !interviewAllowedTools[norm] { return false } // Argument-level constraints + switch norm { case "editfile", "appendfile", "writefile": + path, _ := args["path"].(string) + return strings.HasSuffix(path, "MEMORY.md") + case "exec": + cmd, _ := args["command"].(string) + return isReadOnlyCommand(cmd) } + return true } // isReadOnlyCommand returns true when cmd is a safe, read-only shell command + // that an LLM may run during the interview phase. + func isReadOnlyCommand(cmd string) bool { cmd = strings.TrimSpace(cmd) + if cmd == "" { return false } // Reject write operators anywhere in the command + for _, op := range []string{">", ">>", "| tee "} { if strings.Contains(cmd, op) { return false @@ -4383,10 +6357,13 @@ func isReadOnlyCommand(cmd string) bool { } // Reject path traversal (defense in depth; ExecTool.guardCommand also enforces workspace restriction) + if strings.Contains(cmd, "..") { return false } + // Block absolute paths in arguments (allow "cd /path && cmd" which is stripped later) + for _, field := range strings.Fields(cmd) { if strings.HasPrefix(field, "/") && !strings.HasPrefix(cmd, "cd ") { return false @@ -4394,6 +6371,7 @@ func isReadOnlyCommand(cmd string) bool { } // Strip "cd /path &&" prefix (LLM habit) + if strings.HasPrefix(cmd, "cd ") { if idx := strings.Index(cmd, "&&"); idx >= 0 { cmd = strings.TrimSpace(cmd[idx+2:]) @@ -4401,86 +6379,123 @@ func isReadOnlyCommand(cmd string) bool { } fields := strings.Fields(cmd) + if len(fields) == 0 { return false } + first := filepath.Base(fields[0]) + switch first { case "find", "ls", "cat", "head", "tail", "grep", "rg", + "tree", "wc", "file", "which", "pwd", + "uname", "df", "du", "stat", "realpath", "dirname", + "basename", "date": + return true } + return false } // isWriteTool returns true if the tool can modify files. + func isWriteTool(name string) bool { switch tools.NormalizeToolName(name) { case "writefile", "editfile", "appendfile", "exec": + return true } + return false } // expandPlanCommand detects "/plan <task>" (new plan start) and: + // - writes the interview seed to MEMORY.md + // - rewrites the message content for the LLM + // - returns a compact form for session history + // + // This follows the same pattern as expandSkillCommand: the message is + // rewritten before reaching the LLM, so the AI sees the task description + // while the system prompt contains the interview guide. + func (al *AgentLoop) expandPlanCommand(msg bus.InboundMessage) (expanded string, compact string, ok bool) { content := strings.TrimSpace(msg.Content) + if !strings.HasPrefix(content, "/plan ") { return "", "", false } task := strings.TrimSpace(content[6:]) // len("/plan ") == 6 + if task == "" { return "", "", false } // Known subcommands are handled by handlePlanCommand (fast path). + firstWord := strings.Fields(task)[0] + switch firstWord { case "clear", "done", "add", "start", "next", "worktrees": + return "", "", false } agent := al.registry.GetDefaultAgent() + if agent == nil { return "", "", false } // If a plan is already active, don't expand — handleCommand will + // catch it and return the error on the fast path. + if agent.ContextBuilder.HasActivePlan() { return "", "", false } // Write the interview seed + seed := BuildInterviewSeed(task, agent.Workspace) + if err := agent.ContextBuilder.WriteMemory(seed); err != nil { return "", "", false } + al.notifyStateChange() // Expanded: the task description goes to LLM. + // The system prompt already contains the interview guide. + expanded = task + compact = fmt.Sprintf("[Plan: %s]", utils.Truncate(task, 80)) + return expanded, compact, true } // extractPeer extracts the routing peer from the inbound message's structured Peer field. + func extractPeer(msg bus.InboundMessage) *routing.RoutePeer { if msg.Peer.Kind == "" { return nil } + peerID := msg.Peer.ID + if peerID == "" { if msg.Peer.Kind == "direct" { peerID = msg.SenderID @@ -4488,15 +6503,20 @@ func extractPeer(msg bus.InboundMessage) *routing.RoutePeer { peerID = msg.ChatID } } + return &routing.RoutePeer{Kind: msg.Peer.Kind, ID: peerID} } // extractParentPeer extracts the parent peer (reply-to) from inbound message metadata. + func extractParentPeer(msg bus.InboundMessage) *routing.RoutePeer { parentKind := msg.Metadata["parent_peer_kind"] + parentID := msg.Metadata["parent_peer_id"] + if parentKind == "" || parentID == "" { return nil } + return &routing.RoutePeer{Kind: parentKind, ID: parentID} } diff --git a/pkg/agent/session_recorder.go b/pkg/agent/session_recorder.go index b654911e3..f17bb81fe 100644 --- a/pkg/agent/session_recorder.go +++ b/pkg/agent/session_recorder.go @@ -7,6 +7,7 @@ import ( ) // sessionRecorderImpl bridges tools.SessionRecorder → session.SessionStore. + type sessionRecorderImpl struct { adapter *session.LegacyAdapter } @@ -19,40 +20,102 @@ func newSessionRecorder(adapter *session.LegacyAdapter) *sessionRecorderImpl { func (r *sessionRecorderImpl) RecordFork(conductorKey, subagentKey, taskID, label string) error { store := r.adapter.Store() + return store.Fork(conductorKey, subagentKey, &session.CreateOpts{ ForkTurnID: taskID, - Label: label, + + Label: label, }) } func (r *sessionRecorderImpl) RecordSubagentTurn(subagentKey string, messages []providers.Message) error { store := r.adapter.Store() + turn := &session.Turn{ - Kind: session.TurnNormal, + Kind: session.TurnNormal, + Messages: messages, } + return store.Append(subagentKey, turn) } func (r *sessionRecorderImpl) RecordCompletion(subagentKey, status, result string) error { store := r.adapter.Store() + return store.SetStatus(subagentKey, status) } func (r *sessionRecorderImpl) RecordReport(conductorKey, subagentKey, senderID, content string) error { store := r.adapter.Store() + turn := &session.Turn{ - Kind: session.TurnReport, + Kind: session.TurnReport, + OriginKey: subagentKey, - Author: senderID, + + Author: senderID, + Messages: []providers.Message{ {Role: "user", Content: content}, }, } + if err := store.Append(conductorKey, turn); err != nil { return err } + // Advance LegacyAdapter's stored counter so flush loop doesn't double-write. + r.adapter.AdvanceStored(conductorKey, 1) + + return nil +} + +func (r *sessionRecorderImpl) RecordQuestion(conductorKey, subagentKey, taskID, question string) error { + store := r.adapter.Store() + + turn := &session.Turn{ + Kind: session.TurnQuestion, + + OriginKey: subagentKey, + + Author: taskID, + + Messages: []providers.Message{ + {Role: "user", Content: question}, + }, + } + + if err := store.Append(conductorKey, turn); err != nil { + return err + } + + r.adapter.AdvanceStored(conductorKey, 1) + + return nil +} + +func (r *sessionRecorderImpl) RecordPlanSubmit(conductorKey, subagentKey, taskID, planText string) error { + store := r.adapter.Store() + + turn := &session.Turn{ + Kind: session.TurnPlanSubmit, + + OriginKey: subagentKey, + + Author: taskID, + + Messages: []providers.Message{ + {Role: "user", Content: planText}, + }, + } + + if err := store.Append(conductorKey, turn); err != nil { + return err + } + + r.adapter.AdvanceStored(conductorKey, 1) + return nil } diff --git a/pkg/agent/session_recorder_test.go b/pkg/agent/session_recorder_test.go index 410eb59dc..083d80b8e 100644 --- a/pkg/agent/session_recorder_test.go +++ b/pkg/agent/session_recorder_test.go @@ -11,15 +11,22 @@ import ( func newTestRecorder(t *testing.T) (*sessionRecorderImpl, *session.LegacyAdapter, session.SessionStore) { t.Helper() + tmpDir := t.TempDir() + dbPath := filepath.Join(tmpDir, "test.db") + store, err := session.OpenSQLiteStore(dbPath) if err != nil { t.Fatalf("open store: %v", err) } + adapter := session.NewLegacyAdapter(store) + t.Cleanup(func() { adapter.Close() }) + recorder := newSessionRecorder(adapter) + return recorder, adapter, store } @@ -27,6 +34,7 @@ func TestRecordFork(t *testing.T) { rec, _, store := newTestRecorder(t) // Create conductor session first. + if err := store.Create("conductor:main", nil); err != nil { t.Fatalf("create conductor session: %v", err) } @@ -37,31 +45,39 @@ func TestRecordFork(t *testing.T) { } // Verify child session exists with correct parent. + info, err := store.Get("subagent:subagent-1") if err != nil { t.Fatalf("Get child: %v", err) } + if info == nil { t.Fatal("child session not found") } + if info.ParentKey != "conductor:main" { t.Errorf("ParentKey = %q, want %q", info.ParentKey, "conductor:main") } + if info.ForkTurnID != "subagent-1" { t.Errorf("ForkTurnID = %q, want %q", info.ForkTurnID, "subagent-1") } + if info.Label != "scout" { t.Errorf("Label = %q, want %q", info.Label, "scout") } // Verify parent lists child. + children, err := store.Children("conductor:main") if err != nil { t.Fatalf("Children: %v", err) } + if len(children) != 1 { t.Fatalf("children count = %d, want 1", len(children)) } + if children[0].Key != "subagent:subagent-1" { t.Errorf("child key = %q, want %q", children[0].Key, "subagent:subagent-1") } @@ -71,15 +87,19 @@ func TestRecordSubagentTurn(t *testing.T) { rec, _, store := newTestRecorder(t) // Create subagent session. + if err := store.Create("subagent:subagent-1", nil); err != nil { t.Fatalf("create: %v", err) } msgs := []providers.Message{ {Role: "system", Content: "You are a scout."}, + {Role: "user", Content: "Investigate X."}, + {Role: "assistant", Content: "Found Y."}, } + if err := rec.RecordSubagentTurn("subagent:subagent-1", msgs); err != nil { t.Fatalf("RecordSubagentTurn: %v", err) } @@ -88,15 +108,19 @@ func TestRecordSubagentTurn(t *testing.T) { if err != nil { t.Fatalf("Turns: %v", err) } + if len(turns) != 1 { t.Fatalf("turns count = %d, want 1", len(turns)) } + if turns[0].Kind != session.TurnNormal { t.Errorf("Kind = %d, want TurnNormal", turns[0].Kind) } + if len(turns[0].Messages) != 3 { t.Errorf("messages count = %d, want 3", len(turns[0].Messages)) } + if turns[0].Messages[2].Content != "Found Y." { t.Errorf("last message = %q, want %q", turns[0].Messages[2].Content, "Found Y.") } @@ -117,18 +141,23 @@ func TestRecordCompletion(t *testing.T) { if err != nil { t.Fatalf("Get: %v", err) } + if info.Status != "completed" { t.Errorf("Status = %q, want %q", info.Status, "completed") } // Test failed status. + if err := store.Create("subagent:subagent-2", nil); err != nil { t.Fatalf("create: %v", err) } + if err := rec.RecordCompletion("subagent:subagent-2", "failed", "error"); err != nil { t.Fatalf("RecordCompletion failed: %v", err) } + info2, _ := store.Get("subagent:subagent-2") + if info2.Status != "failed" { t.Errorf("Status = %q, want %q", info2.Status, "failed") } @@ -138,6 +167,7 @@ func TestRecordReport(t *testing.T) { rec, adapter, store := newTestRecorder(t) // Create conductor session via adapter so it's in cache. + _ = adapter.GetOrCreate("conductor:main") if err := store.Create("subagent:subagent-1", nil); err != nil { @@ -145,27 +175,34 @@ func TestRecordReport(t *testing.T) { } content := "[System: subagent:subagent-1] Task 'scout' completed.\n\nResult:\nFound Y." + if err := rec.RecordReport("conductor:main", "subagent:subagent-1", "subagent:subagent-1", content); err != nil { t.Fatalf("RecordReport: %v", err) } // Verify TurnReport in store. + turns, err := store.Turns("conductor:main", 0) if err != nil { t.Fatalf("Turns: %v", err) } + if len(turns) != 1 { t.Fatalf("turns count = %d, want 1", len(turns)) } + if turns[0].Kind != session.TurnReport { t.Errorf("Kind = %d, want TurnReport(%d)", turns[0].Kind, session.TurnReport) } + if turns[0].OriginKey != "subagent:subagent-1" { t.Errorf("OriginKey = %q, want %q", turns[0].OriginKey, "subagent:subagent-1") } + if turns[0].Author != "subagent:subagent-1" { t.Errorf("Author = %q, want %q", turns[0].Author, "subagent:subagent-1") } + if len(turns[0].Messages) != 1 || turns[0].Messages[0].Role != "user" { t.Errorf("unexpected messages: %v", turns[0].Messages) } @@ -175,6 +212,7 @@ func TestAdvanceStoredPreventsDoubleWrite(t *testing.T) { rec, adapter, store := newTestRecorder(t) // Create conductor session via adapter. + _ = adapter.GetOrCreate("conductor:main") if err := store.Create("subagent:subagent-1", nil); err != nil { @@ -182,57 +220,157 @@ func TestAdvanceStoredPreventsDoubleWrite(t *testing.T) { } // Simulate: conductor has 2 messages already flushed. + adapter.AddMessage("conductor:main", "user", "hello") + adapter.AddMessage("conductor:main", "assistant", "hi") + if err := adapter.Save("conductor:main"); err != nil { t.Fatalf("Save: %v", err) } // RecordReport writes directly to store and advances stored counter. + content := "[System: subagent:subagent-1] result" + if err := rec.RecordReport("conductor:main", "subagent:subagent-1", "subagent:subagent-1", content); err != nil { t.Fatalf("RecordReport: %v", err) } // The in-memory cache should also be updated (by loop.go calling AddFullMessage). + // Simulate what loop.go does after RecordReport succeeds. + adapter.AddFullMessage("conductor:main", providers.Message{Role: "user", Content: content}) + // AdvanceStored was already called by RecordReport, so stored = 3 + 1 = 4 + // but we added 1 message to cache making it len=4 as well. No double write. // Save should NOT re-write the report turn. + if err := adapter.Save("conductor:main"); err != nil { t.Fatalf("Save after report: %v", err) } // Count all turns in store for conductor session. + turns, err := store.Turns("conductor:main", 0) if err != nil { t.Fatalf("Turns: %v", err) } // Expected: turn 1 (initial 2 msgs), turn 2 (TurnReport from RecordReport) + // NOT turn 3 (duplicate from flush). + if len(turns) != 2 { t.Errorf("turns count = %d, want 2 (no double-write)", len(turns)) + for i, turn := range turns { t.Logf(" turn[%d]: seq=%d kind=%d msgs=%d", i, turn.Seq, turn.Kind, len(turn.Messages)) } } } +func TestRecordQuestion(t *testing.T) { + rec, adapter, store := newTestRecorder(t) + + // Create conductor session via adapter so it's in cache. + + _ = adapter.GetOrCreate("conductor:main") + + if err := store.Create("subagent:subagent-1", nil); err != nil { + t.Fatalf("create subagent: %v", err) + } + + question := "What database schema should I use for the users table?" + + if err := rec.RecordQuestion("conductor:main", "subagent:subagent-1", "subagent-1", question); err != nil { + t.Fatalf("RecordQuestion: %v", err) + } + + turns, err := store.Turns("conductor:main", 0) + if err != nil { + t.Fatalf("Turns: %v", err) + } + + if len(turns) != 1 { + t.Fatalf("turns count = %d, want 1", len(turns)) + } + + if turns[0].Kind != session.TurnQuestion { + t.Errorf("Kind = %d, want TurnQuestion(%d)", turns[0].Kind, session.TurnQuestion) + } + + if turns[0].OriginKey != "subagent:subagent-1" { + t.Errorf("OriginKey = %q, want %q", turns[0].OriginKey, "subagent:subagent-1") + } + + if turns[0].Author != "subagent-1" { + t.Errorf("Author = %q, want %q", turns[0].Author, "subagent-1") + } + + if len(turns[0].Messages) != 1 || turns[0].Messages[0].Content != question { + t.Errorf("unexpected messages: %v", turns[0].Messages) + } +} + +func TestRecordPlanSubmit(t *testing.T) { + rec, adapter, store := newTestRecorder(t) + + _ = adapter.GetOrCreate("conductor:main") + + if err := store.Create("subagent:subagent-1", nil); err != nil { + t.Fatalf("create subagent: %v", err) + } + + planText := "Goal: Implement auth\nSteps:\n1. Add middleware\n2. Add JWT validation" + + if err := rec.RecordPlanSubmit("conductor:main", "subagent:subagent-1", "subagent-1", planText); err != nil { + t.Fatalf("RecordPlanSubmit: %v", err) + } + + turns, err := store.Turns("conductor:main", 0) + if err != nil { + t.Fatalf("Turns: %v", err) + } + + if len(turns) != 1 { + t.Fatalf("turns count = %d, want 1", len(turns)) + } + + if turns[0].Kind != session.TurnPlanSubmit { + t.Errorf("Kind = %d, want TurnPlanSubmit(%d)", turns[0].Kind, session.TurnPlanSubmit) + } + + if turns[0].OriginKey != "subagent:subagent-1" { + t.Errorf("OriginKey = %q, want %q", turns[0].OriginKey, "subagent:subagent-1") + } + + if len(turns[0].Messages) != 1 || turns[0].Messages[0].Content != planText { + t.Errorf("unexpected messages: %v", turns[0].Messages) + } +} + func TestExtractTaskID(t *testing.T) { tests := []struct { input string - want string + + want string }{ {"subagent:subagent-1", "subagent-1"}, + {"subagent:subagent-42", "subagent-42"}, + {"plain-id", "plain-id"}, + {"a:b:c", "c"}, } + for _, tt := range tests { got := extractTaskID(tt.input) + if got != tt.want { t.Errorf("extractTaskID(%q) = %q, want %q", tt.input, got, tt.want) } @@ -241,5 +379,6 @@ func TestExtractTaskID(t *testing.T) { func init() { // Suppress log output in tests. + os.Setenv("PICOCLAW_LOG_LEVEL", "error") } diff --git a/pkg/session/types.go b/pkg/session/types.go index 32ecdcf39..395ccfa8c 100644 --- a/pkg/session/types.go +++ b/pkg/session/types.go @@ -19,6 +19,15 @@ const ( ) +// Escalation turn kinds — explicit values to keep stable across versions. + +const ( + TurnQuestion TurnKind = 10 // Subagent → conductor question (escalation) + + TurnPlanSubmit TurnKind = 11 // Subagent plan submission for review + +) + // Turn represents a single conversation turn persisted in the store. type Turn struct { diff --git a/pkg/tools/answer_subagent.go b/pkg/tools/answer_subagent.go new file mode 100644 index 000000000..86a84dfbc --- /dev/null +++ b/pkg/tools/answer_subagent.go @@ -0,0 +1,138 @@ +package tools + +import ( + "context" + "fmt" +) + +// AnswerSubagentTool allows the conductor to answer a subagent's question. + +type AnswerSubagentTool struct { + manager *SubagentManager +} + +func NewAnswerSubagentTool(manager *SubagentManager) *AnswerSubagentTool { + return &AnswerSubagentTool{manager: manager} +} + +func (t *AnswerSubagentTool) Name() string { return "answer_subagent" } + +func (t *AnswerSubagentTool) Description() string { + return "Answer a subagent's question or escalation. The subagent is blocked waiting for your response." +} + +func (t *AnswerSubagentTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + + "properties": map[string]any{ + "task_id": map[string]any{ + "type": "string", + + "description": "The task ID of the subagent (e.g. subagent-1)", + }, + + "answer": map[string]any{ + "type": "string", + + "description": "Your answer to the subagent's question", + }, + }, + + "required": []string{"task_id", "answer"}, + } +} + +func (t *AnswerSubagentTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + taskID, _ := args["task_id"].(string) + + if taskID == "" { + return ErrorResult("required parameter \"task_id\" (string) is missing") + } + + answer, _ := args["answer"].(string) + + if answer == "" { + return ErrorResult("required parameter \"answer\" (string) is missing") + } + + if t.manager == nil { + return ErrorResult("subagent manager not available") + } + + if err := t.manager.AnswerQuestion(taskID, answer); err != nil { + return ErrorResult(fmt.Sprintf("failed to answer subagent: %v", err)) + } + + return &ToolResult{ + ForLLM: fmt.Sprintf("Answer sent to %s.", taskID), + + ForUser: fmt.Sprintf("Answered %s", taskID), + } +} + +// ReviewSubagentPlanTool allows the conductor to approve/reject a subagent's plan. + +type ReviewSubagentPlanTool struct { + manager *SubagentManager +} + +func NewReviewSubagentPlanTool(manager *SubagentManager) *ReviewSubagentPlanTool { + return &ReviewSubagentPlanTool{manager: manager} +} + +func (t *ReviewSubagentPlanTool) Name() string { return "review_subagent_plan" } + +func (t *ReviewSubagentPlanTool) Description() string { + return "Approve or reject a subagent's execution plan. Use decision 'approved' to approve, or provide rejection feedback." +} + +func (t *ReviewSubagentPlanTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + + "properties": map[string]any{ + "task_id": map[string]any{ + "type": "string", + + "description": "The task ID of the subagent (e.g. subagent-1)", + }, + + "decision": map[string]any{ + "type": "string", + + "description": "Decision: 'approved' to approve, or rejection feedback text", + }, + }, + + "required": []string{"task_id", "decision"}, + } +} + +func (t *ReviewSubagentPlanTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + taskID, _ := args["task_id"].(string) + + if taskID == "" { + return ErrorResult("required parameter \"task_id\" (string) is missing") + } + + decision, _ := args["decision"].(string) + + if decision == "" { + return ErrorResult("required parameter \"decision\" (string) is missing") + } + + if t.manager == nil { + return ErrorResult("subagent manager not available") + } + + if err := t.manager.AnswerQuestion(taskID, decision); err != nil { + return ErrorResult(fmt.Sprintf("failed to send review decision: %v", err)) + } + + return &ToolResult{ + ForLLM: fmt.Sprintf("Review decision '%s' sent to %s.", decision, taskID), + + ForUser: fmt.Sprintf("Reviewed %s: %s", taskID, decision), + } +} diff --git a/pkg/tools/ask_conductor.go b/pkg/tools/ask_conductor.go new file mode 100644 index 000000000..aa847c645 --- /dev/null +++ b/pkg/tools/ask_conductor.go @@ -0,0 +1,110 @@ +package tools + +import ( + "context" + "fmt" +) + +// AskConductorTool allows a subagent to ask the conductor a question. + +// The subagent blocks until the conductor answers via AnswerSubagentTool. + +type AskConductorTool struct { + taskID string + + conductorKey string + + subagentKey string + + outCh chan<- ContainerMessage + + inCh <-chan string + + recorder SessionRecorder +} + +func NewAskConductorTool( + taskID, conductorKey, subagentKey string, + + outCh chan<- ContainerMessage, + + inCh <-chan string, + + recorder SessionRecorder, +) *AskConductorTool { + return &AskConductorTool{ + taskID: taskID, + + conductorKey: conductorKey, + + subagentKey: subagentKey, + + outCh: outCh, + + inCh: inCh, + + recorder: recorder, + } +} + +func (t *AskConductorTool) Name() string { return "ask_conductor" } + +func (t *AskConductorTool) Description() string { + return "Ask the conductor a clarifying question. Blocks until the conductor responds. Use when you need guidance or a decision before proceeding." +} + +func (t *AskConductorTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + + "properties": map[string]any{ + "question": map[string]any{ + "type": "string", + + "description": "The question to ask the conductor", + }, + }, + + "required": []string{"question"}, + } +} + +func (t *AskConductorTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + question, ok := args["question"].(string) + + if !ok || question == "" { + return ErrorResult("required parameter \"question\" (string) is missing") + } + + // Fire-and-forget: record question in session DAG. + + if t.recorder != nil { + _ = t.recorder.RecordQuestion(t.conductorKey, t.subagentKey, t.taskID, question) + } + + // Send question to conductor (blocking with ctx). + + select { + case t.outCh <- ContainerMessage{Type: "question", Content: question, TaskID: t.taskID}: + + case <-ctx.Done(): + + return ErrorResult(fmt.Sprintf("context canceled while sending question: %v", ctx.Err())) + } + + // Wait for conductor's answer. + + select { + case answer := <-t.inCh: + + return &ToolResult{ + ForLLM: fmt.Sprintf("Conductor answered: %s", answer), + + ForUser: answer, + } + + case <-ctx.Done(): + + return ErrorResult(fmt.Sprintf("context canceled while waiting for answer: %v", ctx.Err())) + } +} diff --git a/pkg/tools/ask_conductor_test.go b/pkg/tools/ask_conductor_test.go new file mode 100644 index 000000000..61b326d9a --- /dev/null +++ b/pkg/tools/ask_conductor_test.go @@ -0,0 +1,77 @@ +package tools + +import ( + "context" + "testing" + "time" +) + +func TestAskConductorTool_Execute(t *testing.T) { + outCh := make(chan ContainerMessage, 4) + + inCh := make(chan string, 1) + + tool := NewAskConductorTool("subagent-1", "conductor:main", "subagent:subagent-1", outCh, inCh, nil) + + if tool.Name() != "ask_conductor" { + t.Errorf("Name() = %q, want %q", tool.Name(), "ask_conductor") + } + + // Simulate conductor answering in background. + + go func() { + msg := <-outCh + + if msg.Type != "question" { + t.Errorf("msg.Type = %q, want %q", msg.Type, "question") + } + + if msg.Content != "What port?" { + t.Errorf("msg.Content = %q, want %q", msg.Content, "What port?") + } + + inCh <- "Use port 8080" + }() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + result := tool.Execute(ctx, map[string]any{"question": "What port?"}) + + if result.IsError { + t.Fatalf("unexpected error: %s", result.ForLLM) + } + + if result.ForUser != "Use port 8080" { + t.Errorf("ForUser = %q, want %q", result.ForUser, "Use port 8080") + } +} + +func TestAskConductorTool_MissingQuestion(t *testing.T) { + tool := NewAskConductorTool("subagent-1", "conductor:main", "subagent:subagent-1", nil, nil, nil) + + result := tool.Execute(context.Background(), map[string]any{}) + + if !result.IsError { + t.Error("expected error for missing question") + } +} + +func TestAskConductorTool_ContextCanceled(t *testing.T) { + outCh := make(chan ContainerMessage) // unbuffered, will block + + inCh := make(chan string) + + tool := NewAskConductorTool("subagent-1", "conductor:main", "subagent:subagent-1", outCh, inCh, nil) + + ctx, cancel := context.WithCancel(context.Background()) + + cancel() // cancel immediately + + result := tool.Execute(ctx, map[string]any{"question": "test?"}) + + if !result.IsError { + t.Error("expected error on canceled context") + } +} diff --git a/pkg/tools/session_recorder.go b/pkg/tools/session_recorder.go index 1ac339ecb..05104b1df 100644 --- a/pkg/tools/session_recorder.go +++ b/pkg/tools/session_recorder.go @@ -16,4 +16,10 @@ type SessionRecorder interface { // RecordReport injects a TurnReport into the conductor session. RecordReport(conductorSessionKey, subagentSessionKey, senderID, content string) error + + // RecordQuestion injects a TurnQuestion into the conductor session (subagent escalation). + RecordQuestion(conductorKey, subagentKey, taskID, question string) error + + // RecordPlanSubmit injects a TurnPlanSubmit into the conductor session (plan review request). + RecordPlanSubmit(conductorKey, subagentKey, taskID, planText string) error } diff --git a/pkg/tools/spawn.go b/pkg/tools/spawn.go index e83157616..af19ca86d 100644 --- a/pkg/tools/spawn.go +++ b/pkg/tools/spawn.go @@ -7,22 +7,29 @@ import ( ) type SpawnTool struct { - manager *SubagentManager - originChannel string - originChatID string + manager *SubagentManager + + originChannel string + + originChatID string + allowlistCheck func(targetAgentID string) bool - callback AsyncCallback // For async completion notification + + callback AsyncCallback // For async completion notification } func NewSpawnTool(manager *SubagentManager) *SpawnTool { return &SpawnTool{ - manager: manager, + manager: manager, + originChannel: "cli", - originChatID: "direct", + + originChatID: "direct", } } // SetCallback implements AsyncTool interface for async completion notification + func (t *SpawnTool) SetCallback(cb AsyncCallback) { t.callback = cb } @@ -38,31 +45,42 @@ func (t *SpawnTool) Description() string { func (t *SpawnTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "task": map[string]any{ - "type": "string", + "type": "string", + "description": "The task for subagent to complete", }, + "label": map[string]any{ - "type": "string", + "type": "string", + "description": "Optional short label for the task (for display)", }, + "agent_id": map[string]any{ - "type": "string", + "type": "string", + "description": "Optional target agent ID to delegate the task to", }, + "preset": map[string]any{ - "type": "string", - "enum": []string{"scout", "analyst", "coder", "worker", "coordinator"}, + "type": "string", + + "enum": []string{"scout", "analyst", "coder", "worker", "coordinator"}, + "description": "Optional capability tier: scout (explore), analyst (analyze), coder (code), worker (build), coordinator (orchestrate)", }, }, + "required": []string{"task"}, } } func (t *SpawnTool) SetContext(channel, chatID string) { t.originChannel = channel + t.originChatID = chatID } @@ -72,20 +90,28 @@ func (t *SpawnTool) SetAllowlistChecker(check func(targetAgentID string) bool) { func (t *SpawnTool) Execute(ctx context.Context, args map[string]any) *ToolResult { task, ok := args["task"].(string) + if !ok || strings.TrimSpace(task) == "" { return ErrorResult( + `Required parameter "task" (string) is missing. ` + + `Example: {"task": "describe what you need done", "preset": "scout"}`, ) } label, _ := args["label"].(string) + agentID, _ := args["agent_id"].(string) + preset, _ := args["preset"].(string) // Check allowlist if targeting a specific agent ID. + // Presets (scout, analyst, etc.) are NOT agent IDs — they are validated + // separately by IsValidPreset() in the subagent manager. + if agentID != "" && t.allowlistCheck != nil { if !t.allowlistCheck(agentID) { return ErrorResult(fmt.Sprintf("agent %q is not in the allowed agents list", agentID)) @@ -93,9 +119,12 @@ func (t *SpawnTool) Execute(ctx context.Context, args map[string]any) *ToolResul } // Validate preset name if provided + if preset != "" && !IsValidPreset(Preset(preset)) { return ErrorResult(fmt.Sprintf( + "preset %q is not valid. Available presets: scout, analyst, coder, worker, coordinator", + preset, )) } @@ -105,11 +134,13 @@ func (t *SpawnTool) Execute(ctx context.Context, args map[string]any) *ToolResul } // Pass callback to manager for async completion notification + result, err := t.manager.Spawn(ctx, task, label, agentID, t.originChannel, t.originChatID, preset, t.callback) if err != nil { return ErrorResult(fmt.Sprintf("failed to spawn subagent: %v", err)) } // Return AsyncResult since the task runs in background + return AsyncResult(result) } diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index b5266345f..1926e2eb7 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -3,6 +3,8 @@ package tools import ( "context" "fmt" + "os" + "path/filepath" "sort" "strconv" "strings" @@ -16,243 +18,441 @@ import ( ) // spawnTimeout is the hard upper bound for a single spawn goroutine. + // MaxIterations × HTTP timeout provides the soft limit; this is a safety net. + const spawnTimeout = 30 * time.Minute +// ContainerMessage is sent from a subagent to the conductor via outCh. + +type ContainerMessage struct { + Type string // "question" or "plan_review" + + Content string + + TaskID string +} + +// isDeliberatePreset returns true for presets that use the deliberate + +// (clarifying → review → executing) workflow with escalation channels. + +func isDeliberatePreset(p Preset) bool { + switch p { + case PresetCoder, PresetWorker, PresetCoordinator: + + return true + } + + return false +} + +// SubagentPlanState represents the deliberate workflow phase. + +type SubagentPlanState int + +const ( + PlanNone SubagentPlanState = iota // Not a deliberate preset + + PlanClarifying // Gathering info, asking questions + + PlanReview // Plan submitted, awaiting approval + + PlanExecuting // Plan approved, executing + + PlanCompleted // Done + +) + +// String returns a human-readable label for the plan state. + +func (s SubagentPlanState) String() string { + switch s { + case PlanClarifying: + + return "clarifying" + + case PlanReview: + + return "review" + + case PlanExecuting: + + return "executing" + + case PlanCompleted: + + return "completed" + + default: + + return "none" + } +} + type SubagentTask struct { - ID string - Task string - Label string - AgentID string + ID string + + Task string + + Label string + + AgentID string + OriginChannel string - OriginChatID string - Status string - Result string - Created int64 - CompletedAt int64 `json:"-"` - Iterations int `json:"-"` - ToolCalls int `json:"-"` - ToolStats map[string]int `json:"-"` - cancel context.CancelFunc + + OriginChatID string + + Status string + + Result string + + Created int64 + + CompletedAt int64 `json:"-"` + + Iterations int `json:"-"` + + ToolCalls int `json:"-"` + + ToolStats map[string]int `json:"-"` + + cancel context.CancelFunc + + // Escalation channels for deliberate presets (nil for exploratory). + + inCh chan string // conductor → subagent answers + + outCh chan ContainerMessage // subagent → conductor questions/plan reviews + + // Plan mode state (deliberate presets only). + + PlanState SubagentPlanState + + PlanGoal string + + PlanSteps []string } type SubagentManager struct { - tasks map[string]*SubagentTask - mu sync.RWMutex - wg sync.WaitGroup // tracks running spawn goroutines - provider providers.LLMProvider - defaultModel string - bus *bus.MessageBus - workspace string - tools *ToolRegistry - webSearchOpts WebSearchToolOptions - maxIterations int - maxTokens int - temperature float64 - hasMaxTokens bool + tasks map[string]*SubagentTask + + mu sync.RWMutex + + wg sync.WaitGroup // tracks running spawn goroutines + + provider providers.LLMProvider + + defaultModel string + + bus *bus.MessageBus + + workspace string + + tools *ToolRegistry + + webSearchOpts WebSearchToolOptions + + maxIterations int + + maxTokens int + + temperature float64 + + hasMaxTokens bool + hasTemperature bool - nextID int - reporter orch.AgentReporter - recorder SessionRecorder + + nextID int + + reporter orch.AgentReporter + + recorder SessionRecorder + conductorSessionKey string } func NewSubagentManager( provider providers.LLMProvider, + defaultModel, workspace string, + bus *bus.MessageBus, + reporter orch.AgentReporter, + webSearchOpts WebSearchToolOptions, ) *SubagentManager { if reporter == nil { reporter = orch.Noop } + return &SubagentManager{ - tasks: make(map[string]*SubagentTask), - provider: provider, - defaultModel: defaultModel, - bus: bus, - workspace: workspace, - tools: NewToolRegistry(), + tasks: make(map[string]*SubagentTask), + + provider: provider, + + defaultModel: defaultModel, + + bus: bus, + + workspace: workspace, + + tools: NewToolRegistry(), + webSearchOpts: webSearchOpts, + maxIterations: 10, - nextID: 1, - reporter: reporter, + + nextID: 1, + + reporter: reporter, } } // SetLLMOptions sets max tokens and temperature for subagent LLM calls. + func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) { sm.mu.Lock() + defer sm.mu.Unlock() + sm.maxTokens = maxTokens + sm.hasMaxTokens = true + sm.temperature = temperature + sm.hasTemperature = true } // SetTools sets the tool registry for subagent execution. + // If not set, subagent will have access to the provided tools. + func (sm *SubagentManager) SetTools(tools *ToolRegistry) { sm.mu.Lock() + defer sm.mu.Unlock() + sm.tools = tools } // SetSessionRecorder configures session recording for DAG persistence. + func (sm *SubagentManager) SetSessionRecorder(r SessionRecorder, conductorSessionKey string) { sm.mu.Lock() + defer sm.mu.Unlock() + sm.recorder = r + sm.conductorSessionKey = conductorSessionKey } // RegisterTool registers a tool for subagent execution. + func (sm *SubagentManager) RegisterTool(tool Tool) { sm.mu.Lock() + defer sm.mu.Unlock() + sm.tools.Register(tool) } func (sm *SubagentManager) Spawn( ctx context.Context, + task, label, agentID, originChannel, originChatID, preset string, + callback AsyncCallback, ) (string, error) { sm.mu.Lock() + defer sm.mu.Unlock() taskID := fmt.Sprintf("subagent-%d", sm.nextID) + sm.nextID++ subagentTask := &SubagentTask{ - ID: taskID, - Task: task, - Label: label, - AgentID: agentID, + ID: taskID, + + Task: task, + + Label: label, + + AgentID: agentID, + OriginChannel: originChannel, - OriginChatID: originChatID, - Status: "running", - Created: time.Now().UnixMilli(), + + OriginChatID: originChatID, + + Status: "running", + + Created: time.Now().UnixMilli(), } + + // Create escalation channels for deliberate presets. + + p := Preset(preset) + + if IsValidPreset(p) && isDeliberatePreset(p) { + subagentTask.inCh = make(chan string, 1) + + subagentTask.outCh = make(chan ContainerMessage, 4) + } + sm.tasks[taskID] = subagentTask sm.reporter.ReportSpawn(taskID, label, task) // Record fork in session DAG (conductor → subagent). + if sm.recorder != nil && sm.conductorSessionKey != "" { subSessionKey := routing.BuildSubagentSessionKey(taskID) + _ = sm.recorder.RecordFork(sm.conductorSessionKey, subSessionKey, taskID, label) } // Start task in background with a detached context that has a hard timeout. + // The spawned goroutine must outlive the parent (e.g. heartbeat session) + // which may finish before the subagent completes. + // The cancel func is stored on the task so CancelTask() can stop it. + spawnCtx, spawnCancel := context.WithTimeout(context.Background(), spawnTimeout) + subagentTask.cancel = spawnCancel + sm.wg.Add(1) + go func() { defer sm.wg.Done() + sm.runTask(spawnCtx, subagentTask, preset, callback) }() if label != "" { return fmt.Sprintf("Spawned subagent '%s' for task: %s", label, task), nil } + return fmt.Sprintf("Spawned subagent for task: %s", task), nil } func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, preset string, callback AsyncCallback) { task.Status = "running" - // Build system prompt based on preset type - systemPrompt := `You are a subagent. Complete the given task independently and report the result. -You have access to tools - use them as needed to complete your task. -After completing the task, provide a clear summary of what was done.` - - // Select prompt based on preset (exploratory vs deliberate) p := Preset(preset) - if IsValidPreset(p) { - switch p { - case PresetScout, PresetAnalyst: - // Exploratory presets - systemPrompt = `You are an exploratory subagent. Investigate the task and report your findings. + + if IsValidPreset(p) && isDeliberatePreset(p) { + sm.runDeliberateTask(ctx, task, p, callback) + } else { + sm.runExploratoryTask(ctx, task, p, callback) + } +} + +// clarifyingSystemPrompt returns the system prompt for the clarifying phase. + +func clarifyingSystemPrompt() string { + return `You are a deliberate subagent in the CLARIFYING phase. + +Your job is to understand the task fully before acting. You MUST: + +1. Read relevant files and gather context using your tools. + +2. If anything is unclear, use ask_conductor to ask the conductor. + +3. When you have a clear plan, use submit_plan with a goal and steps. + + + +Do NOT execute any changes yet. Only investigate and plan. + +Available escalation tools: ask_conductor, submit_plan.` +} + +// executingSystemPrompt returns the system prompt for the executing phase. + +func executingSystemPrompt() string { + return `You are a deliberate subagent in the EXECUTING phase. Your plan was approved. + +Execute the plan steps methodically. Use all available tools to complete the work. + +After completing, provide a clear summary of what was done and how it was verified. + + + +If you encounter a blocker, use ask_conductor to escalate.` +} + +// exploratorySystemPrompt returns the system prompt for exploratory presets. + +func exploratorySystemPrompt(p Preset) string { + switch p { + case PresetScout, PresetAnalyst: + + return `You are an exploratory subagent. Investigate the task and report your findings. + Use your best judgment when encountering ambiguity. Use tools as needed. + Return clear findings and observations.` - case PresetCoder, PresetWorker, PresetCoordinator: - // Deliberate presets - systemPrompt = `You are a deliberate subagent. Complete the task methodically and verify your work. -Before executing significant actions, think through your approach. -After completing, provide a clear summary of what was done and how it was verified.` - } - } - messages := []providers.Message{ - { - Role: "system", - Content: systemPrompt, - }, - { - Role: "user", - Content: task.Task, - }, - } - - // Check if context is already canceled before starting - select { - case <-ctx.Done(): - sm.mu.Lock() - task.Status = "canceled" - task.Result = "Task canceled before execution" - sm.mu.Unlock() - return default: - } - // Run tool loop with access to tools + return `You are a subagent. Complete the given task independently and report the result. + +You have access to tools - use them as needed to complete your task. + +After completing the task, provide a clear summary of what was done.` + } +} + +// getLLMOptions returns the LLM options snapshot under read lock. + +func (sm *SubagentManager) getLLMOptions() map[string]any { sm.mu.RLock() - // Use preset registry if preset is valid, otherwise use default registry - tools := sm.tools - if IsValidPreset(p) { - tools = sm.buildPresetRegistry(p, sm.workspace) - } - maxIter := sm.maxIterations - maxTokens := sm.maxTokens - temperature := sm.temperature - hasMaxTokens := sm.hasMaxTokens - hasTemperature := sm.hasTemperature - sm.mu.RUnlock() - var llmOptions map[string]any - if hasMaxTokens || hasTemperature { - llmOptions = map[string]any{} - if hasMaxTokens { - llmOptions["max_tokens"] = maxTokens + defer sm.mu.RUnlock() + + var opts map[string]any + + if sm.hasMaxTokens || sm.hasTemperature { + opts = map[string]any{} + + if sm.hasMaxTokens { + opts["max_tokens"] = sm.maxTokens } - if hasTemperature { - llmOptions["temperature"] = temperature + + if sm.hasTemperature { + opts["temperature"] = sm.temperature } } - // Notify conductor that the subagent is starting - sm.reporter.ReportConversation("conductor", task.ID, task.Task) + return opts +} - loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - Model: sm.defaultModel, - Tools: tools, - MaxIterations: maxIter, - LLMOptions: llmOptions, - Reporter: sm.reporter, - AgentID: task.ID, - }, messages, task.OriginChannel, task.OriginChatID) +// finishTask records completion, sends bus announcement, and invokes callback. +// Must NOT hold sm.mu on entry. + +func (sm *SubagentManager) finishTask( + ctx context.Context, + task *SubagentTask, + messages []providers.Message, + loopResult *ToolLoopResult, + err error, + callback AsyncCallback, +) { 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) } @@ -260,125 +460,372 @@ After completing, provide a clear summary of what was done and how it was verifi if err != nil { task.Status = "failed" + task.Result = fmt.Sprintf("Error: %v", err) - // Check if it was canceled + gcReason := "failed" + if ctx.Err() != nil { task.Status = "canceled" + task.Result = "Task canceled during execution" + gcReason = "canceled" } + sm.reporter.ReportGC(task.ID, gcReason) - // Record failure/cancellation in session DAG. + if sm.recorder != nil { subKey := routing.BuildSubagentSessionKey(task.ID) + _ = sm.recorder.RecordCompletion(subKey, task.Status, task.Result) } + result = &ToolResult{ - ForLLM: task.Result, + ForLLM: task.Result, + ForUser: "", - Silent: false, + IsError: true, - Async: false, - Err: err, + + Err: err, } } else { task.Status = "completed" + task.Result = loopResult.Content + task.CompletedAt = time.Now().UnixMilli() + task.Iterations = loopResult.Iterations + task.ToolCalls = loopResult.ToolCalls + task.ToolStats = loopResult.ToolStats - // Notify conductor of the result + + task.PlanState = PlanCompleted + sm.reporter.ReportConversation(task.ID, "conductor", loopResult.Content) + sm.reporter.ReportGC(task.ID, "completed") - // Record subagent turn and completion in session DAG. + if sm.recorder != nil { subKey := routing.BuildSubagentSessionKey(task.ID) + _ = sm.recorder.RecordSubagentTurn(subKey, messages) + _ = sm.recorder.RecordCompletion(subKey, "completed", loopResult.Content) } + result = &ToolResult{ ForLLM: fmt.Sprintf( + "Subagent '%s' completed (iterations: %d, tool calls: %d): %s", + task.Label, + loopResult.Iterations, + loopResult.ToolCalls, + loopResult.Content, ), + ForUser: loopResult.Content, - Silent: false, - IsError: false, - Async: false, } } // 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) + metadata := map[string]string{ "duration_ms": strconv.FormatInt(task.CompletedAt-task.Created, 10), - "iterations": strconv.Itoa(task.Iterations), - "tool_calls": strconv.Itoa(task.ToolCalls), + + "iterations": strconv.Itoa(task.Iterations), + + "tool_calls": strconv.Itoa(task.ToolCalls), } + if len(task.ToolStats) > 0 { metadata["tool_stats"] = formatToolStats(task.ToolStats) } + pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + sm.bus.PublishInbound(pubCtx, bus.InboundMessage{ - Channel: "system", + 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, + + ChatID: fmt.Sprintf("%s:%s", task.OriginChannel, task.OriginChatID), + + Content: announceContent, + Metadata: metadata, }) } } +// setPlanState updates task's plan state in memory and records status in session DAG. + +func (sm *SubagentManager) setPlanState(task *SubagentTask, state SubagentPlanState) { + task.PlanState = state + + if sm.recorder != nil { + subKey := routing.BuildSubagentSessionKey(task.ID) + + _ = sm.recorder.RecordCompletion(subKey, state.String(), "") + } +} + +// runExploratoryTask runs a single-phase tool loop for exploratory presets. + +func (sm *SubagentManager) runExploratoryTask( + ctx context.Context, + task *SubagentTask, + preset Preset, + callback AsyncCallback, +) { + systemPrompt := buildSubagentSystemPrompt(exploratorySystemPrompt(preset), sm.workspace) + + messages := []providers.Message{ + {Role: "system", Content: systemPrompt}, + + {Role: "user", Content: task.Task}, + } + + select { + case <-ctx.Done(): + + sm.mu.Lock() + + task.Status = "canceled" + + task.Result = "Task canceled before execution" + + sm.mu.Unlock() + + return + + default: + } + + sm.mu.RLock() + + reg := sm.tools + + if IsValidPreset(preset) { + reg = sm.buildPresetRegistry(preset, sm.workspace, task) + } + + maxIter := sm.maxIterations + + sm.mu.RUnlock() + + sm.reporter.ReportConversation("conductor", task.ID, task.Task) + + loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ + Provider: sm.provider, + + Model: sm.defaultModel, + + Tools: reg, + + MaxIterations: maxIter, + + LLMOptions: sm.getLLMOptions(), + + Reporter: sm.reporter, + + AgentID: task.ID, + }, messages, task.OriginChannel, task.OriginChatID) + + sm.finishTask(ctx, task, messages, loopResult, err, callback) +} + +// runDeliberateTask runs the clarifying → review → executing workflow. + +func (sm *SubagentManager) runDeliberateTask( + ctx context.Context, + task *SubagentTask, + preset Preset, + callback AsyncCallback, +) { + select { + case <-ctx.Done(): + + sm.mu.Lock() + + task.Status = "canceled" + + task.Result = "Task canceled before execution" + + sm.mu.Unlock() + + return + + default: + } + + sm.mu.RLock() + + reg := sm.buildPresetRegistry(preset, sm.workspace, task) + + maxIter := sm.maxIterations + + sm.mu.RUnlock() + + sm.reporter.ReportConversation("conductor", task.ID, task.Task) + + sm.setPlanState(task, PlanClarifying) + + // Phase 1: Clarifying — subagent gathers info and submits a plan. + + clarifyMsgs := []providers.Message{ + {Role: "system", Content: buildSubagentSystemPrompt(clarifyingSystemPrompt(), sm.workspace)}, + + {Role: "user", Content: task.Task}, + } + + clarifyResult, err := RunToolLoop(ctx, ToolLoopConfig{ + Provider: sm.provider, + + Model: sm.defaultModel, + + Tools: reg, + + MaxIterations: maxIter, + + LLMOptions: sm.getLLMOptions(), + + Reporter: sm.reporter, + + AgentID: task.ID, + }, clarifyMsgs, task.OriginChannel, task.OriginChatID) + if err != nil { + sm.finishTask(ctx, task, clarifyMsgs, nil, err, callback) + + return + } + + // After clarifying, the subagent should have used submit_plan. + + // If it didn't produce a plan, treat the clarifying result as direct completion. + + if task.PlanGoal == "" { + sm.finishTask(ctx, task, clarifyMsgs, clarifyResult, nil, callback) + + return + } + + // Phase 2: Executing — plan was approved, now execute it. + + sm.setPlanState(task, PlanExecuting) + + executeMsgs := []providers.Message{ + {Role: "system", Content: buildSubagentSystemPrompt(executingSystemPrompt(), sm.workspace)}, + + {Role: "user", Content: fmt.Sprintf("Execute the approved plan:\nGoal: %s\nSteps:\n%s", + + task.PlanGoal, formatPlanSteps(task.PlanSteps))}, + } + + execResult, err := RunToolLoop(ctx, ToolLoopConfig{ + Provider: sm.provider, + + Model: sm.defaultModel, + + Tools: reg, + + MaxIterations: maxIter * 2, // Executing gets more iterations + + LLMOptions: sm.getLLMOptions(), + + Reporter: sm.reporter, + + AgentID: task.ID, + }, executeMsgs, task.OriginChannel, task.OriginChatID) + + sm.finishTask(ctx, task, executeMsgs, execResult, err, callback) +} + +// formatPlanSteps formats plan steps as a numbered list. + +func formatPlanSteps(steps []string) string { + var b strings.Builder + + for i, step := range steps { + fmt.Fprintf(&b, "%d. %s\n", i+1, step) + } + + return b.String() +} + // buildPresetRegistry constructs a ToolRegistry for the given preset with appropriate restrictions. -func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string) *ToolRegistry { + +// If task is non-nil and has escalation channels, ask_conductor and submit_plan are registered. + +func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string, task ...*SubagentTask) *ToolRegistry { registry := NewToolRegistry() + config := SandboxConfigForPreset(preset, writeRoot) readRoot := writeRoot + if readRoot == "" { readRoot = sm.workspace } // Register read_file and list_dir with restrict=true + if config.AllowedTools["read_file"] { registry.Register(NewReadFileTool(readRoot, true)) } + if config.AllowedTools["list_dir"] { registry.Register(NewListDirTool(readRoot, true)) } // Register write tools only if allowed and writeRoot is set + if config.AllowedTools["write_file"] && writeRoot != "" { registry.Register(NewWriteFileTool(writeRoot, true)) + registry.Register(NewEditFileTool(writeRoot, true)) + registry.Register(NewAppendFileTool(writeRoot, true)) } // Register exec and bg_monitor if allowed. + // Each subagent gets its own ExecTool to avoid mutating the shared instance's + // allowRules (which would leak sandbox restrictions to the conductor). + if config.AllowedTools["exec"] { execWorkDir := writeRoot + if execWorkDir == "" { execWorkDir = sm.workspace } + execTool, err := NewExecTool(execWorkDir, true) if err != nil { // exec disabled for this subagent; skip registration + return registry } + if config.ExecPolicy != nil { execTool.SetAllowRules(config.ExecPolicy.AllowRules) + execTool.SetLocalNetOnly(config.ExecPolicy.LocalNetOnly) } + registry.Register(execTool) if config.AllowedTools["bg_monitor"] { @@ -387,58 +834,171 @@ func (sm *SubagentManager) buildPresetRegistry(preset Preset, writeRoot string) } // Register git tools (worktree-safe push and PR creation) + if config.AllowedTools["git_push"] { registry.Register(NewGitPushTool()) } + if config.AllowedTools["create_pr"] { registry.Register(NewCreatePRTool()) } // Register web tools + if config.AllowedTools["web_search"] { webSearchTool, _ := NewWebSearchTool(sm.webSearchOpts) + if webSearchTool != nil { registry.Register(webSearchTool) } } + if config.AllowedTools["web_fetch"] { registry.Register(NewWebFetchTool(50000)) } // Register message tool (always available) + registry.Register(NewMessageTool()) // Register spawn tool only for coordinator preset + if config.AllowedTools["spawn"] && preset == PresetCoordinator { spawnTool := NewSpawnTool(sm) + registry.Register(spawnTool) } + // Register escalation tools for deliberate presets with channels. + + if len(task) > 0 && task[0] != nil && task[0].outCh != nil { + t := task[0] + + subKey := "subagent:" + t.ID + + registry.Register(NewAskConductorTool( + + t.ID, sm.conductorSessionKey, subKey, + + t.outCh, t.inCh, sm.recorder, + )) + + submitPlan := NewSubmitPlanTool( + + t.ID, sm.conductorSessionKey, subKey, + + t.outCh, t.inCh, sm.recorder, + ) + + submitPlan.SetPlanCallback(func(goal string, steps []string) { + t.PlanGoal = goal + + t.PlanSteps = steps + }) + + registry.Register(submitPlan) + } + return registry } +// PendingQuestions drains all outCh channels and returns pending container messages. + +// Non-blocking: reads all available messages without waiting. + +func (sm *SubagentManager) PendingQuestions() []ContainerMessage { + sm.mu.RLock() + + defer sm.mu.RUnlock() + + var msgs []ContainerMessage + + for _, task := range sm.tasks { + if task.outCh == nil { + continue + } + + for { + select { + case msg := <-task.outCh: + + msgs = append(msgs, msg) + + default: + + goto nextTask + } + } + + nextTask: + } + + return msgs +} + +// AnswerQuestion sends an answer to a subagent's inCh (non-blocking). + +func (sm *SubagentManager) AnswerQuestion(taskID, answer string) error { + sm.mu.RLock() + + task, ok := sm.tasks[taskID] + + sm.mu.RUnlock() + + if !ok { + return fmt.Errorf("task %q not found", taskID) + } + + if task.inCh == nil { + return fmt.Errorf("task %q has no escalation channel", taskID) + } + + select { + case task.inCh <- answer: + + return nil + + default: + + return fmt.Errorf("task %q answer channel full", taskID) + } +} + // WaitAll blocks until all spawned subagent goroutines have finished + // or the timeout expires. Returns true if all goroutines finished, + // false on timeout. + func (sm *SubagentManager) WaitAll(timeout time.Duration) bool { done := make(chan struct{}) + go func() { sm.wg.Wait() + close(done) }() + select { case <-done: + return true + case <-time.After(timeout): + return false } } // CancelTask cancels the context for a running subagent task. + func (sm *SubagentManager) CancelTask(taskID string) { sm.mu.RLock() + task, ok := sm.tasks[taskID] + sm.mu.RUnlock() + if ok && task.cancel != nil { task.cancel() } @@ -446,36 +1006,49 @@ func (sm *SubagentManager) CancelTask(taskID string) { func (sm *SubagentManager) GetTask(taskID string) (*SubagentTask, bool) { sm.mu.RLock() + defer sm.mu.RUnlock() + task, ok := sm.tasks[taskID] + return task, ok } func (sm *SubagentManager) ListTasks() []*SubagentTask { sm.mu.RLock() + defer sm.mu.RUnlock() tasks := make([]*SubagentTask, 0, len(sm.tasks)) + for _, task := range sm.tasks { tasks = append(tasks, task) } + return tasks } // 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. + type SubagentTool struct { - manager *SubagentManager + manager *SubagentManager + originChannel string - originChatID string + + originChatID string } func NewSubagentTool(manager *SubagentManager) *SubagentTool { return &SubagentTool{ - manager: manager, + manager: manager, + originChannel: "cli", - originChatID: "direct", + + originChatID: "direct", } } @@ -490,30 +1063,39 @@ func (t *SubagentTool) Description() string { func (t *SubagentTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "task": map[string]any{ - "type": "string", + "type": "string", + "description": "The task for subagent to complete", }, + "label": map[string]any{ - "type": "string", + "type": "string", + "description": "Optional short label for the task (for display)", }, }, + "required": []string{"task"}, } } func (t *SubagentTool) SetContext(channel, chatID string) { t.originChannel = channel + t.originChatID = chatID } func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolResult { task, ok := args["task"].(string) + if !ok { return ErrorResult( + `Required parameter "task" (string) is missing. ` + + `Example: {"task": "describe what you need done"}`, ).WithError(fmt.Errorf("task parameter is required")) } @@ -526,85 +1108,218 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe } // Build messages for subagent + messages := []providers.Message{ { - Role: "system", + Role: "system", + Content: "You are a subagent. Complete the given task independently and provide a clear, concise result.", }, + { - Role: "user", + Role: "user", + Content: task, }, } // 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 + sm.mu.RUnlock() var llmOptions map[string]any + if hasMaxTokens || hasTemperature { llmOptions = map[string]any{} + if hasMaxTokens { llmOptions["max_tokens"] = maxTokens } + if hasTemperature { llmOptions["temperature"] = temperature } } loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - Model: sm.defaultModel, - Tools: tools, + Provider: sm.provider, + + Model: sm.defaultModel, + + Tools: tools, + MaxIterations: maxIter, - LLMOptions: llmOptions, + + LLMOptions: llmOptions, }, messages, t.originChannel, t.originChatID) if err != nil { return ErrorResult(fmt.Sprintf("Subagent execution failed: %v", err)).WithError(err) } // ForUser: Brief summary for user (truncated if too long) + userContent := loopResult.Content + maxUserLen := 500 + if len(userContent) > maxUserLen { userContent = userContent[:maxUserLen] + "..." } // ForLLM: Full execution details + labelStr := label + if labelStr == "" { labelStr = "(unnamed)" } + llmContent := fmt.Sprintf("Subagent task completed:\nLabel: %s\nIterations: %d\nTool calls: %d\nResult: %s", + labelStr, loopResult.Iterations, loopResult.ToolCalls, loopResult.Content) return &ToolResult{ - ForLLM: llmContent, + ForLLM: llmContent, + ForUser: userContent, - Silent: false, + + Silent: false, + IsError: false, - Async: false, + + Async: false, } } // formatToolStats formats a tool stats map as a compact string: "exec:3,read_file:5". + // Keys are sorted alphabetically for deterministic output. + func formatToolStats(stats map[string]int) string { keys := make([]string, 0, len(stats)) + for k := range stats { keys = append(keys, k) } + sort.Strings(keys) + parts := make([]string, 0, len(keys)) + for _, k := range keys { parts = append(parts, k+":"+strconv.Itoa(stats[k])) } + return strings.Join(parts, ",") } + +// extractPlanContext reads MEMORY.md from the workspace and extracts relevant + +// sections (Task, Context, Commands) to provide as subagent environment. + +func extractPlanContext(workspace string) string { + memPath := filepath.Join(workspace, "memory", "MEMORY.md") + + data, err := os.ReadFile(memPath) + if err != nil { + return "" + } + + content := string(data) + + var sections []string + + // Extract key sections by header. + + for _, header := range []string{"## Context", "## Commands", "## Orchestration"} { + if section := extractSection(content, header); section != "" { + sections = append(sections, section) + } + } + + // Also extract the task line from the header block. + + for _, line := range strings.Split(content, "\n") { + if strings.HasPrefix(line, "> Task:") { + sections = append([]string{strings.TrimSpace(line)}, sections...) + + break + } + } + + if len(sections) == 0 { + return "" + } + + return strings.Join(sections, "\n\n") +} + +// extractSection extracts a markdown section by header (including its content + +// until the next section of the same or higher level). + +func extractSection(content, header string) string { + idx := strings.Index(content, header) + + if idx < 0 { + return "" + } + + // Determine header level. + + level := 0 + + for _, c := range header { + if c == '#' { + level++ + } else { + break + } + } + + start := idx + + rest := content[idx+len(header):] + + // Find next section at same or higher level. + + nextHeader := "\n" + strings.Repeat("#", level) + " " + + end := strings.Index(rest, nextHeader) + + if end < 0 { + return strings.TrimSpace(content[start:]) + } + + return strings.TrimSpace(content[start : start+len(header)+end]) +} + +// buildSubagentSystemPrompt builds an enriched system prompt for a subagent + +// by combining the base prompt with environment context from MEMORY.md. + +func buildSubagentSystemPrompt(basePrompt, workspace string) string { + envContext := extractPlanContext(workspace) + + if envContext == "" { + return basePrompt + } + + return basePrompt + "\n\n## Environment Context\n\n" + envContext +} diff --git a/pkg/tools/subagent_container_test.go b/pkg/tools/subagent_container_test.go new file mode 100644 index 000000000..4d9b26a32 --- /dev/null +++ b/pkg/tools/subagent_container_test.go @@ -0,0 +1,125 @@ +package tools + +import ( + "testing" +) + +func TestIsDeliberatePreset(t *testing.T) { + tests := []struct { + preset Preset + want bool + }{ + {PresetScout, false}, + {PresetAnalyst, false}, + {PresetCoder, true}, + {PresetWorker, true}, + {PresetCoordinator, true}, + {"unknown", false}, + } + for _, tt := range tests { + got := isDeliberatePreset(tt.preset) + if got != tt.want { + t.Errorf("isDeliberatePreset(%q) = %v, want %v", tt.preset, got, tt.want) + } + } +} + +func TestContainerMessageChannels(t *testing.T) { + // Simulate channel creation for a deliberate preset task. + task := &SubagentTask{ + ID: "subagent-1", + inCh: make(chan string, 1), + outCh: make(chan ContainerMessage, 4), + } + + // Subagent sends a question. + task.outCh <- ContainerMessage{ + Type: "question", + Content: "Which DB schema?", + TaskID: task.ID, + } + + // Drain pending messages. + var msgs []ContainerMessage + for { + select { + case msg := <-task.outCh: + msgs = append(msgs, msg) + default: + goto done + } + } +done: + if len(msgs) != 1 { + t.Fatalf("msgs count = %d, want 1", len(msgs)) + } + if msgs[0].Type != "question" { + t.Errorf("Type = %q, want %q", msgs[0].Type, "question") + } + if msgs[0].Content != "Which DB schema?" { + t.Errorf("Content = %q, want %q", msgs[0].Content, "Which DB schema?") + } + + // Conductor answers. + task.inCh <- "Use PostgreSQL" + answer := <-task.inCh + if answer != "Use PostgreSQL" { + t.Errorf("answer = %q, want %q", answer, "Use PostgreSQL") + } +} + +func TestPendingQuestionsAndAnswerQuestion(t *testing.T) { + mgr := &SubagentManager{ + tasks: map[string]*SubagentTask{ + "subagent-1": { + ID: "subagent-1", + outCh: make(chan ContainerMessage, 4), + inCh: make(chan string, 1), + }, + "subagent-2": { + ID: "subagent-2", + // No channels — exploratory preset. + }, + }, + } + + // Send question from subagent-1. + mgr.tasks["subagent-1"].outCh <- ContainerMessage{ + Type: "question", + Content: "What port?", + TaskID: "subagent-1", + } + + msgs := mgr.PendingQuestions() + if len(msgs) != 1 { + t.Fatalf("pending count = %d, want 1", len(msgs)) + } + if msgs[0].TaskID != "subagent-1" { + t.Errorf("TaskID = %q, want %q", msgs[0].TaskID, "subagent-1") + } + + // Second call should return empty (already drained). + msgs2 := mgr.PendingQuestions() + if len(msgs2) != 0 { + t.Errorf("second pending count = %d, want 0", len(msgs2)) + } + + // Answer the question. + if err := mgr.AnswerQuestion("subagent-1", "8080"); err != nil { + t.Fatalf("AnswerQuestion: %v", err) + } + answer := <-mgr.tasks["subagent-1"].inCh + if answer != "8080" { + t.Errorf("answer = %q, want %q", answer, "8080") + } + + // Answer non-existent task. + if err := mgr.AnswerQuestion("subagent-99", "x"); err == nil { + t.Error("expected error for non-existent task") + } + + // Answer task without channels. + if err := mgr.AnswerQuestion("subagent-2", "x"); err == nil { + t.Error("expected error for task without escalation channel") + } +} diff --git a/pkg/tools/subagent_env_test.go b/pkg/tools/subagent_env_test.go new file mode 100644 index 000000000..aea6e6b91 --- /dev/null +++ b/pkg/tools/subagent_env_test.go @@ -0,0 +1,192 @@ +package tools + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestExtractPlanContext(t *testing.T) { + tmpDir := t.TempDir() + + memDir := filepath.Join(tmpDir, "memory") + + if err := os.MkdirAll(memDir, 0o755); err != nil { + t.Fatal(err) + } + + memContent := `# Active Plan + +> Task: Implement authentication + +> Status: executing + +> Phase: 1 + + + +## Context + +The project uses JWT tokens for auth. + +Database is PostgreSQL. + + + +## Phase 1: Setup + +- [x] Add middleware + +- [ ] Add JWT validation + + + +## Commands + +build: go build ./... + +test: go test ./... + + + +## Orchestration + +### Delegated + +- auth-scout: investigate patterns + +### Findings + +- Found existing middleware in pkg/auth + +` + + if err := os.WriteFile(filepath.Join(memDir, "MEMORY.md"), []byte(memContent), 0o644); err != nil { + t.Fatal(err) + } + + ctx := extractPlanContext(tmpDir) + + if ctx == "" { + t.Fatal("extractPlanContext returned empty") + } + + // Should contain the task line. + + if !strings.Contains(ctx, "> Task: Implement authentication") { + t.Error("missing task line") + } + + // Should contain Context section. + + if !strings.Contains(ctx, "JWT tokens") { + t.Error("missing Context section content") + } + + // Should contain Commands section. + + if !strings.Contains(ctx, "go build") { + t.Error("missing Commands section content") + } + + // Should contain Orchestration section. + + if !strings.Contains(ctx, "auth-scout") { + t.Error("missing Orchestration section content") + } +} + +func TestExtractPlanContext_NoFile(t *testing.T) { + ctx := extractPlanContext(t.TempDir()) + + if ctx != "" { + t.Errorf("expected empty for missing MEMORY.md, got %q", ctx) + } +} + +func TestExtractSection(t *testing.T) { + content := `## Context + +Some context here. + + + +## Commands + +build: go build + + + +## Other + +stuff` + + section := extractSection(content, "## Context") + + if !strings.Contains(section, "Some context here.") { + t.Errorf("Context section = %q, missing content", section) + } + + if strings.Contains(section, "## Commands") { + t.Errorf("Context section leaked into next section") + } + + section = extractSection(content, "## Commands") + + if !strings.Contains(section, "go build") { + t.Errorf("Commands section = %q, missing content", section) + } + + section = extractSection(content, "## Nonexistent") + + if section != "" { + t.Errorf("expected empty for nonexistent section, got %q", section) + } +} + +func TestBuildSubagentSystemPrompt(t *testing.T) { + // With no workspace/MEMORY.md, should return base prompt unchanged. + + base := "You are a subagent." + + got := buildSubagentSystemPrompt(base, t.TempDir()) + + if got != base { + t.Errorf("expected base prompt unchanged, got %q", got) + } + + // With MEMORY.md, should append environment context. + + tmpDir := t.TempDir() + + memDir := filepath.Join(tmpDir, "memory") + + os.MkdirAll(memDir, 0o755) + + os.WriteFile(filepath.Join(memDir, "MEMORY.md"), []byte(`# Plan + +> Task: Test task + + + +## Context + +Test context info. + +`), 0o644) + + got = buildSubagentSystemPrompt(base, tmpDir) + + if !strings.Contains(got, base) { + t.Error("result should contain base prompt") + } + + if !strings.Contains(got, "Environment Context") { + t.Error("result should contain Environment Context header") + } + + if !strings.Contains(got, "Test context info") { + t.Error("result should contain MEMORY.md context") + } +} diff --git a/pkg/tools/subagent_plan_test.go b/pkg/tools/subagent_plan_test.go new file mode 100644 index 000000000..dd4cc8533 --- /dev/null +++ b/pkg/tools/subagent_plan_test.go @@ -0,0 +1,98 @@ +package tools + +import ( + "testing" +) + +func TestSubagentPlanStateString(t *testing.T) { + tests := []struct { + state SubagentPlanState + want string + }{ + {PlanNone, "none"}, + {PlanClarifying, "clarifying"}, + {PlanReview, "review"}, + {PlanExecuting, "executing"}, + {PlanCompleted, "completed"}, + } + for _, tt := range tests { + got := tt.state.String() + if got != tt.want { + t.Errorf("SubagentPlanState(%d).String() = %q, want %q", tt.state, got, tt.want) + } + } +} + +func TestFormatPlanSteps(t *testing.T) { + steps := []string{"Read config", "Add middleware", "Write tests"} + got := formatPlanSteps(steps) + want := "1. Read config\n2. Add middleware\n3. Write tests\n" + if got != want { + t.Errorf("formatPlanSteps = %q, want %q", got, want) + } +} + +func TestClarifyingSystemPrompt(t *testing.T) { + prompt := clarifyingSystemPrompt() + if prompt == "" { + t.Error("clarifyingSystemPrompt returned empty string") + } + // Should mention ask_conductor and submit_plan. + for _, keyword := range []string{"ask_conductor", "submit_plan", "CLARIFYING"} { + if !containsString(prompt, keyword) { + t.Errorf("clarifyingSystemPrompt missing keyword %q", keyword) + } + } +} + +func TestExecutingSystemPrompt(t *testing.T) { + prompt := executingSystemPrompt() + if prompt == "" { + t.Error("executingSystemPrompt returned empty string") + } + if !containsString(prompt, "EXECUTING") { + t.Error("executingSystemPrompt missing keyword EXECUTING") + } +} + +func TestExploratorySystemPrompt(t *testing.T) { + scoutPrompt := exploratorySystemPrompt(PresetScout) + if scoutPrompt == "" { + t.Error("exploratorySystemPrompt(scout) returned empty") + } + defaultPrompt := exploratorySystemPrompt("unknown") + if defaultPrompt == "" { + t.Error("exploratorySystemPrompt(unknown) returned empty") + } + if scoutPrompt == defaultPrompt { + t.Error("scout and unknown prompts should differ") + } +} + +func TestDeliberateTaskChannelsCreated(t *testing.T) { + // Verify that channels and initial state are correct for deliberate presets. + task := &SubagentTask{ + inCh: make(chan string, 1), + outCh: make(chan ContainerMessage, 4), + } + if task.inCh == nil || task.outCh == nil { + t.Fatal("expected channels to be non-nil for deliberate task") + } + if task.PlanState != PlanNone { + t.Errorf("initial PlanState = %v, want PlanNone", task.PlanState) + } +} + +// containsString checks if s contains substr. +func containsString(s, substr string) bool { + return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsSubstr(s, substr)) +} + +func containsSubstr(s, sub string) bool { + for i := 0; i <= len(s)-len(sub); i++ { + if s[i:i+len(sub)] == sub { + return true + } + } + return false +} diff --git a/pkg/tools/submit_plan.go b/pkg/tools/submit_plan.go new file mode 100644 index 000000000..ddda70cdd --- /dev/null +++ b/pkg/tools/submit_plan.go @@ -0,0 +1,186 @@ +package tools + +import ( + "context" + "encoding/json" + "fmt" + "strings" +) + +// SubmitPlanTool allows a subagent to submit a plan for conductor review. + +// The subagent blocks until the conductor approves or rejects. + +type SubmitPlanTool struct { + taskID string + + conductorKey string + + subagentKey string + + outCh chan<- ContainerMessage + + inCh <-chan string + + recorder SessionRecorder + + setPlan func(goal string, steps []string) // callback to record plan on task +} + +func NewSubmitPlanTool( + taskID, conductorKey, subagentKey string, + + outCh chan<- ContainerMessage, + + inCh <-chan string, + + recorder SessionRecorder, +) *SubmitPlanTool { + return &SubmitPlanTool{ + taskID: taskID, + + conductorKey: conductorKey, + + subagentKey: subagentKey, + + outCh: outCh, + + inCh: inCh, + + recorder: recorder, + } +} + +// SetPlanCallback sets the function called when a plan is approved to record + +// the goal and steps on the parent SubagentTask. + +func (t *SubmitPlanTool) SetPlanCallback(fn func(goal string, steps []string)) { + t.setPlan = fn +} + +func (t *SubmitPlanTool) Name() string { return "submit_plan" } + +func (t *SubmitPlanTool) Description() string { + return "Submit your execution plan for conductor review. Blocks until the conductor approves or rejects. On rejection, revise and resubmit." +} + +func (t *SubmitPlanTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + + "properties": map[string]any{ + "goal": map[string]any{ + "type": "string", + + "description": "The goal of the plan", + }, + + "steps": map[string]any{ + "type": "array", + + "description": "Ordered list of steps to execute", + + "items": map[string]any{"type": "string"}, + }, + }, + + "required": []string{"goal", "steps"}, + } +} + +func (t *SubmitPlanTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + goal, _ := args["goal"].(string) + + if goal == "" { + return ErrorResult("required parameter \"goal\" (string) is missing") + } + + stepsRaw, _ := args["steps"] + + var steps []string + + switch v := stepsRaw.(type) { + case []any: + + steps = make([]string, 0, len(v)) + + for _, s := range v { + if str, ok := s.(string); ok { + steps = append(steps, str) + } + } + + case []string: + + steps = v + } + + if len(steps) == 0 { + return ErrorResult("required parameter \"steps\" (array of strings) is missing or empty") + } + + // Build plan text for recording and display. + + var b strings.Builder + + b.WriteString("Goal: ") + + b.WriteString(goal) + + b.WriteByte('\n') + + for i, step := range steps { + fmt.Fprintf(&b, "%d. %s\n", i+1, step) + } + + planText := b.String() + + // Record in session DAG. + + if t.recorder != nil { + _ = t.recorder.RecordPlanSubmit(t.conductorKey, t.subagentKey, t.taskID, planText) + } + + // Encode plan as JSON for the conductor. + + planJSON, _ := json.Marshal(map[string]any{"goal": goal, "steps": steps}) + + // Send plan_review to conductor. + + select { + case t.outCh <- ContainerMessage{Type: "plan_review", Content: string(planJSON), TaskID: t.taskID}: + + case <-ctx.Done(): + + return ErrorResult(fmt.Sprintf("context canceled while submitting plan: %v", ctx.Err())) + } + + // Wait for conductor's decision. + + select { + case decision := <-t.inCh: + + if strings.HasPrefix(decision, "approved") { + if t.setPlan != nil { + t.setPlan(goal, steps) + } + + return &ToolResult{ + ForLLM: "Plan approved by conductor. Proceed with execution.", + + ForUser: "Plan approved", + } + } + + return &ToolResult{ + ForLLM: fmt.Sprintf("Plan rejected by conductor: %s\nRevise your plan and resubmit.", decision), + + ForUser: fmt.Sprintf("Plan rejected: %s", decision), + } + + case <-ctx.Done(): + + return ErrorResult(fmt.Sprintf("context canceled while waiting for review: %v", ctx.Err())) + } +} diff --git a/pkg/tools/submit_plan_test.go b/pkg/tools/submit_plan_test.go new file mode 100644 index 000000000..0345c8278 --- /dev/null +++ b/pkg/tools/submit_plan_test.go @@ -0,0 +1,229 @@ +package tools + +import ( + "context" + "testing" + "time" +) + +func TestSubmitPlanTool_Execute_Approved(t *testing.T) { + outCh := make(chan ContainerMessage, 4) + + inCh := make(chan string, 1) + + tool := NewSubmitPlanTool("subagent-1", "conductor:main", "subagent:subagent-1", outCh, inCh, nil) + + var gotGoal string + + var gotSteps []string + + tool.SetPlanCallback(func(goal string, steps []string) { + gotGoal = goal + + gotSteps = steps + }) + + if tool.Name() != "submit_plan" { + t.Errorf("Name() = %q, want %q", tool.Name(), "submit_plan") + } + + // Simulate conductor approving in background. + + go func() { + msg := <-outCh + + if msg.Type != "plan_review" { + t.Errorf("msg.Type = %q, want %q", msg.Type, "plan_review") + } + + inCh <- "approved" + }() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + result := tool.Execute(ctx, map[string]any{ + "goal": "Add auth", + + "steps": []any{"Add middleware", "Add JWT"}, + }) + + if result.IsError { + t.Fatalf("unexpected error: %s", result.ForLLM) + } + + if gotGoal != "Add auth" { + t.Errorf("setPlan goal = %q, want %q", gotGoal, "Add auth") + } + + if len(gotSteps) != 2 { + t.Errorf("setPlan steps count = %d, want 2", len(gotSteps)) + } +} + +func TestSubmitPlanTool_Execute_Rejected(t *testing.T) { + outCh := make(chan ContainerMessage, 4) + + inCh := make(chan string, 1) + + tool := NewSubmitPlanTool("subagent-1", "conductor:main", "subagent:subagent-1", outCh, inCh, nil) + + var planSet bool + + tool.SetPlanCallback(func(goal string, steps []string) { + planSet = true + }) + + go func() { + <-outCh + + inCh <- "rejected: needs more detail on step 2" + }() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + result := tool.Execute(ctx, map[string]any{ + "goal": "Add auth", + + "steps": []any{"Add middleware"}, + }) + + if result.IsError { + t.Fatalf("unexpected error: %s", result.ForLLM) + } + + if planSet { + t.Error("setPlan should NOT be called on rejection") + } + + if result.ForLLM == "" { + t.Error("ForLLM should contain rejection message") + } +} + +func TestSubmitPlanTool_MissingParams(t *testing.T) { + tool := NewSubmitPlanTool("subagent-1", "", "", nil, nil, nil) + + result := tool.Execute(context.Background(), map[string]any{}) + + if !result.IsError { + t.Error("expected error for missing goal") + } + + result = tool.Execute(context.Background(), map[string]any{"goal": "test"}) + + if !result.IsError { + t.Error("expected error for missing steps") + } +} + +func TestSubmitPlanTool_ContextCanceled(t *testing.T) { + outCh := make(chan ContainerMessage) // unbuffered + + inCh := make(chan string) + + tool := NewSubmitPlanTool("subagent-1", "", "", outCh, inCh, nil) + + ctx, cancel := context.WithCancel(context.Background()) + + cancel() + + result := tool.Execute(ctx, map[string]any{ + "goal": "test", + + "steps": []any{"step1"}, + }) + + if !result.IsError { + t.Error("expected error on canceled context") + } +} + +func TestAnswerSubagentTool_Execute(t *testing.T) { + mgr := &SubagentManager{ + tasks: map[string]*SubagentTask{ + "subagent-1": { + ID: "subagent-1", + + inCh: make(chan string, 1), + }, + }, + } + + tool := NewAnswerSubagentTool(mgr) + + if tool.Name() != "answer_subagent" { + t.Errorf("Name() = %q, want %q", tool.Name(), "answer_subagent") + } + + result := tool.Execute(context.Background(), map[string]any{ + "task_id": "subagent-1", + + "answer": "Use port 8080", + }) + + if result.IsError { + t.Fatalf("unexpected error: %s", result.ForLLM) + } + + // Verify the answer was sent. + + answer := <-mgr.tasks["subagent-1"].inCh + + if answer != "Use port 8080" { + t.Errorf("answer = %q, want %q", answer, "Use port 8080") + } +} + +func TestAnswerSubagentTool_MissingParams(t *testing.T) { + tool := NewAnswerSubagentTool(nil) + + result := tool.Execute(context.Background(), map[string]any{}) + + if !result.IsError { + t.Error("expected error for missing task_id") + } + + result = tool.Execute(context.Background(), map[string]any{"task_id": "x"}) + + if !result.IsError { + t.Error("expected error for missing answer") + } +} + +func TestReviewSubagentPlanTool_Execute(t *testing.T) { + mgr := &SubagentManager{ + tasks: map[string]*SubagentTask{ + "subagent-1": { + ID: "subagent-1", + + inCh: make(chan string, 1), + }, + }, + } + + tool := NewReviewSubagentPlanTool(mgr) + + if tool.Name() != "review_subagent_plan" { + t.Errorf("Name() = %q, want %q", tool.Name(), "review_subagent_plan") + } + + result := tool.Execute(context.Background(), map[string]any{ + "task_id": "subagent-1", + + "decision": "approved", + }) + + if result.IsError { + t.Fatalf("unexpected error: %s", result.ForLLM) + } + + decision := <-mgr.tasks["subagent-1"].inCh + + if decision != "approved" { + t.Errorf("decision = %q, want %q", decision, "approved") + } +} diff --git a/todo/TASKS-2.md b/todo/TASKS-2.md index 9a54c9341..e20446678 100644 --- a/todo/TASKS-2.md +++ b/todo/TASKS-2.md @@ -1,4 +1,4 @@ -# TASKS-2: Subagent Orchestration (Container Model) +# TASKS-2: Subagent Orchestration (Container Model) ✅ 実装済み ## TASKS-1 反映メモ (2026-03-05) @@ -78,15 +78,15 @@ TASKS-2 の下地はかなり実装済み。以下を前提として差分のみ | **Async Callback** | `pkg/agent/loop.go` `processRequest` | spawn 完了 → MessageBus → conductor に結果注入 | | **Orchestration Nudge** | `pkg/agent/loop.go` `buildOrchReminder()` | plan 実行中に spawn/subagent 使用を促すリマインダ | -### 未実装 ❌ → TASKS-2 スコープ +### 実装完了 ✅ (2026-03-05) | # | 要素 | 概要 | |---|---|---| -| 1 | **ContainerMessage channel** | subagent→conductor の question/status/result 双方向通信 | -| 2 | **Escalation chain** | subagent question → conductor 回答 or → human escalate | -| 3 | **Deliberate Plan Mode** | coder/worker/coordinator の clarifying→review→executing 状態遷移 | -| 4 | **SubagentEnvironment injection** | MEMORY.md からの自動コンテキスト注入 | -| 5 | **MEMORY.md Orchestration section** | conductor guidance に delegated/findings/decisions 追記 | +| 1 | **ContainerMessage channel** | ✅ `ContainerMessage` + `inCh`/`outCh` on `SubagentTask` | +| 2 | **Escalation chain** | ✅ `ask_conductor` / `answer_subagent` tools + conductor question injection | +| 3 | **Deliberate Plan Mode** | ✅ `SubagentPlanState` + `runDeliberateTask()` (clarifying→review→executing) | +| 4 | **SubagentEnvironment injection** | ✅ `extractPlanContext()` + `buildSubagentSystemPrompt()` | +| 5 | **MEMORY.md Orchestration section** | ✅ `orchestrationGuidance` 拡張 (Delegated/Findings/Decisions) | --- From fa546ad7333a337138d9a44e98f3fd9c174f7896 Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Thu, 5 Mar 2026 01:37:13 +0900 Subject: [PATCH 2/2] style: fix gci import ordering across pkg/{tools,agent,session} MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Auto-fix via golangci-lint --fix to satisfy gci formatter rules (standard → default → localmodule import grouping). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --- pkg/agent/context_cache_test.go | 262 ++++- pkg/agent/context_test.go | 62 + pkg/agent/instance_test.go | 64 +- pkg/agent/loop_reporter_test.go | 51 +- pkg/agent/loop_test.go | 1630 ++++++++++++++++++++++---- pkg/agent/memory.go | 432 ++++++- pkg/agent/memory_test.go | 279 ++++- pkg/agent/mock_provider_test.go | 7 +- pkg/agent/registry.go | 56 +- pkg/agent/registry_test.go | 63 +- pkg/agent/session_tracker.go | 107 +- pkg/agent/session_tracker_test.go | 35 + pkg/session/graph.go | 54 +- pkg/session/graph_test.go | 38 +- pkg/session/legacy_adapter.go | 52 + pkg/session/legacy_adapter_test.go | 56 +- pkg/session/manager.go | 171 ++- pkg/session/manager_test.go | 57 +- pkg/session/sqlite.go | 60 + pkg/tools/bg_monitor.go | 109 +- pkg/tools/bg_monitor_test.go | 96 +- pkg/tools/createpr.go | 154 ++- pkg/tools/createpr_test.go | 99 +- pkg/tools/cron.go | 162 ++- pkg/tools/dev_preview.go | 56 +- pkg/tools/dev_preview_test.go | 158 ++- pkg/tools/edit.go | 44 +- pkg/tools/edit_test.go | 230 +++- pkg/tools/filesystem.go | 85 +- pkg/tools/filesystem_test.go | 177 ++- pkg/tools/gitpush.go | 60 +- pkg/tools/gitpush_test.go | 85 +- pkg/tools/i2c.go | 70 +- pkg/tools/i2c_linux.go | 226 +++- pkg/tools/logs.go | 35 +- pkg/tools/logs_test.go | 50 +- pkg/tools/message.go | 37 +- pkg/tools/message_test.go | 64 + pkg/tools/registry.go | 115 +- pkg/tools/registry_test.go | 206 +++- pkg/tools/result_test.go | 62 +- pkg/tools/shell.go | 446 ++++++- pkg/tools/shell_process_unix.go | 54 + pkg/tools/shell_process_windows.go | 3 + pkg/tools/shell_test.go | 352 +++++- pkg/tools/shell_timeout_unix_test.go | 49 + pkg/tools/skills_install.go | 120 +- pkg/tools/skills_install_test.go | 39 +- pkg/tools/skills_search.go | 41 +- pkg/tools/skills_search_test.go | 40 +- pkg/tools/spawn_test.go | 22 +- pkg/tools/spi.go | 67 +- pkg/tools/spi_linux.go | 183 ++- pkg/tools/subagent_reporter_test.go | 98 ++ pkg/tools/subagent_tool_test.go | 132 ++- pkg/tools/toolloop.go | 128 +- pkg/tools/toolloop_reporter_test.go | 139 ++- pkg/tools/web.go | 320 +++-- pkg/tools/web_test.go | 200 +++- pkg/tools/workspace_ctx.go | 26 +- 60 files changed, 7487 insertions(+), 888 deletions(-) diff --git a/pkg/agent/context_cache_test.go b/pkg/agent/context_cache_test.go index 0905e8a46..ab0d098cd 100644 --- a/pkg/agent/context_cache_test.go +++ b/pkg/agent/context_cache_test.go @@ -12,70 +12,103 @@ import ( ) // setupWorkspace creates a temporary workspace with standard directories and optional files. + // Returns the tmpDir path; caller should defer os.RemoveAll(tmpDir). + func setupWorkspace(t *testing.T, files map[string]string) string { t.Helper() + tmpDir, err := os.MkdirTemp("", "picoclaw-test-*") if err != nil { t.Fatal(err) } + os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755) + os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755) + for name, content := range files { dir := filepath.Dir(filepath.Join(tmpDir, name)) + os.MkdirAll(dir, 0o755) + if err := os.WriteFile(filepath.Join(tmpDir, name), []byte(content), 0o644); err != nil { t.Fatal(err) } } + return tmpDir } // TestSingleSystemMessage verifies that BuildMessages always produces exactly one + // system message regardless of summary/history variations. + // Fix: multiple system messages break Anthropic (top-level system param) and + // Codex (only reads last system message as instructions). + func TestSingleSystemMessage(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ "IDENTITY.md": "# Identity\nTest agent.", }) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) tests := []struct { - name string + name string + history []providers.Message + summary string + message string }{ { - name: "no summary, no history", + name: "no summary, no history", + summary: "", + message: "hello", }, + { - name: "with summary", + name: "with summary", + summary: "Previous conversation discussed X", + message: "hello", }, + { name: "with history and summary", + history: []providers.Message{ {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "hello"}, }, + summary: strings.Repeat("Long summary text. ", 50), + message: "new message", }, + { name: "system message in history is filtered", + history: []providers.Message{ {Role: "system", Content: "stale system prompt from previous session"}, + {Role: "user", Content: "hi"}, + {Role: "assistant", Content: "hello"}, }, + summary: "", + message: "new message", }, } @@ -85,35 +118,44 @@ func TestSingleSystemMessage(t *testing.T) { msgs := cb.BuildMessages(tt.history, tt.summary, tt.message, nil, "test", "chat1") systemCount := 0 + for _, m := range msgs { if m.Role == "system" { systemCount++ } } + if systemCount != 1 { t.Errorf("expected exactly 1 system message, got %d", systemCount) } + if msgs[0].Role != "system" { t.Errorf("first message should be system, got %s", msgs[0].Role) } + if msgs[len(msgs)-1].Role != "user" { t.Errorf("last message should be user, got %s", msgs[len(msgs)-1].Role) } // System message must contain identity (static) and time (dynamic) + sys := msgs[0].Content + if !strings.Contains(sys, "picoclaw") { t.Error("system message missing identity") } + if !strings.Contains(sys, "Current Time") { t.Error("system message missing dynamic time context") } // Summary handling + if tt.summary != "" { if !strings.Contains(sys, "CONTEXT_SUMMARY:") { t.Error("summary present but CONTEXT_SUMMARY prefix missing") } + if !strings.Contains(sys, tt.summary[:20]) { t.Error("summary content not found in system message") } @@ -127,29 +169,46 @@ func TestSingleSystemMessage(t *testing.T) { } // TestMtimeAutoInvalidation verifies that the cache detects source file changes + // via mtime without requiring explicit InvalidateCache(). + // Fix: original implementation had no auto-invalidation — edits to bootstrap files, + // memory, or skills were invisible until process restart. + func TestMtimeAutoInvalidation(t *testing.T) { tests := []struct { - name string - file string // relative path inside workspace - contentV1 string - contentV2 string + name string + + file string // relative path inside workspace + + contentV1 string + + contentV2 string + checkField string // substring to verify in rebuilt prompt }{ { - name: "bootstrap file change", - file: "IDENTITY.md", - contentV1: "# Original Identity", - contentV2: "# Updated Identity", + name: "bootstrap file change", + + file: "IDENTITY.md", + + contentV1: "# Original Identity", + + contentV2: "# Updated Identity", + checkField: "Updated Identity", }, + { - name: "memory file change", - file: "memory/MEMORY.md", - contentV1: "# Memory\nUser likes Go.", - contentV2: "# Memory\nUser likes Rust.", + name: "memory file change", + + file: "memory/MEMORY.md", + + contentV1: "# Memory\nUser likes Go.", + + contentV2: "# Memory\nUser likes Rust.", + checkField: "User likes Rust", }, } @@ -157,6 +216,7 @@ func TestMtimeAutoInvalidation(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{tt.file: tt.contentV1}) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) @@ -164,26 +224,39 @@ func TestMtimeAutoInvalidation(t *testing.T) { sp1 := cb.BuildSystemPromptWithCache() // Overwrite file and set future mtime to ensure detection. + // Use 2s offset for filesystem mtime resolution safety (some FS + // have 1s or coarser granularity, especially in CI containers). + fullPath := filepath.Join(tmpDir, tt.file) + os.WriteFile(fullPath, []byte(tt.contentV2), 0o644) + future := time.Now().Add(2 * time.Second) + os.Chtimes(fullPath, future, future) // Verify sourceFilesChangedLocked detects the mtime change + cb.systemPromptMutex.RLock() + changed := cb.sourceFilesChangedLocked() + cb.systemPromptMutex.RUnlock() + if !changed { t.Fatalf("sourceFilesChangedLocked() should detect %s change", tt.file) } // Should auto-rebuild without explicit InvalidateCache() + sp2 := cb.BuildSystemPromptWithCache() + if sp1 == sp2 { t.Errorf("cache not rebuilt after %s change", tt.file) } + if !strings.Contains(sp2, tt.checkField) { t.Errorf("rebuilt prompt missing expected content %q", tt.checkField) } @@ -191,23 +264,34 @@ func TestMtimeAutoInvalidation(t *testing.T) { } // Skills directory mtime change + t.Run("skills dir change", func(t *testing.T) { tmpDir := setupWorkspace(t, nil) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) + _ = cb.BuildSystemPromptWithCache() // populate cache // Touch skills directory (simulate new skill installed) + skillsDir := filepath.Join(tmpDir, "skills") + future := time.Now().Add(2 * time.Second) + os.Chtimes(skillsDir, future, future) // Verify sourceFilesChangedLocked detects it (cache is rebuilt) + // We confirm by checking internal state: a second call should rebuild. + cb.systemPromptMutex.RLock() + changed := cb.sourceFilesChangedLocked() + cb.systemPromptMutex.RUnlock() + if !changed { t.Error("sourceFilesChangedLocked() should detect skills dir mtime change") } @@ -215,17 +299,22 @@ func TestMtimeAutoInvalidation(t *testing.T) { } // TestExplicitInvalidateCache verifies that InvalidateCache() forces a rebuild + // even when source files haven't changed (useful for tests and reload commands). + func TestExplicitInvalidateCache(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ "IDENTITY.md": "# Test Identity", }) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) sp1 := cb.BuildSystemPromptWithCache() + cb.InvalidateCache() + sp2 := cb.BuildSystemPromptWithCache() if sp1 != sp2 { @@ -233,29 +322,39 @@ func TestExplicitInvalidateCache(t *testing.T) { } // Verify cachedAt was reset + cb.InvalidateCache() + cb.systemPromptMutex.RLock() + if !cb.cachedAt.IsZero() { t.Error("cachedAt should be zero after InvalidateCache()") } + cb.systemPromptMutex.RUnlock() } // TestCacheStability verifies that the static prompt is stable across repeated calls + // when no files change (regression test for issue #607). + func TestCacheStability(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ "IDENTITY.md": "# Identity\nContent", - "SOUL.md": "# Soul\nContent", + + "SOUL.md": "# Soul\nContent", }) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) results := make([]string, 5) + for i := range results { results[i] = cb.BuildSystemPromptWithCache() } + for i := 1; i < len(results); i++ { if results[i] != results[0] { t.Errorf("cached prompt changed between call 0 and %d", i) @@ -263,32 +362,47 @@ func TestCacheStability(t *testing.T) { } // Static prompt must NOT contain per-request data + if strings.Contains(results[0], "Current Time") { t.Error("static cached prompt should not contain time (added dynamically)") } } // TestNewFileCreationInvalidatesCache verifies that creating a source file that + // did not exist when the cache was built triggers a cache rebuild. + // This catches the "from nothing to something" edge case that the old + // modifiedSince (return false on stat error) would miss. + func TestNewFileCreationInvalidatesCache(t *testing.T) { tests := []struct { - name string - file string // relative path inside workspace - content string + name string + + file string // relative path inside workspace + + content string + checkField string // substring to verify in rebuilt prompt }{ { - name: "new bootstrap file", - file: "SOUL.md", - content: "# Soul\nBe kind and helpful.", + name: "new bootstrap file", + + file: "SOUL.md", + + content: "# Soul\nBe kind and helpful.", + checkField: "Be kind and helpful", }, + { - name: "new memory file", - file: "memory/MEMORY.md", - content: "# Memory\nUser prefers dark mode.", + name: "new memory file", + + file: "memory/MEMORY.md", + + content: "# Memory\nUser prefers dark mode.", + checkField: "User prefers dark mode", }, } @@ -296,29 +410,41 @@ func TestNewFileCreationInvalidatesCache(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Start with an empty workspace (no bootstrap/memory files) + tmpDir := setupWorkspace(t, nil) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) // Populate cache — file does not exist yet + sp1 := cb.BuildSystemPromptWithCache() + if strings.Contains(sp1, tt.checkField) { t.Fatalf("prompt should not contain %q before file is created", tt.checkField) } // Create the file after cache was built + fullPath := filepath.Join(tmpDir, tt.file) + os.MkdirAll(filepath.Dir(fullPath), 0o755) + if err := os.WriteFile(fullPath, []byte(tt.content), 0o644); err != nil { t.Fatal(err) } + // Set future mtime to guarantee detection + future := time.Now().Add(2 * time.Second) + os.Chtimes(fullPath, future, future) // Cache should auto-invalidate because file went from absent -> present + sp2 := cb.BuildSystemPromptWithCache() + if !strings.Contains(sp2, tt.checkField) { t.Errorf("cache not invalidated on new file creation: expected %q in prompt", tt.checkField) } @@ -327,110 +453,163 @@ func TestNewFileCreationInvalidatesCache(t *testing.T) { } // TestSkillFileContentChange verifies that modifying a skill file's content + // (not just the directory structure) invalidates the cache. + // This is the scenario where directory mtime alone is insufficient — on most + // filesystems, editing a file inside a directory does NOT update the parent + // directory's mtime. + func TestSkillFileContentChange(t *testing.T) { skillMD := `--- + name: test-skill + description: "A test skill" + --- + # Test Skill v1 + Original content.` tmpDir := setupWorkspace(t, map[string]string{ "skills/test-skill/SKILL.md": skillMD, }) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) // Populate cache + sp1 := cb.BuildSystemPromptWithCache() + _ = sp1 // cache is warm // Modify the skill file content (without touching the skills/ directory) + updatedSkillMD := `--- + name: test-skill + description: "An updated test skill" + --- + # Test Skill v2 + Updated content.` skillPath := filepath.Join(tmpDir, "skills", "test-skill", "SKILL.md") + if err := os.WriteFile(skillPath, []byte(updatedSkillMD), 0o644); err != nil { t.Fatal(err) } + // Set future mtime on the skill file only (NOT the directory) + future := time.Now().Add(2 * time.Second) + os.Chtimes(skillPath, future, future) // Verify that sourceFilesChangedLocked detects the content change + cb.systemPromptMutex.RLock() + changed := cb.sourceFilesChangedLocked() + cb.systemPromptMutex.RUnlock() + if !changed { t.Error("sourceFilesChangedLocked() should detect skill file content change") } // Verify cache is actually rebuilt with new content + sp2 := cb.BuildSystemPromptWithCache() + if sp1 == sp2 && strings.Contains(sp1, "test-skill") { // If the skill appeared in the prompt and the prompt didn't change, + // the cache was not invalidated. + t.Error("cache should be invalidated when skill file content changes") } } // TestConcurrentBuildSystemPromptWithCache verifies that multiple goroutines + // can safely call BuildSystemPromptWithCache concurrently without producing + // empty results, panics, or data races. + // Run with: go test -race ./pkg/agent/ -run TestConcurrentBuildSystemPromptWithCache + func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ - "IDENTITY.md": "# Identity\nConcurrency test agent.", - "SOUL.md": "# Soul\nBe helpful.", - "memory/MEMORY.md": "# Memory\nUser prefers Go.", + "IDENTITY.md": "# Identity\nConcurrency test agent.", + + "SOUL.md": "# Soul\nBe helpful.", + + "memory/MEMORY.md": "# Memory\nUser prefers Go.", + "skills/demo/SKILL.md": "---\nname: demo\ndescription: \"demo skill\"\n---\n# Demo", }) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) const goroutines = 20 + const iterations = 50 var wg sync.WaitGroup + errs := make(chan string, goroutines*iterations) for g := range goroutines { wg.Add(1) + go func(id int) { defer wg.Done() + for i := range iterations { result := cb.BuildSystemPromptWithCache() + if result == "" { errs <- "empty prompt returned" + return } + if !strings.Contains(result, "picoclaw") { errs <- "prompt missing identity" + return } // Also exercise BuildMessages concurrently + msgs := cb.BuildMessages(nil, "", "hello", nil, "test", "chat") + if len(msgs) < 2 { errs <- "BuildMessages returned fewer than 2 messages" + return } + if msgs[0].Role != "system" { errs <- "first message not system" + return } // Occasionally invalidate to exercise the write path + if i%10 == 0 { cb.InvalidateCache() } @@ -439,6 +618,7 @@ func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { } wg.Wait() + close(errs) for errMsg := range errs { @@ -449,64 +629,90 @@ func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { // BenchmarkBuildMessagesWithCache measures caching performance. // TestEmptyWorkspaceBaselineDetectsNewFiles verifies that when the cache is + // built on an empty workspace (no tracked files exist), creating a file + // afterwards still triggers cache invalidation. This validates the + // time.Unix(1, 0) fallback for maxMtime: any real file's mtime is after epoch, + // so fileChangedSince correctly detects the absent -> present transition AND + // the mtime comparison succeeds even without artificially inflated Chtimes. + func TestEmptyWorkspaceBaselineDetectsNewFiles(t *testing.T) { // Empty workspace: no bootstrap files, no memory, no skills content. + tmpDir := setupWorkspace(t, nil) + defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) // Build cache — all tracked files are absent, maxMtime falls back to epoch. + sp1 := cb.BuildSystemPromptWithCache() // Create a bootstrap file with natural mtime (no Chtimes manipulation). + // The file's mtime should be the current wall-clock time, which is + // strictly after time.Unix(1, 0). + soulPath := filepath.Join(tmpDir, "SOUL.md") + if err := os.WriteFile(soulPath, []byte("# Soul\nNewly created."), 0o644); err != nil { t.Fatal(err) } // Cache should detect the new file via existedAtCache (absent -> present). + cb.systemPromptMutex.RLock() + changed := cb.sourceFilesChangedLocked() + cb.systemPromptMutex.RUnlock() + if !changed { t.Fatal("sourceFilesChangedLocked should detect newly created file on empty workspace") } sp2 := cb.BuildSystemPromptWithCache() + if !strings.Contains(sp2, "Newly created") { t.Error("rebuilt prompt should contain new file content") } + if sp1 == sp2 { t.Error("cache should have been invalidated after file creation") } } // BenchmarkBuildMessagesWithCache measures caching performance. + func BenchmarkBuildMessagesWithCache(b *testing.B) { tmpDir, _ := os.MkdirTemp("", "picoclaw-bench-*") + defer os.RemoveAll(tmpDir) os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755) + os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755) + for _, name := range []string{"IDENTITY.md", "SOUL.md", "USER.md"} { os.WriteFile(filepath.Join(tmpDir, name), []byte(strings.Repeat("Content.\n", 10)), 0o644) } cb := NewContextBuilder(tmpDir) + history := []providers.Message{ {Role: "user", Content: "previous message"}, + {Role: "assistant", Content: "previous response"}, } b.ResetTimer() + for i := 0; i < b.N; i++ { _ = cb.BuildMessages(history, "summary", "new message", nil, "cli", "test") } diff --git a/pkg/agent/context_test.go b/pkg/agent/context_test.go index e023c9c30..d72d0aca9 100644 --- a/pkg/agent/context_test.go +++ b/pkg/agent/context_test.go @@ -12,9 +12,11 @@ func msg(role, content string) providers.Message { func assistantWithTools(toolIDs ...string) providers.Message { calls := make([]providers.ToolCall, len(toolIDs)) + for i, id := range toolIDs { calls[i] = providers.ToolCall{ID: id, Type: "function"} } + return providers.Message{Role: "assistant", ToolCalls: calls} } @@ -24,11 +26,13 @@ func toolResult(id string) providers.Message { func TestSanitizeHistoryForProvider_EmptyHistory(t *testing.T) { result := sanitizeHistoryForProvider(nil) + if len(result) != 0 { t.Fatalf("expected empty, got %d messages", len(result)) } result = sanitizeHistoryForProvider([]providers.Message{}) + if len(result) != 0 { t.Fatalf("expected empty, got %d messages", len(result)) } @@ -37,170 +41,228 @@ func TestSanitizeHistoryForProvider_EmptyHistory(t *testing.T) { func TestSanitizeHistoryForProvider_SingleToolCall(t *testing.T) { history := []providers.Message{ msg("user", "hello"), + assistantWithTools("A"), + toolResult("A"), + msg("assistant", "done"), } result := sanitizeHistoryForProvider(history) + if len(result) != 4 { t.Fatalf("expected 4 messages, got %d", len(result)) } + assertRoles(t, result, "user", "assistant", "tool", "assistant") } func TestSanitizeHistoryForProvider_MultiToolCalls(t *testing.T) { history := []providers.Message{ msg("user", "do two things"), + assistantWithTools("A", "B"), + toolResult("A"), + toolResult("B"), + msg("assistant", "both done"), } result := sanitizeHistoryForProvider(history) + if len(result) != 5 { t.Fatalf("expected 5 messages, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant") } func TestSanitizeHistoryForProvider_AssistantToolCallAfterPlainAssistant(t *testing.T) { history := []providers.Message{ msg("user", "hi"), + msg("assistant", "thinking"), + assistantWithTools("A"), + toolResult("A"), } result := sanitizeHistoryForProvider(history) + if len(result) != 2 { t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user", "assistant") } func TestSanitizeHistoryForProvider_OrphanedLeadingTool(t *testing.T) { history := []providers.Message{ toolResult("A"), + msg("user", "hello"), } result := sanitizeHistoryForProvider(history) + if len(result) != 1 { t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user") } func TestSanitizeHistoryForProvider_ToolAfterUserDropped(t *testing.T) { history := []providers.Message{ msg("user", "hello"), + toolResult("A"), } result := sanitizeHistoryForProvider(history) + if len(result) != 1 { t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user") } func TestSanitizeHistoryForProvider_ToolAfterAssistantNoToolCalls(t *testing.T) { history := []providers.Message{ msg("user", "hello"), + msg("assistant", "hi"), + toolResult("A"), } result := sanitizeHistoryForProvider(history) + if len(result) != 2 { t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user", "assistant") } func TestSanitizeHistoryForProvider_AssistantToolCallAtStart(t *testing.T) { history := []providers.Message{ assistantWithTools("A"), + toolResult("A"), + msg("user", "hello"), } result := sanitizeHistoryForProvider(history) + if len(result) != 1 { t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user") } func TestSanitizeHistoryForProvider_MultiToolCallsThenNewRound(t *testing.T) { history := []providers.Message{ msg("user", "do two things"), + assistantWithTools("A", "B"), + toolResult("A"), + toolResult("B"), + msg("assistant", "done"), + msg("user", "hi"), + assistantWithTools("C"), + toolResult("C"), + msg("assistant", "done again"), } result := sanitizeHistoryForProvider(history) + if len(result) != 9 { t.Fatalf("expected 9 messages, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "user", "assistant", "tool", "assistant") } func TestSanitizeHistoryForProvider_ConsecutiveMultiToolRounds(t *testing.T) { history := []providers.Message{ msg("user", "start"), + assistantWithTools("A", "B"), + toolResult("A"), + toolResult("B"), + assistantWithTools("C", "D"), + toolResult("C"), + toolResult("D"), + msg("assistant", "all done"), } result := sanitizeHistoryForProvider(history) + if len(result) != 8 { t.Fatalf("expected 8 messages, got %d: %+v", len(result), roles(result)) } + assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "tool", "tool", "assistant") } func TestSanitizeHistoryForProvider_PlainConversation(t *testing.T) { history := []providers.Message{ msg("user", "hello"), + msg("assistant", "hi"), + msg("user", "how are you"), + msg("assistant", "fine"), } result := sanitizeHistoryForProvider(history) + if len(result) != 4 { t.Fatalf("expected 4 messages, got %d", len(result)) } + assertRoles(t, result, "user", "assistant", "user", "assistant") } func roles(msgs []providers.Message) []string { r := make([]string, len(msgs)) + for i, m := range msgs { r[i] = m.Role } + return r } func assertRoles(t *testing.T, msgs []providers.Message, expected ...string) { t.Helper() + if len(msgs) != len(expected) { t.Fatalf("role count mismatch: got %v, want %v", roles(msgs), expected) } + for i, exp := range expected { if msgs[i].Role != exp { t.Errorf("message[%d]: got role %q, want %q", i, msgs[i].Role, exp) diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index af1bf2ead..9762d0a3c 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -12,28 +12,35 @@ func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 1234, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 1234, + MaxToolIterations: 5, }, }, } configuredTemp := 1.0 + cfg.Agents.Defaults.Temperature = &configuredTemp provider := &mockProvider{} + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if agent.MaxTokens != 1234 { t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234) } + if agent.Temperature != 1.0 { t.Fatalf("Temperature = %f, want %f", agent.Temperature, 1.0) } @@ -44,23 +51,29 @@ func TestNewAgentInstance_DefaultsTemperatureWhenZero(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 1234, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 1234, + MaxToolIterations: 5, }, }, } configuredTemp := 0.0 + cfg.Agents.Defaults.Temperature = &configuredTemp provider := &mockProvider{} + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if agent.Temperature != 0.0 { @@ -73,20 +86,25 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 1234, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 1234, + MaxToolIterations: 5, }, }, } provider := &mockProvider{} + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if agent.Temperature != 0.7 { @@ -99,33 +117,41 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ Workspace: tmpDir, - Model: "step-3.5-flash", + + Model: "step-3.5-flash", }, }, + ModelList: []config.ModelConfig{ { ModelName: "step-3.5-flash", - Model: "openrouter/stepfun/step-3.5-flash:free", - APIBase: "https://openrouter.ai/api/v1", + + Model: "openrouter/stepfun/step-3.5-flash:free", + + APIBase: "https://openrouter.ai/api/v1", }, }, } provider := &mockProvider{} + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if len(agent.Candidates) != 1 { t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) } + if agent.Candidates[0].Provider != "openrouter" { t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openrouter") } + if agent.Candidates[0].Model != "stepfun/step-3.5-flash:free" { t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "stepfun/step-3.5-flash:free") } @@ -136,33 +162,41 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAliasWithoutProtocol(t * if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ Workspace: tmpDir, - Model: "glm-5", + + Model: "glm-5", }, }, + ModelList: []config.ModelConfig{ { ModelName: "glm-5", - Model: "glm-5", - APIBase: "https://api.z.ai/api/coding/paas/v4", + + Model: "glm-5", + + APIBase: "https://api.z.ai/api/coding/paas/v4", }, }, } provider := &mockProvider{} + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if len(agent.Candidates) != 1 { t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) } + if agent.Candidates[0].Provider != "openai" { t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openai") } + if agent.Candidates[0].Model != "glm-5" { t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "glm-5") } diff --git a/pkg/agent/loop_reporter_test.go b/pkg/agent/loop_reporter_test.go index ef7ea0905..f6f3059d0 100644 --- a/pkg/agent/loop_reporter_test.go +++ b/pkg/agent/loop_reporter_test.go @@ -12,63 +12,93 @@ import ( ) // makeOrchTestLoop creates a minimal AgentLoop with a temp workspace and + // a real Broadcaster wired as the reporter. + // Returns the loop, the broadcaster, and a cleanup function. + func makeOrchTestLoop(t *testing.T) (*AgentLoop, *orch.Broadcaster) { t.Helper() + tmpDir, err := os.MkdirTemp("", "agent-orch-test-*") if err != nil { t.Fatalf("MkdirTemp: %v", err) } + t.Cleanup(func() { os.RemoveAll(tmpDir) }) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 512, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 512, + MaxToolIterations: 5, }, }, } + al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) + b := orch.NewBroadcaster() + al.SetOrchReporter(b) + return al, b } // collectOrchEvents drains the subscriber channel until an agent_gc event + // arrives or the deadline is exceeded. + func collectOrchEvents(t *testing.T, ch <-chan orch.Event, timeout time.Duration) []orch.Event { t.Helper() + var events []orch.Event + deadline := time.After(timeout) + for { select { case ev := <-ch: + events = append(events, ev) + if ev.Type == "agent_gc" { return events } + case <-deadline: + t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) } } } // TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC verifies that a main + // session processed via ProcessDirect emits the full lifecycle: + // + // agent_spawn(sessionKey) → agent_state(waiting) → agent_gc(completed) + // + // and that the Broadcaster snapshot is empty after the call returns. + func TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC(t *testing.T) { al, b := makeOrchTestLoop(t) + sub := b.Subscribe() + defer b.Unsubscribe(sub) const sessionKey = "orch-test-session" + _, err := al.ProcessDirect(context.Background(), "hello", sessionKey) if err != nil { t.Fatalf("ProcessDirect: %v", err) @@ -77,39 +107,51 @@ func TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC(t *testing.T) { events := collectOrchEvents(t, sub.Ch, 5*time.Second) // First event: agent_spawn with correct ID. + if events[0].Type != "agent_spawn" || events[0].ID != sessionKey { t.Errorf("first event must be agent_spawn(%s), got: %+v", sessionKey, events[0]) } // At least one agent_state(waiting) for this session. + var hasWaiting bool + for _, ev := range events { if ev.Type == "agent_state" && ev.ID == sessionKey && ev.State == "waiting" { hasWaiting = true + break } } + if !hasWaiting { t.Errorf("missing agent_state(waiting) for %s; events: %+v", sessionKey, events) } // Last event: agent_gc(completed) for this session. + last := events[len(events)-1] + if last.Type != "agent_gc" || last.ID != sessionKey || last.Reason != "completed" { t.Errorf("last event must be agent_gc(completed,%s), got: %+v", sessionKey, last) } // Snapshot must be empty — session removed on GC. + if snap := b.Snapshot(); len(snap) != 0 { t.Errorf("snapshot must be empty after GC, got: %v", snap) } } // TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC verifies that heartbeat + // sessions appear on canvas with sessionKey = "heartbeat". + func TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC(t *testing.T) { al, b := makeOrchTestLoop(t) + sub := b.Subscribe() + defer b.Unsubscribe(sub) _, err := al.ProcessHeartbeat(context.Background(), "check system", "heartbeat-chan", "none") @@ -120,12 +162,15 @@ func TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC(t *testing.T) { events := collectOrchEvents(t, sub.Ch, 5*time.Second) // ProcessHeartbeat always uses sessionKey = "heartbeat". + const want = "heartbeat" + if events[0].Type != "agent_spawn" || events[0].ID != want { t.Errorf("first event must be agent_spawn(%s), got: %+v", want, events[0]) } last := events[len(events)-1] + if last.Type != "agent_gc" || last.ID != want || last.Reason != "completed" { t.Errorf("last event must be agent_gc(completed,%s), got: %+v", want, last) } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 7c1781253..ee7f45e81 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -21,55 +21,77 @@ import ( type fakeChannel struct{ id string } -func (f *fakeChannel) Name() string { return "fake" } -func (f *fakeChannel) Start(ctx context.Context) error { return nil } -func (f *fakeChannel) Stop(ctx context.Context) error { return nil } +func (f *fakeChannel) Name() string { return "fake" } + +func (f *fakeChannel) Start(ctx context.Context) error { return nil } + +func (f *fakeChannel) Stop(ctx context.Context) error { return nil } + func (f *fakeChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { return nil } -func (f *fakeChannel) IsRunning() bool { return true } -func (f *fakeChannel) IsAllowed(string) bool { return true } -func (f *fakeChannel) IsAllowedSender(sender bus.SenderInfo) bool { return true } -func (f *fakeChannel) ReasoningChannelID() string { return f.id } + +func (f *fakeChannel) IsRunning() bool { return true } + +func (f *fakeChannel) IsAllowed(string) bool { return true } + +func (f *fakeChannel) IsAllowedSender(sender bus.SenderInfo) bool { return true } + +func (f *fakeChannel) ReasoningChannelID() string { return f.id } func TestRecordLastChannel(t *testing.T) { // Create temp workspace + tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) // Create test config + cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } // Create agent loop + msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) // Test RecordLastChannel + testChannel := "test-channel" + err = al.RecordLastChannel(testChannel) if err != nil { t.Fatalf("RecordLastChannel failed: %v", err) } // Verify channel was saved + lastChannel := al.state.GetLastChannel() + if lastChannel != testChannel { t.Errorf("Expected channel '%s', got '%s'", testChannel, lastChannel) } // Verify persistence by creating a new agent loop + al2 := NewAgentLoop(cfg, msgBus, provider) + if al2.state.GetLastChannel() != testChannel { t.Errorf("Expected persistent channel '%s', got '%s'", testChannel, al2.state.GetLastChannel()) } @@ -77,44 +99,59 @@ func TestRecordLastChannel(t *testing.T) { func TestRecordLastChatID(t *testing.T) { // Create temp workspace + tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) // Create test config + cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } // Create agent loop + msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) // Test RecordLastChatID + testChatID := "test-chat-id-123" + err = al.RecordLastChatID(testChatID) if err != nil { t.Fatalf("RecordLastChatID failed: %v", err) } // Verify chat ID was saved + lastChatID := al.state.GetLastChatID() + if lastChatID != testChatID { t.Errorf("Expected chat ID '%s', got '%s'", testChatID, lastChatID) } // Verify persistence by creating a new agent loop + al2 := NewAgentLoop(cfg, msgBus, provider) + if al2.state.GetLastChatID() != testChatID { t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, al2.state.GetLastChatID()) } @@ -125,24 +162,31 @@ func TestRecordLastHeartbeatTarget(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) target := "telegram:-100123/42" + if err := al.RecordLastHeartbeatTarget(target); err != nil { t.Fatalf("RecordLastHeartbeatTarget failed: %v", err) } @@ -154,226 +198,299 @@ func TestRecordLastHeartbeatTarget(t *testing.T) { func TestNewAgentLoop_StateInitialized(t *testing.T) { // Create temp workspace + tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) // Create test config + cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } // Create agent loop + msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) // Verify state manager is initialized + if al.state == nil { t.Error("Expected state manager to be initialized") } // Verify state directory was created + stateDir := filepath.Join(tmpDir, "state") + if _, err := os.Stat(stateDir); os.IsNotExist(err) { t.Error("Expected state directory to exist") } } // TestToolRegistry_ToolRegistration verifies tools can be registered and retrieved + func TestToolRegistry_ToolRegistration(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) // Register a custom tool + customTool := &mockCustomTool{} + al.RegisterTool(customTool) // Verify tool is registered by checking it doesn't panic on GetStartupInfo + // (actual tool retrieval is tested in tools package tests) + info := al.GetStartupInfo() + toolsInfo := info["tools"].(map[string]any) + toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list + found := slices.Contains(toolsList, "mock_custom") + if !found { t.Error("Expected custom tool to be registered") } } // TestToolContext_Updates verifies tool context is updated with channel/chatID + func TestToolContext_Updates(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &simpleMockProvider{response: "OK"} + _ = NewAgentLoop(cfg, msgBus, provider) // Verify that ContextualTool interface is defined and can be implemented + // This test validates the interface contract exists + ctxTool := &mockContextualTool{} // Verify the tool implements the interface correctly + var _ tools.ContextualTool = ctxTool } // TestToolRegistry_GetDefinitions verifies tool definitions can be retrieved + func TestToolRegistry_GetDefinitions(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) // Register a test tool and verify it shows up in startup info + testTool := &mockCustomTool{} + al.RegisterTool(testTool) info := al.GetStartupInfo() + toolsInfo := info["tools"].(map[string]any) + toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list + found := slices.Contains(toolsList, "mock_custom") + if !found { t.Error("Expected custom tool to be registered") } } // TestAgentLoop_GetStartupInfo verifies startup info contains tools + func TestAgentLoop_GetStartupInfo(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) info := al.GetStartupInfo() // Verify tools info exists + toolsInfo, ok := info["tools"] + if !ok { t.Fatal("Expected 'tools' key in startup info") } toolsMap, ok := toolsInfo.(map[string]any) + if !ok { t.Fatal("Expected 'tools' to be a map") } count, ok := toolsMap["count"] + if !ok { t.Fatal("Expected 'count' in tools info") } // Should have default tools registered + if count.(int) == 0 { t.Error("Expected at least some tools to be registered") } } // TestAgentLoop_Stop verifies Stop() sets running to false + func TestAgentLoop_Stop(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) // Note: running is only set to true when Run() is called + // We can't test that without starting the event loop + // Instead, verify the Stop method can be called safely + al.Stop() // Verify running is false (initial state or after Stop) + if al.running.Load() { t.Error("Expected agent to be stopped (or never started)") } @@ -387,13 +504,18 @@ type simpleMockProvider struct { func (m *simpleMockProvider) Chat( ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, ) (*providers.LLMResponse, error) { return &providers.LLMResponse{ - Content: m.response, + Content: m.response, + ToolCalls: []providers.ToolCall{}, }, nil } @@ -403,6 +525,7 @@ func (m *simpleMockProvider) GetDefaultModel() string { } // mockCustomTool is a simple mock tool for registration testing + type mockCustomTool struct{} func (m *mockCustomTool) Name() string { @@ -415,7 +538,8 @@ func (m *mockCustomTool) Description() string { func (m *mockCustomTool) Parameters() map[string]any { return map[string]any{ - "type": "object", + "type": "object", + "properties": map[string]any{}, } } @@ -425,9 +549,11 @@ func (m *mockCustomTool) Execute(ctx context.Context, args map[string]any) *tool } // mockContextualTool tracks context updates + type mockContextualTool struct { lastChannel string - lastChatID string + + lastChatID string } func (m *mockContextualTool) Name() string { @@ -440,7 +566,8 @@ func (m *mockContextualTool) Description() string { func (m *mockContextualTool) Parameters() map[string]any { return map[string]any{ - "type": "object", + "type": "object", + "properties": map[string]any{}, } } @@ -451,133 +578,179 @@ func (m *mockContextualTool) Execute(ctx context.Context, args map[string]any) * func (m *mockContextualTool) SetContext(channel, chatID string) { m.lastChannel = channel + m.lastChatID = chatID } // testHelper executes a message and returns the response + type testHelper struct { al *AgentLoop } func (h testHelper) executeAndGetResponse(tb testing.TB, ctx context.Context, msg bus.InboundMessage) string { // Use a short timeout to avoid hanging + timeoutCtx, cancel := context.WithTimeout(ctx, responseTimeout) + defer cancel() response, err := h.al.processMessage(timeoutCtx, msg) if err != nil { tb.Fatalf("processMessage failed: %v", err) } + return response } const responseTimeout = 3 * time.Second // TestToolResult_SilentToolDoesNotSendUserMessage verifies silent tools don't trigger outbound + func TestToolResult_SilentToolDoesNotSendUserMessage(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &simpleMockProvider{response: "File operation complete"} + al := NewAgentLoop(cfg, msgBus, provider) + helper := testHelper{al: al} // ReadFileTool returns SilentResult, which should not send user message + ctx := context.Background() + msg := bus.InboundMessage{ - Channel: "test", - SenderID: "user1", - ChatID: "chat1", - Content: "read test.txt", + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "read test.txt", + SessionKey: "test-session", } response := helper.executeAndGetResponse(t, ctx, msg) // Silent tool should return the LLM's response directly + if response != "File operation complete" { t.Errorf("Expected 'File operation complete', got: %s", response) } } // TestToolResult_UserFacingToolDoesSendMessage verifies user-facing tools trigger outbound + func TestToolResult_UserFacingToolDoesSendMessage(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &simpleMockProvider{response: "Command output: hello world"} + al := NewAgentLoop(cfg, msgBus, provider) + helper := testHelper{al: al} // ExecTool returns UserResult, which should send user message + ctx := context.Background() + msg := bus.InboundMessage{ - Channel: "test", - SenderID: "user1", - ChatID: "chat1", - Content: "run hello", + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "run hello", + SessionKey: "test-session", } response := helper.executeAndGetResponse(t, ctx, msg) // User-facing tool should include the output in final response + if response != "Command output: hello world" { t.Errorf("Expected 'Command output: hello world', got: %s", response) } } // failFirstMockProvider fails on the first N calls with a specific error + type failFirstMockProvider struct { - failures int + failures int + currentCall int - failError error + + failError error + successResp string } func (m *failFirstMockProvider) Chat( ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, ) (*providers.LLMResponse, error) { m.currentCall++ + if m.currentCall <= m.failures { return nil, m.failError } + return &providers.LLMResponse{ - Content: m.successResp, + Content: m.successResp, + ToolCalls: []providers.ToolCall{}, }, nil } @@ -587,19 +760,24 @@ func (m *failFirstMockProvider) GetDefaultModel() string { } // TestAgentLoop_ContextExhaustionRetry verify that the agent retries on context errors + func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, @@ -608,39 +786,61 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { msgBus := bus.NewMessageBus() // Create a provider that fails once with a context error + contextErr := fmt.Errorf("InvalidParameter: Total tokens of image and text exceed max message tokens") + provider := &failFirstMockProvider{ - failures: 1, - failError: contextErr, + failures: 1, + + failError: contextErr, + successResp: "Recovered from context error", } al := NewAgentLoop(cfg, msgBus, provider) // Inject some history to simulate a full context + sessionKey := "test-session-context" + // Create dummy history + history := []providers.Message{ {Role: "system", Content: "System prompt"}, + {Role: "user", Content: "Old message 1"}, + {Role: "assistant", Content: "Old response 1"}, + {Role: "user", Content: "Old message 2"}, + {Role: "assistant", Content: "Old response 2"}, + {Role: "user", Content: "Trigger message"}, } + defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { t.Fatal("No default agent found") } + defaultAgent.Sessions.SetHistory(sessionKey, history) // Call ProcessDirectWithChannel + // Note: ProcessDirectWithChannel calls processMessage which will execute runLLMIteration + response, err := al.ProcessDirectWithChannel( + context.Background(), + "Trigger message", + sessionKey, + "test", + "test-chat", ) if err != nil { @@ -652,17 +852,25 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { } // We expect 2 calls: 1st failed, 2nd succeeded + if provider.currentCall != 2 { t.Errorf("Expected 2 calls (1 fail + 1 success), got %d", provider.currentCall) } // Check final history length + finalHistory := defaultAgent.Sessions.GetHistory(sessionKey) + // We verify that the history has been modified (compressed) + // Original length: 6 + // Expected behavior: compression drops ~50% of history (mid slice) + // We can assert that the length is NOT what it would be without compression. + // Without compression: 6 + 1 (new user msg) + 1 (assistant msg) = 8 + if len(finalHistory) >= 8 { t.Errorf("Expected history to be compressed (len < 8), got %d", len(finalHistory)) } @@ -670,22 +878,33 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { func TestShouldInjectReminder(t *testing.T) { tests := []struct { - name string + name string + iteration int - interval int - want bool + + interval int + + want bool }{ {"first iteration skipped", 1, 5, false}, + {"iteration 5 interval 5", 5, 5, true}, + {"iteration 10 interval 5", 10, 5, true}, + {"iteration 3 interval 5", 3, 5, false}, + {"interval zero disabled", 5, 0, false}, + {"interval negative disabled", 5, -1, false}, + {"iteration 2 interval 1", 2, 1, true}, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := shouldInjectReminder(tt.iteration, tt.interval) + if got != tt.want { t.Errorf("shouldInjectReminder(%d, %d) = %v, want %v", tt.iteration, tt.interval, got, tt.want) } @@ -698,63 +917,96 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) + chManager, err := channels.NewManager(&config.Config{}, bus.NewMessageBus(), nil) if err != nil { t.Fatalf("Failed to create channel manager: %v", err) } + for name, id := range map[string]string{ - "whatsapp": "rid-whatsapp", - "telegram": "rid-telegram", - "feishu": "rid-feishu", - "discord": "rid-discord", - "maixcam": "rid-maixcam", - "qq": "rid-qq", - "dingtalk": "rid-dingtalk", - "slack": "rid-slack", - "line": "rid-line", - "onebot": "rid-onebot", - "wecom": "rid-wecom", + "whatsapp": "rid-whatsapp", + + "telegram": "rid-telegram", + + "feishu": "rid-feishu", + + "discord": "rid-discord", + + "maixcam": "rid-maixcam", + + "qq": "rid-qq", + + "dingtalk": "rid-dingtalk", + + "slack": "rid-slack", + + "line": "rid-line", + + "onebot": "rid-onebot", + + "wecom": "rid-wecom", + "wecom_app": "rid-wecom-app", } { chManager.RegisterChannel(name, &fakeChannel{id: id}) } + al.SetChannelManager(chManager) + tests := []struct { channel string - wantID string + + wantID string }{ {channel: "whatsapp", wantID: "rid-whatsapp"}, + {channel: "telegram", wantID: "rid-telegram"}, + {channel: "feishu", wantID: "rid-feishu"}, + {channel: "discord", wantID: "rid-discord"}, + {channel: "maixcam", wantID: "rid-maixcam"}, + {channel: "qq", wantID: "rid-qq"}, + {channel: "dingtalk", wantID: "rid-dingtalk"}, + {channel: "slack", wantID: "rid-slack"}, + {channel: "line", wantID: "rid-line"}, + {channel: "onebot", wantID: "rid-onebot"}, + {channel: "wecom", wantID: "rid-wecom"}, + {channel: "wecom_app", wantID: "rid-wecom-app"}, + {channel: "unknown", wantID: ""}, } for _, tt := range tests { t.Run(tt.channel, func(t *testing.T) { got := al.targetReasoningChannelID(tt.channel) + if got != tt.wantID { t.Fatalf("targetReasoningChannelID(%q) = %q, want %q", tt.channel, got, tt.wantID) } @@ -764,18 +1016,23 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) { func TestBuildTaskReminder_WithoutBlocker(t *testing.T) { msg := buildTaskReminder("implement feature X", "") + if msg.Role != "user" { t.Errorf("expected role 'user', got %q", msg.Role) } + if !strings.Contains(msg.Content, "[TASK REMINDER]") { t.Error("expected content to contain '[TASK REMINDER]'") } + if !strings.Contains(msg.Content, "implement feature X") { t.Error("expected content to contain original message") } + if strings.Contains(msg.Content, "blocker") { t.Error("expected content NOT to contain 'blocker' when no blocker provided") } + if !strings.Contains(msg.Content, "move on") { t.Error("expected content to contain completion prompt") } @@ -783,18 +1040,23 @@ func TestBuildTaskReminder_WithoutBlocker(t *testing.T) { func TestBuildTaskReminder_WithBlocker(t *testing.T) { msg := buildTaskReminder("implement feature X", "ModuleNotFoundError: No module named 'foo'") + if msg.Role != "user" { t.Errorf("expected role 'user', got %q", msg.Role) } + if !strings.Contains(msg.Content, "[TASK REMINDER]") { t.Error("expected content to contain '[TASK REMINDER]'") } + if !strings.Contains(msg.Content, "implement feature X") { t.Error("expected content to contain original message") } + if !strings.Contains(msg.Content, "Last blocker") { t.Error("expected content to contain 'Last blocker'") } + if !strings.Contains(msg.Content, "ModuleNotFoundError") { t.Error("expected content to contain blocker text") } @@ -805,38 +1067,51 @@ func TestResolveProvider_CachesProviders(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - Provider: "vllm", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + Provider: "vllm", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, + Providers: config.ProvidersConfig{ VLLM: config.ProviderConfig{ - APIKey: "test-key", + APIKey: "test-key", + APIBase: "https://example.com/v1", }, }, } msgBus := bus.NewMessageBus() + primary := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, primary) // First call creates and caches a provider for "vllm/test-model" + p1 := al.resolveProvider("vllm", "test-model", primary) + if p1 == primary { t.Fatal("expected a new provider from legacy providers config, not the fallback") } // Calling again should return the same instance (cached) + p2 := al.resolveProvider("vllm", "test-model", primary) + if p1 != p2 { t.Fatal("expected same cached instance on second call") } @@ -847,31 +1122,41 @@ func TestResolveProvider_FallsBackOnError(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - Provider: "vllm", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + Provider: "vllm", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + primary := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, primary) // Request a provider that can't be created (no config for "nonexistent") + p := al.resolveProvider("nonexistent", "unknown-model", primary) + if p != primary { t.Fatal("expected fallback to primary provider on creation error") } // Ensure the failed provider is NOT cached + if _, ok := al.providerCache["nonexistent"]; ok { t.Fatal("failed provider should not be cached") } @@ -882,55 +1167,72 @@ func TestResolveProvider_EmptyNameReturnsFallback(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + primary := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, primary) p := al.resolveProvider("", "", primary) + if p != primary { t.Fatal("expected fallback provider for empty name") } } // TestSlashCommandResponseSkipsPlaceholder verifies that slash command responses + // are published with SkipPlaceholder=true so they don't overwrite the ongoing task + // status bubble. + func TestSlashCommandResponseSkipsPlaceholder(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() go func() { @@ -938,15 +1240,21 @@ func TestSlashCommandResponseSkipsPlaceholder(t *testing.T) { }() // Send a slash command + msgBus.PublishInbound(context.Background(), bus.InboundMessage{ - Channel: "telegram", + Channel: "telegram", + SenderID: "user1", - ChatID: "chat1", - Content: "/skills", + + ChatID: "chat1", + + Content: "/skills", }) // Read the outbound message + outMsg, ok := msgBus.SubscribeOutbound(ctx) + if !ok { t.Fatal("expected outbound message from slash command") } @@ -958,26 +1266,35 @@ func TestSlashCommandResponseSkipsPlaceholder(t *testing.T) { func TestBuildTaskReminder_Truncation(t *testing.T) { // Build a long message (1000 runes) + longMsg := strings.Repeat("あ", 1000) + longBlocker := strings.Repeat("X", 500) msg := buildTaskReminder(longMsg, longBlocker) // The full message should NOT contain 1000 'あ' characters + runeCount := strings.Count(msg.Content, "あ") + if runeCount >= 1000 { t.Errorf("expected task message to be truncated, got %d 'あ' runes", runeCount) } + // Should be at most taskReminderMaxChars (500) runes for the task part + if runeCount > taskReminderMaxChars { t.Errorf("expected at most %d task runes, got %d", taskReminderMaxChars, runeCount) } // Blocker should be truncated too + xCount := strings.Count(msg.Content, "X") + if xCount >= 500 { t.Errorf("expected blocker to be truncated, got %d 'X' chars", xCount) } + if xCount > blockerMaxChars { t.Errorf("expected at most %d blocker chars, got %d", blockerMaxChars, xCount) } @@ -985,28 +1302,39 @@ func TestBuildTaskReminder_Truncation(t *testing.T) { func TestBuildPlanReminder(t *testing.T) { tests := []struct { - name string - status string - wantOK bool + name string + + status string + + wantOK bool + wantSubstr string }{ {"interviewing", "interviewing", true, "interviewing the user"}, + {"review", "review", true, "under review"}, + {"executing returns false", "executing", false, ""}, + {"empty returns false", "", false, ""}, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { msg, ok := buildPlanReminder(tt.status) + if ok != tt.wantOK { t.Fatalf("buildPlanReminder(%q) ok = %v, want %v", tt.status, ok, tt.wantOK) } + if !ok { return } + if msg.Role != "user" { t.Errorf("expected role 'user', got %q", msg.Role) } + if !strings.Contains(msg.Content, tt.wantSubstr) { t.Errorf("expected content to contain %q, got %q", tt.wantSubstr, msg.Content) } @@ -1018,34 +1346,46 @@ func TestBuildPlanReminder(t *testing.T) { func newTestAgentLoop(t *testing.T) (*AgentLoop, func()) { t.Helper() + tmpDir, err := os.MkdirTemp("", "agent-plan-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } + msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + return al, func() { os.RemoveAll(tmpDir) } } func TestPlanCommand_ShowNoPlan(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() response, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan"}) + if !handled { t.Fatal("expected /plan to be handled") } + if !strings.Contains(response, "No active plan") { t.Errorf("expected 'No active plan', got %q", response) } @@ -1053,27 +1393,40 @@ func TestPlanCommand_ShowNoPlan(t *testing.T) { func TestSplitChatAndThread(t *testing.T) { tests := []struct { - name string - chatID string + name string + + chatID string + wantChatID string + wantThread int }{ {name: "plain chat", chatID: "-100123", wantChatID: "-100123", wantThread: 0}, + {name: "chat with thread", chatID: "-100123/77", wantChatID: "-100123", wantThread: 77}, + {name: "invalid thread", chatID: "-100123/abc", wantChatID: "-100123", wantThread: 0}, + {name: "empty", chatID: "", wantChatID: "", wantThread: 0}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { gotChatID, gotThread := splitChatAndThread(tt.chatID) + if gotChatID != tt.wantChatID || gotThread != tt.wantThread { t.Fatalf( + "splitChatAndThread(%q) = (%q, %d), want (%q, %d)", + tt.chatID, + gotChatID, + gotThread, + tt.wantChatID, + tt.wantThread, ) } @@ -1083,40 +1436,55 @@ func TestSplitChatAndThread(t *testing.T) { func TestHeartbeatCommandThreadHerePersistsConfig(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() var saved bool + var updatedThread int + al.SetConfigSaver(func(cfg *config.Config) error { saved = true + if cfg.Channels.Telegram.HeartbeatThreadID != 42 { t.Fatalf("HeartbeatThreadID in saver = %d, want 42", cfg.Channels.Telegram.HeartbeatThreadID) } + return nil }) + al.SetHeartbeatThreadUpdater(func(threadID int) { updatedThread = threadID }) msg := bus.InboundMessage{ Content: "/heartbeat thread here", + Channel: "telegram", - ChatID: "-100500/42", + + ChatID: "-100500/42", } + resp, handled := al.handleCommand(context.Background(), msg) + if !handled { t.Fatal("expected /heartbeat command to be handled") } + if !strings.Contains(resp, "Heartbeat thread set to 42") { t.Fatalf("unexpected response: %q", resp) } + if !saved { t.Fatal("expected config saver to be called") } + if updatedThread != 42 { t.Fatalf("updatedThread = %d, want 42", updatedThread) } + if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 42 { t.Fatalf("cfg heartbeat thread = %d, want 42", got) } + if got := al.state.GetHeartbeatTarget(); got != "telegram:-100500" { t.Fatalf("state heartbeat target = %q, want %q", got, "telegram:-100500") } @@ -1124,20 +1492,27 @@ func TestHeartbeatCommandThreadHerePersistsConfig(t *testing.T) { func TestHeartbeatCommandThreadOff(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() al.cfg.Channels.Telegram.HeartbeatThreadID = 99 + resp, handled := al.handleCommand(context.Background(), bus.InboundMessage{ Content: "/heartbeat thread off", + Channel: "telegram", - ChatID: "-100500/42", + + ChatID: "-100500/42", }) + if !handled { t.Fatal("expected /heartbeat command to be handled") } + if !strings.Contains(resp, "disabled") { t.Fatalf("unexpected response: %q", resp) } + if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 0 { t.Fatalf("cfg heartbeat thread = %d, want 0", got) } @@ -1145,33 +1520,45 @@ func TestHeartbeatCommandThreadOff(t *testing.T) { func TestPlanCommand_StartNewPlan(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() // /plan <task> should NOT be handled by handleCommand — it falls through + // to the LLM queue via expandPlanCommand. + _, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan Set up monitoring"}) + if handled { t.Fatal("expected /plan <task> NOT to be handled (should fall through to LLM)") } // expandPlanCommand writes the seed and rewrites the message + msg := bus.InboundMessage{Content: "/plan Set up monitoring"} + expanded, compact, ok := al.expandPlanCommand(msg) + if !ok { t.Fatal("expected expandPlanCommand to succeed") } + if expanded != "Set up monitoring" { t.Errorf("expected expanded = 'Set up monitoring', got %q", expanded) } + if !strings.Contains(compact, "Set up monitoring") { t.Errorf("expected compact to contain task, got %q", compact) } // Verify plan was created + agent := al.registry.GetDefaultAgent() + if !agent.ContextBuilder.HasActivePlan() { t.Error("expected active plan after expandPlanCommand") } + if status := agent.ContextBuilder.GetPlanStatus(); status != "interviewing" { t.Errorf("expected 'interviewing', got %q", status) } @@ -1179,16 +1566,21 @@ func TestPlanCommand_StartNewPlan(t *testing.T) { func TestPlanCommand_StartBlockedByExisting(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() // Start first plan via expandPlanCommand + al.expandPlanCommand(bus.InboundMessage{Content: "/plan First task"}) // Try to start another — handleCommand should block it on the fast path + response, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan Second task"}) + if !handled { t.Fatal("expected second /plan to be handled (blocked)") } + if !strings.Contains(response, "already active") { t.Errorf("expected 'already active', got %q", response) } @@ -1196,16 +1588,21 @@ func TestPlanCommand_StartBlockedByExisting(t *testing.T) { func TestPlanCommand_Clear(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() // Start plan then clear + al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan clear"}) + if !strings.Contains(response, "Plan cleared") { t.Errorf("expected 'Plan cleared', got %q", response) } agent := al.registry.GetDefaultAgent() + if agent.ContextBuilder.HasActivePlan() { t.Error("expected no plan after clear") } @@ -1213,9 +1610,11 @@ func TestPlanCommand_Clear(t *testing.T) { func TestPlanCommand_ClearNoPlan(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan clear"}) + if !strings.Contains(response, "No active plan") { t.Errorf("expected 'No active plan', got %q", response) } @@ -1223,16 +1622,21 @@ func TestPlanCommand_ClearNoPlan(t *testing.T) { func TestPlanCommand_Start(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() // Create interviewing plan with phases (start requires phases) + plan := "# Active Plan\n\n> Task: Test task\n> Status: interviewing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + _ = agent.ContextBuilder.WriteMemory(plan) // Transition to executing via /plan start + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + if !strings.Contains(response, "approved") { t.Errorf("expected 'approved', got %q", response) } @@ -1242,6 +1646,7 @@ func TestPlanCommand_Start(t *testing.T) { } // planStartPending must be set so Run() enqueues an LLM trigger + if !al.planStartPending { t.Error("expected planStartPending to be true after /plan start") } @@ -1249,16 +1654,21 @@ func TestPlanCommand_Start(t *testing.T) { func TestPlanCommand_StartFromReview(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() // Create a plan in review status + plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + _ = agent.ContextBuilder.WriteMemory(plan) // Approve via /plan start + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + if !strings.Contains(response, "approved") { t.Errorf("expected 'approved', got %q", response) } @@ -1274,18 +1684,23 @@ func TestPlanCommand_StartFromReview(t *testing.T) { func TestPlanCommand_StartNoPhases(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() // Create interviewing plan without phases + al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) // Should be blocked because no phases exist + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + if !strings.Contains(response, "no phases") { t.Errorf("expected 'no phases' error, got %q", response) } agent := al.registry.GetDefaultAgent() + if status := agent.ContextBuilder.GetPlanStatus(); status != "interviewing" { t.Errorf("expected status to remain 'interviewing', got %q", status) } @@ -1297,20 +1712,27 @@ func TestPlanCommand_StartNoPhases(t *testing.T) { func TestPlanCommand_StartAlreadyExecuting(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() // Create interviewing plan with phases, then start + plan := "# Active Plan\n\n> Task: Test task\n> Status: interviewing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + _ = agent.ContextBuilder.WriteMemory(plan) + al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) // Clear the flag from the first call (simulating Run() consuming it) + al.planStartPending = false // Try start again — should be rejected + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + if !strings.Contains(response, "already executing") { t.Errorf("expected 'already executing', got %q", response) } @@ -1322,26 +1744,43 @@ func TestPlanCommand_StartAlreadyExecuting(t *testing.T) { func TestPlanCommand_Done(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() + // Write a plan directly with phases + plan := `# Active Plan + + > Task: Test task + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [ ] Step one + - [ ] Step two + + ## Context + Test context + ` + agent.ContextBuilder.WriteMemory(plan) response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan done 1"}) + if !strings.Contains(response, "Marked step 1") { t.Errorf("expected confirmation, got %q", response) } @@ -1349,10 +1788,13 @@ Test context func TestPlanCommand_DoneInvalidStep(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan done abc"}) + if !strings.Contains(response, "positive integer") { t.Errorf("expected step validation error, got %q", response) } @@ -1360,29 +1802,45 @@ func TestPlanCommand_DoneInvalidStep(t *testing.T) { func TestPlanCommand_Add(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() + plan := `# Active Plan + + > Task: Test task + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [ ] Step one + + ## Context + Test context + ` + agent.ContextBuilder.WriteMemory(plan) response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan add New step here"}) + if !strings.Contains(response, "Added step") { t.Errorf("expected 'Added step', got %q", response) } content := agent.ContextBuilder.ReadMemory() + if !strings.Contains(content, "New step here") { t.Error("expected new step in plan content") } @@ -1390,27 +1848,45 @@ Test context func TestPlanCommand_Next(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() + plan := `# Active Plan + + > Task: Test task + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [x] Step one + + ## Phase 2: Deploy + - [ ] Step two + + ## Context + Test + ` + agent.ContextBuilder.WriteMemory(plan) response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan next"}) + if !strings.Contains(response, "phase 2") { t.Errorf("expected 'phase 2', got %q", response) } @@ -1422,83 +1898,128 @@ Test func TestPlanCommand_ShowActivePlan(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() + plan := `# Active Plan + + > Task: Deploy app + > Status: executing + > Phase: 1 + + ## Phase 1: Build + - [x] Compile code + - [ ] Run tests + + ## Context + Production server + ` + agent.ContextBuilder.WriteMemory(plan) response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan"}) + if !strings.Contains(response, "Deploy app") { t.Errorf("expected task name in display, got %q", response) } + if !strings.Contains(response, "Phase 1") { t.Errorf("expected phase info in display, got %q", response) } } // TestAutoPhaseAdvance verifies that auto-advance sends notification after LLM iteration + // when current phase is complete. + func TestAutoPhaseAdvance(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-auto-advance-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &simpleMockProvider{response: "OK"} + al := NewAgentLoop(cfg, msgBus, provider) agent := al.registry.GetDefaultAgent() + if agent == nil { t.Fatal("No default agent") } // Write plan with phase 1 complete + plan := `# Active Plan + + > Task: Test auto advance + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [x] Step one + - [x] Step two + + ## Phase 2: Deploy + - [ ] Step three + + ## Context + Test + ` + agent.ContextBuilder.WriteMemory(plan) // Process a message which triggers runAgentLoop + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() _, err = al.ProcessDirectWithChannel(ctx, "continue", "auto-advance-test", "test", "chat1") @@ -1507,56 +2028,80 @@ Test } // After processing, phase should be auto-advanced + if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 2 { t.Errorf("expected phase auto-advanced to 2, got %d", phase) } } // TestAutoCompleteClears verifies that plan is marked completed with correct phase when all phases are complete. + func TestAutoCompleteClears(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-auto-complete-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() + provider := &simpleMockProvider{response: "All done"} + al := NewAgentLoop(cfg, msgBus, provider) agent := al.registry.GetDefaultAgent() + if agent == nil { t.Fatal("No default agent") } // Write fully complete plan + plan := `# Active Plan + + > Task: Test auto complete + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [x] Step one + - [x] Step two + + ## Context + Test + ` + agent.ContextBuilder.WriteMemory(plan) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() _, err = al.ProcessDirectWithChannel(ctx, "finish up", "auto-complete-test", "test", "chat1") @@ -1565,12 +2110,15 @@ Test } // Plan should be kept with status "completed" and phase set to total + if !agent.ContextBuilder.HasActivePlan() { t.Error("expected plan to be retained after completion") } + if status := agent.ContextBuilder.GetPlanStatus(); status != "completed" { t.Errorf("expected plan status 'completed', got %q", status) } + if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 1 { t.Errorf("expected phase 1 (total phases), got %d", phase) } @@ -1579,60 +2127,111 @@ Test func TestIsToolAllowedDuringInterview_FuzzyNames(t *testing.T) { tests := []struct { name string + args map[string]any + want bool }{ // Exact names — read tools allowed + {"read_file", nil, true}, + {"list_dir", nil, true}, + {"web_search", nil, true}, + {"web_fetch", nil, true}, + // Fuzzy variants — should also be allowed + {"readfile", nil, true}, + {"ReadFile", nil, true}, + {"listdir", nil, true}, + {"websearch", nil, true}, + {"webfetch", nil, true}, + // Message tool — allowed (needed for interview questions) + {"message", nil, true}, + {"Message", nil, true}, + // Write to MEMORY.md — allowed + {"edit_file", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, + {"editfile", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, + {"EditFile", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, + // Write to non-MEMORY.md — blocked + {"edit_file", map[string]any{"path": "/ws/main.go"}, false}, + {"editfile", map[string]any{"path": "/ws/main.go"}, false}, + // exec — read-only commands allowed + {"exec", map[string]any{"command": "find . -name '*.py'"}, true}, + {"exec", map[string]any{"command": "ls -la"}, true}, + {"exec", map[string]any{"command": "grep -r TODO ."}, true}, + {"exec", map[string]any{"command": "cat README.md"}, true}, + // exec — cd prefix stripped + {"exec", map[string]any{"command": "cd /home/user/project && find . -type f"}, true}, + {"exec", map[string]any{"command": "cd /tmp && rm -rf *"}, false}, + // exec — write operators blocked + {"exec", map[string]any{"command": "find . > output.txt"}, false}, + {"exec", map[string]any{"command": "ls -la >> log.txt"}, false}, + {"exec", map[string]any{"command": "cat foo | tee bar.txt"}, false}, + // exec — path traversal blocked + {"exec", map[string]any{"command": "cat ../../etc/passwd"}, false}, + {"exec", map[string]any{"command": "find ../../"}, false}, + {"exec", map[string]any{"command": "ls ../secret"}, false}, + // exec — absolute paths blocked + {"exec", map[string]any{"command": "cat /etc/passwd"}, false}, + {"exec", map[string]any{"command": "find /etc -name '*.conf'"}, false}, + {"exec", map[string]any{"command": "ls /root"}, false}, + // exec — write commands blocked + {"exec", map[string]any{"command": "rm -rf /"}, false}, + {"exec", map[string]any{"command": "mv a b"}, false}, + // exec — no args / empty command blocked + {"exec", nil, false}, + {"exec", map[string]any{"command": ""}, false}, + {"Exec", nil, false}, } + for _, tt := range tests { got := isToolAllowedDuringInterview(tt.name, tt.args) + if got != tt.want { t.Errorf("isToolAllowedDuringInterview(%q, %v) = %v, want %v", tt.name, tt.args, got, tt.want) } @@ -1641,69 +2240,109 @@ func TestIsToolAllowedDuringInterview_FuzzyNames(t *testing.T) { func TestBuildArgsSnippet_ExecStripsCD(t *testing.T) { tests := []struct { - name string - tool string - args map[string]any + name string + + tool string + + args map[string]any + workspace string - wantSnip string + + wantSnip string }{ { name: "exec strips cd prefix", + tool: "exec", + args: map[string]any{ "command": "cd /home/user/workspace/project/my-projects && pytest tests/test_integration.py", }, + workspace: "/home/user/workspace", - wantSnip: "pytest tests/test_integration.py", + + wantSnip: "pytest tests/test_integration.py", }, + { - name: "exec no cd prefix, flags stripped", - tool: "exec", - args: map[string]any{"command": "ls -la"}, + name: "exec no cd prefix, flags stripped", + + tool: "exec", + + args: map[string]any{"command": "ls -la"}, + workspace: "/ws", - wantSnip: "ls", + + wantSnip: "ls", }, + { - name: "exec empty command", - tool: "exec", - args: map[string]any{}, + name: "exec empty command", + + tool: "exec", + + args: map[string]any{}, + workspace: "/ws", - wantSnip: "{}", + + wantSnip: "{}", }, + { - name: "read_file strips workspace", - tool: "read_file", - args: map[string]any{"path": "/home/user/workspace/src/main.go"}, + name: "read_file strips workspace", + + tool: "read_file", + + args: map[string]any{"path": "/home/user/workspace/src/main.go"}, + workspace: "/home/user/workspace", - wantSnip: "src/main.go", + + wantSnip: "src/main.go", }, + { - name: "edit_file shows path", - tool: "edit_file", - args: map[string]any{"path": "/ws/config.json", "old_text": "old value here"}, + name: "edit_file shows path", + + tool: "edit_file", + + args: map[string]any{"path": "/ws/config.json", "old_text": "old value here"}, + workspace: "/ws", - wantSnip: "config.json", + + wantSnip: "config.json", }, + { name: "file tool long path prioritizes filename", + tool: "read_file", + args: map[string]any{ "path": "/ws/projects/terra-py-form/src/terra_py_form/hot/state/backend.py", }, + workspace: "/ws", - wantSnip: "projects/terra-py-form/src/terra_py_form/hot/sta\u2026/backend.py", + + wantSnip: "projects/terra-py-form/src/terra_py_form/hot/sta\u2026/backend.py", }, + { - name: "unknown tool shows raw JSON", - tool: "web_search", - args: map[string]any{"query": "hello"}, + name: "unknown tool shows raw JSON", + + tool: "web_search", + + args: map[string]any{"query": "hello"}, + workspace: "/ws", - wantSnip: `{"query":"hello"}`, + + wantSnip: `{"query":"hello"}`, }, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := buildArgsSnippet(tt.tool, tt.args, tt.workspace) + if got != tt.wantSnip { t.Errorf("buildArgsSnippet(%q) = %q, want %q", tt.tool, got, tt.wantSnip) } @@ -1713,63 +2352,96 @@ func TestBuildArgsSnippet_ExecStripsCD(t *testing.T) { func TestFormatCompactEntry(t *testing.T) { tests := []struct { - name string - entry toolLogEntry - wantSub string // must be a substring + name string + + entry toolLogEntry + + wantSub string // must be a substring + wantMark string // result marker must appear - noTime bool // if true, duration should NOT appear + + noTime bool // if true, duration should NOT appear }{ { - name: "exec short entry", - entry: toolLogEntry{Name: "[1] exec", ArgsSnip: "ls", Result: "✓ 1.0s"}, - wantSub: "exec ls", + name: "exec short entry", + + entry: toolLogEntry{Name: "[1] exec", ArgsSnip: "ls", Result: "✓ 1.0s"}, + + wantSub: "exec ls", + wantMark: "✓ 1.0s", // exec keeps duration + }, + { name: "exec long entry truncated from end", + entry: toolLogEntry{ - Name: "[2] exec", + Name: "[2] exec", + ArgsSnip: "pytest tests/integration/test_very_long_name.py", - Result: "✗ 3.0s", + + Result: "✗ 3.0s", }, + wantMark: "✗", }, + { name: "file tool omits duration, shows filename", + entry: toolLogEntry{ - Name: "[3] edit_file", + Name: "[3] edit_file", + ArgsSnip: "projects/terra/src/deep/nested/backend.py", - Result: "✓ 0.0s", + + Result: "✓ 0.0s", }, - wantSub: "backend.py", + + wantSub: "backend.py", + wantMark: "✓", - noTime: true, + + noTime: true, }, + { name: "file tool path truncates from start", + entry: toolLogEntry{ - Name: "[4] read_file", + Name: "[4] read_file", + ArgsSnip: "projects/terra-py-form/src/terra_py_form/hot/state/backend.py", - Result: "✓ 0.1s", + + Result: "✓ 0.1s", }, - wantSub: "backend.py", + + wantSub: "backend.py", + wantMark: "✓", - noTime: true, + + noTime: true, }, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := formatCompactEntry(tt.entry) + if tt.wantSub != "" && !strings.Contains(got, tt.wantSub) { t.Errorf("expected to contain %q, got: %q", tt.wantSub, got) } + if !strings.Contains(got, tt.wantMark) { t.Errorf("result marker %q missing from: %q", tt.wantMark, got) } + if tt.noTime && strings.Contains(got, "0s") { t.Errorf("file tool should omit duration, got: %q", got) } + // Must not exceed maxEntryLineWidth + if runeLen := len([]rune(got)); runeLen > maxEntryLineWidth { t.Errorf("entry too wide: %d runes (max %d): %q", runeLen, maxEntryLineWidth, got) } @@ -1780,10 +2452,14 @@ func TestFormatCompactEntry(t *testing.T) { func TestBuildRichStatus(t *testing.T) { task := &activeTask{ Iteration: 3, - MaxIter: 20, + + MaxIter: 20, + toolLog: []toolLogEntry{ {Name: "exec", ArgsSnip: "ls -la", Result: "✓ 1.2s"}, + {Name: "exec", ArgsSnip: "pytest tests/", Result: "✓ 5.0s"}, + {Name: "read_file", ArgsSnip: "src/main.go", Result: "⏳"}, }, } @@ -1792,10 +2468,15 @@ func TestBuildRichStatus(t *testing.T) { mustContain := []string{ "Task in progress (3/20)", + "my-projects", + "read_file", // latest entry + "No errors", // no error yet + } + for _, s := range mustContain { if !strings.Contains(got, s) { t.Errorf("expected output to contain %q, got:\n%s", s, got) @@ -1803,12 +2484,15 @@ func TestBuildRichStatus(t *testing.T) { } // Non-background: should NOT have reply prompt + if strings.Contains(got, "Reply to intervene") { t.Error("non-background task should not have reply prompt") } // Background: should have reply prompt + bgGot := buildRichStatus(task, true, "/home/user/my-projects") + if !strings.Contains(bgGot, "Reply to intervene") { t.Error("background task should have reply prompt") } @@ -1816,43 +2500,60 @@ func TestBuildRichStatus(t *testing.T) { func TestBuildRichStatus_ProjectDir(t *testing.T) { // exec-based projectDir takes priority + task := &activeTask{ - Iteration: 1, - MaxIter: 10, + Iteration: 1, + + MaxIter: 10, + projectDir: "terra-py-form", + toolLog: []toolLogEntry{ {Name: "exec", ArgsSnip: "ls", Result: "✓ 0.1s"}, }, } + got := buildRichStatus(task, false, "/home/user/.picoclaw/workspace") + if !strings.Contains(got, "terra-py-form") { t.Errorf("expected projectDir in output, got:\n%s", got) } // fileCommonDir fallback + task2 := &activeTask{ - Iteration: 1, - MaxIter: 10, + Iteration: 1, + + MaxIter: 10, + fileCommonDir: "projects/terra-py-form", + toolLog: []toolLogEntry{ {Name: "read_file", ArgsSnip: "src/main.py", Result: "✓ 0.1s"}, }, } + got2 := buildRichStatus(task2, false, "/home/user/.picoclaw/workspace") + if !strings.Contains(got2, "terra-py-form") { t.Errorf("expected fileCommonDir basename in output, got:\n%s", got2) } // workspace basename fallback with trailing slash + task3 := &activeTask{ Iteration: 1, - MaxIter: 10, + + MaxIter: 10, + toolLog: []toolLogEntry{ {Name: "exec", ArgsSnip: "ls", Result: "✓ 0.1s"}, }, } + for _, ws := range []string{"/home/user/my-project/", "/home/user/my-project"} { got := buildRichStatus(task3, false, ws) + if !strings.Contains(got, "my-project") { t.Errorf("workspace %q: expected 'my-project' in output, got:\n%s", ws, got) } @@ -1862,20 +2563,30 @@ func TestBuildRichStatus_ProjectDir(t *testing.T) { func TestExtractExecProjectDir(t *testing.T) { tests := []struct { name string - cmd string + + cmd string + want string }{ {"cd deep path", "cd /ws/projects/terra-py-form && pytest", "terra-py-form"}, + {"cd direct subdir", "cd /ws/my-app && make build", "my-app"}, + {"cd trailing slash", "cd /ws/my-app/ && ls", "my-app"}, + {"cd to workspace", "cd /ws && ls", "ws"}, + {"no cd prefix", "pytest tests/", ""}, + {"empty command", "", ""}, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { args := map[string]any{"command": tt.cmd} + got := extractExecProjectDir(args) + if got != tt.want { t.Errorf("extractExecProjectDir(%q) = %q, want %q", tt.cmd, got, tt.want) } @@ -1885,20 +2596,29 @@ func TestExtractExecProjectDir(t *testing.T) { func TestFileParentRelDir(t *testing.T) { ws := "/home/user/.picoclaw/workspace" + tests := []struct { name string + path string + want string }{ {"deep path", ws + "/projects/terra/src/main.py", "projects/terra/src"}, + {"direct subdir", ws + "/my-app/README.md", "my-app"}, + {"workspace root file", ws + "/notes.txt", ""}, + {"outside workspace", "/tmp/foo.txt", ""}, + {"trailing slash ws", ws + "/projects/terra/src/main.py", "projects/terra/src"}, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := fileParentRelDir(tt.path, ws) + if got != tt.want { t.Errorf("fileParentRelDir(%q, ws) = %q, want %q", tt.path, got, tt.want) } @@ -1909,18 +2629,26 @@ func TestFileParentRelDir(t *testing.T) { func TestCommonDirPrefix(t *testing.T) { tests := []struct { name string + a, b string + want string }{ {"same dir", "projects/terra/src", "projects/terra/src", "projects/terra/src"}, + {"converge to project", "projects/terra/src", "projects/terra/tests", "projects/terra"}, + {"converge to top", "projects/terra/src", "projects/other/tests", "projects"}, + {"no common", "aaa/bbb", "ccc/ddd", ""}, + {"one is prefix", "projects/terra", "projects/terra/src", "projects/terra"}, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := commonDirPrefix(tt.a, tt.b) + if got != tt.want { t.Errorf("commonDirPrefix(%q, %q) = %q, want %q", tt.a, tt.b, got, tt.want) } @@ -1930,22 +2658,33 @@ func TestCommonDirPrefix(t *testing.T) { func TestDisplayProjectDir(t *testing.T) { // exec projectDir wins + task1 := &activeTask{projectDir: "my-app", fileCommonDir: "projects/other"} + if got := displayProjectDir(task1); got != "my-app" { t.Errorf("expected 'my-app', got %q", got) } + // fileCommonDir fallback: basename + task2 := &activeTask{fileCommonDir: "projects/terra-py-form"} + if got := displayProjectDir(task2); got != "terra-py-form" { t.Errorf("expected 'terra-py-form', got %q", got) } + // single component + task3 := &activeTask{fileCommonDir: "my-app"} + if got := displayProjectDir(task3); got != "my-app" { t.Errorf("expected 'my-app', got %q", got) } + // empty + task4 := &activeTask{} + if got := displayProjectDir(task4); got != "" { t.Errorf("expected empty, got %q", got) } @@ -1953,78 +2692,105 @@ func TestDisplayProjectDir(t *testing.T) { func TestBuildRichStatus_FixedHeight(t *testing.T) { // Test that output has the same number of lines regardless of entry count + countLines := func(s string) int { return strings.Count(s, "\n") } // 0 entries + task0 := &activeTask{Iteration: 1, MaxIter: 10} + lines0 := countLines(buildRichStatus(task0, true, "/ws/p")) // 1 entry + task1 := &activeTask{ Iteration: 1, MaxIter: 10, + toolLog: []toolLogEntry{{Name: "exec", ArgsSnip: "ls", Result: "⏳"}}, } + lines1 := countLines(buildRichStatus(task1, true, "/ws/p")) // 5 entries + task5 := &activeTask{Iteration: 5, MaxIter: 10} + for i := 0; i < 5; i++ { task5.toolLog = append(task5.toolLog, toolLogEntry{ Name: fmt.Sprintf("[%d] exec", i), ArgsSnip: "cmd", Result: "✓ 1.0s", }) } + lines5 := countLines(buildRichStatus(task5, true, "/ws/p")) // 5 entries + sticky error + task5err := &activeTask{Iteration: 5, MaxIter: 10} + for i := 0; i < 5; i++ { task5err.toolLog = append(task5err.toolLog, toolLogEntry{ Name: fmt.Sprintf("[%d] exec", i), ArgsSnip: "cmd", Result: "✓ 1.0s", }) } + errEntry := toolLogEntry{ Name: "[3] exec", ArgsSnip: "pytest", Result: "✗ 2.0s", + ErrDetail: "FAILED test\nExit code: 1", } + task5err.lastError = &errEntry + lines5err := countLines(buildRichStatus(task5err, true, "/ws/p")) if lines0 != lines1 || lines1 != lines5 || lines5 != lines5err { t.Errorf("line counts should be equal: 0=%d, 1=%d, 5=%d, 5+err=%d", + lines0, lines1, lines5, lines5err) } } func TestBuildRichStatus_StickyError(t *testing.T) { // Error from a past entry sticks in the error section + errEntry := toolLogEntry{ Name: "[2] exec", ArgsSnip: "pytest", Result: "✗ 3.2s", + ErrDetail: "FAILED test_login\nExit code: 1", } + task := &activeTask{ Iteration: 5, - MaxIter: 10, + + MaxIter: 10, + toolLog: []toolLogEntry{ {Name: "[3] read_file", ArgsSnip: "src/auth.py", Result: "✓ 0.1s"}, + {Name: "[4] edit_file", ArgsSnip: "src/auth.py", Result: "✓ 0.2s"}, + {Name: "[5] exec", ArgsSnip: "pytest --retry", Result: "⏳"}, }, + lastError: &errEntry, } got := buildRichStatus(task, false, "/ws/p") // Error section should show the sticky error in code block + if !strings.Contains(got, "FAILED test_login") { t.Errorf("expected sticky error detail in error section, got:\n%s", got) } + if !strings.Contains(got, "\u274C") { // ❌ t.Errorf("expected ❌ error header, got:\n%s", got) } // Latest entry is NOT the error + if !strings.Contains(got, "pytest --retry") { t.Errorf("expected latest entry command, got:\n%s", got) } @@ -2032,11 +2798,15 @@ func TestBuildRichStatus_StickyError(t *testing.T) { func TestBuildRichStatus_LatestEntryNoInlineResult(t *testing.T) { longCmd := "uv run pytest tests/hot/test_state_backend_integration.py" + task := &activeTask{ Iteration: 2, - MaxIter: 10, + + MaxIter: 10, + toolLog: []toolLogEntry{ {Name: "exec", ArgsSnip: "ls -la", Result: "\u2713 0.5s"}, + {Name: "exec", ArgsSnip: longCmd, Result: "\u23F3"}, }, } @@ -2044,25 +2814,35 @@ func TestBuildRichStatus_LatestEntryNoInlineResult(t *testing.T) { got := buildRichStatus(task, false, "/ws/my-project") // Latest entry shows command (possibly truncated) with filename visible + if !strings.Contains(got, "integration.py") { t.Errorf("latest entry should show filename, got:\n%s", got) } + // Result on separate indented line + if !strings.Contains(got, " \u23F3") { t.Errorf("latest entry result should be on indented line, got:\n%s", got) } + // Project name shown + if !strings.Contains(got, "my-project") { t.Errorf("should show project name, got:\n%s", got) } + // No second separator before error section + lines := strings.Split(got, "\n") + sepCount := 0 + for _, l := range lines { if strings.HasPrefix(l, "\u2501") { sepCount++ } } + if sepCount != 1 { t.Errorf("expected exactly 1 separator, got %d in:\n%s", sepCount, got) } @@ -2070,35 +2850,49 @@ func TestBuildRichStatus_LatestEntryNoInlineResult(t *testing.T) { func TestSanitizeHistoryForProvider_MultiToolCall(t *testing.T) { // Regression: assistant with 2+ tool_calls had 2nd+ tool results dropped + // because the check only allowed tool after assistant, not after sibling tool. + history := []providers.Message{ {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ {ID: "a", Function: &providers.FunctionCall{Name: "exec"}}, + {ID: "b", Function: &providers.FunctionCall{Name: "read_file"}}, }}, + {Role: "tool", Content: "ok", ToolCallID: "a"}, + {Role: "tool", Content: "ok", ToolCallID: "b"}, + {Role: "assistant", Content: "done"}, } got := sanitizeHistoryForProvider(history) // All 5 messages must survive + if len(got) != 5 { roles := make([]string, len(got)) + for i, m := range got { roles[i] = m.Role } + t.Fatalf("expected 5 messages, got %d: %v", len(got), roles) } + // Verify both tool results present + toolCount := 0 + for _, m := range got { if m.Role == "tool" { toolCount++ } } + if toolCount != 2 { t.Errorf("expected 2 tool results, got %d", toolCount) } @@ -2107,20 +2901,27 @@ func TestSanitizeHistoryForProvider_MultiToolCall(t *testing.T) { // ---------- plan nudge tests ---------- // countingMockProvider counts Chat calls and always returns text-only responses. + type countingMockProvider struct { callCount int } func (m *countingMockProvider) Chat( ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, ) (*providers.LLMResponse, error) { m.callCount++ + return &providers.LLMResponse{ - Content: fmt.Sprintf("Response %d", m.callCount), + Content: fmt.Sprintf("Response %d", m.callCount), + ToolCalls: []providers.ToolCall{}, }, nil } @@ -2134,51 +2935,70 @@ func TestPlanNudge_ForegroundExecution(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } provider := &countingMockProvider{} + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) agent := al.registry.GetDefaultAgent() + if agent == nil { t.Fatal("no default agent") } // Write a plan in executing status with unchecked steps + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n\n## Context\n" + agent.ContextBuilder.WriteMemory(plan) // Process a foreground message (no background metadata) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() msg := bus.InboundMessage{ - Channel: "test", - SenderID: "user1", - ChatID: "chat1", - Content: "continue working", + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "continue working", + SessionKey: "nudge-test", } + _, err = al.processMessage(ctx, msg) if err != nil { t.Fatalf("processMessage failed: %v", err) } // The provider should have been called at least 2 times: + // 1st call: returns text-only → nudge fires (unchecked steps remain) + // 2nd call: returns text-only → nudge already fired, loop exits + if provider.callCount < 2 { t.Errorf("expected at least 2 provider calls (nudge should trigger continuation), got %d", provider.callCount) } @@ -2189,48 +3009,64 @@ func TestPlanNudge_NoNudgeWhenAllStepsComplete(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } provider := &countingMockProvider{} + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) agent := al.registry.GetDefaultAgent() + if agent == nil { t.Fatal("no default agent") } // Write a plan where all steps are already checked + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [x] Step one\n- [x] Step two\n\n## Context\n" + agent.ContextBuilder.WriteMemory(plan) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() msg := bus.InboundMessage{ - Channel: "test", - SenderID: "user1", - ChatID: "chat1", - Content: "all done", + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "all done", + SessionKey: "nudge-test-complete", } + _, err = al.processMessage(ctx, msg) if err != nil { t.Fatalf("processMessage failed: %v", err) } // No unchecked steps → preUnchecked=0 → no nudge → only 1 provider call + if provider.callCount != 1 { t.Errorf("expected exactly 1 provider call (no nudge needed), got %d", provider.callCount) } @@ -2241,98 +3077,135 @@ func TestPlanNudge_ProgressMessage(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } // Provider that checks the nudge message content on the 2nd call + var nudgeContent string + provider := &nudgeCaptureMockProvider{onSecondCall: func(msgs []providers.Message) { // The last user message should be the nudge + for i := len(msgs) - 1; i >= 0; i-- { if msgs[i].Role == "user" { nudgeContent = msgs[i].Content + break } } }} + msgBus := bus.NewMessageBus() + al := NewAgentLoop(cfg, msgBus, provider) agent := al.registry.GetDefaultAgent() + if agent == nil { t.Fatal("no default agent") } // Write a plan with 3 unchecked steps; the provider edits memory to + // mark one step between calls (simulated by the first-call hook). + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n- [ ] Step three\n\n## Context\n" + agent.ContextBuilder.WriteMemory(plan) // After the first LLM response (no tool calls), simulate that + // one step was marked [x] externally (as if the AI did it via tool). + // We do this by hooking the provider's first call to mutate memory. + provider.onFirstCall = func() { updated := strings.Replace(agent.ContextBuilder.ReadMemory(), "- [ ] Step one", "- [x] Step one", 1) + agent.ContextBuilder.WriteMemory(updated) } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() msg := bus.InboundMessage{ - Channel: "test", - SenderID: "user1", - ChatID: "chat1", - Content: "work on the plan", + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "work on the plan", + SessionKey: "nudge-progress-test", } + _, err = al.processMessage(ctx, msg) if err != nil { t.Fatalf("processMessage failed: %v", err) } // Should have gotten the "Progress recorded" nudge (not the "none were marked" one) + if !strings.Contains(nudgeContent, "Progress recorded") { t.Errorf("expected 'Progress recorded' nudge, got %q", nudgeContent) } + if !strings.Contains(nudgeContent, "2 unchecked steps remain") { t.Errorf("expected '2 unchecked steps remain' in nudge, got %q", nudgeContent) } } // nudgeCaptureMockProvider calls hooks on 1st and 2nd Chat invocations. + type nudgeCaptureMockProvider struct { - callCount int - onFirstCall func() + callCount int + + onFirstCall func() + onSecondCall func([]providers.Message) } func (m *nudgeCaptureMockProvider) Chat( ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, ) (*providers.LLMResponse, error) { m.callCount++ + if m.callCount == 1 && m.onFirstCall != nil { m.onFirstCall() } + if m.callCount == 2 && m.onSecondCall != nil { m.onSecondCall(messages) } + return &providers.LLMResponse{ - Content: fmt.Sprintf("Response %d", m.callCount), + Content: fmt.Sprintf("Response %d", m.callCount), + ToolCalls: []providers.ToolCall{}, }, nil } @@ -2345,62 +3218,85 @@ func (m *nudgeCaptureMockProvider) GetDefaultModel() string { func TestConsumeStream_NormalCompletion(t *testing.T) { ch := make(chan protocoltypes.StreamEvent, 8) + go func() { ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} + ch <- protocoltypes.StreamEvent{ContentDelta: "world!"} + ch <- protocoltypes.StreamEvent{ FinishReason: "stop", - Usage: &providers.UsageInfo{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7}, + + Usage: &providers.UsageInfo{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7}, } + close(ch) }() ctx, cancel := context.WithCancel(context.Background()) + defer cancel() resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) if err != nil { t.Fatalf("unexpected error: %v", err) } + if detected { t.Fatal("expected detected=false for normal content") } + if resp.Content != "Hello world!" { t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") } + if resp.FinishReason != "stop" { t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") } + if resp.Usage == nil || resp.Usage.TotalTokens != 7 { t.Errorf("Usage.TotalTokens = %v, want 7", resp.Usage) } + _ = ctx // keep linter happy } func TestConsumeStream_DetectsRepetition(t *testing.T) { ch := make(chan protocoltypes.StreamEvent, 64) + cancelCalled := false ctx, cancel := context.WithCancel(context.Background()) + wrappedCancel := func() { cancelCalled = true + cancel() } // Send enough repetitive content to trigger detection. + // The pattern "abcdefghij" repeated many times will have very low n-gram uniqueness. + repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk + go func() { // Send 6 chunks of repetitive content = 3000 chars total, + // each with 500 runes. The check triggers after every 1000 runes + // when content > 2000 chars. + for i := 0; i < 6; i++ { ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} } + // Send more data that should be ignored after detection. + for i := 0; i < 10; i++ { ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} } + close(ch) }() @@ -2408,53 +3304,68 @@ func TestConsumeStream_DetectsRepetition(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } + if !detected { t.Fatal("expected repetition detection to trigger") } + if !cancelCalled { t.Error("expected cancelFn to be called") } + // The response should be shorter than the full 3000+ chars + // because detection triggers early. + if len(resp.Content) >= 3000+10*len("more data") { t.Errorf("Content length = %d, expected less than full output", len(resp.Content)) } + _ = ctx } func TestConsumeStream_ToolCallAccumulation(t *testing.T) { ch := make(chan protocoltypes.StreamEvent, 8) + go func() { ch <- protocoltypes.StreamEvent{ ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ {Index: 0, ID: "call_1", Name: "test_fn", ArgumentsDelta: `{"ke`}, }, } + ch <- protocoltypes.StreamEvent{ ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ {Index: 0, ArgumentsDelta: `y":"val"}`}, }, } + ch <- protocoltypes.StreamEvent{FinishReason: "tool_calls"} + close(ch) }() _, cancel := context.WithCancel(context.Background()) + defer cancel() resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) if err != nil { t.Fatalf("unexpected error: %v", err) } + if detected { t.Fatal("expected no repetition detection for tool calls") } + if len(resp.ToolCalls) != 1 { t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls)) } + if resp.ToolCalls[0].Name != "test_fn" { t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_fn") } + if resp.ToolCalls[0].Arguments["key"] != "val" { t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "val") } @@ -2462,19 +3373,25 @@ func TestConsumeStream_ToolCallAccumulation(t *testing.T) { func TestConsumeStream_StreamError(t *testing.T) { ch := make(chan protocoltypes.StreamEvent, 4) + go func() { ch <- protocoltypes.StreamEvent{ContentDelta: "partial"} + ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("read error")} + close(ch) }() _, cancel := context.WithCancel(context.Background()) + defer cancel() _, _, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) + if err == nil { t.Fatal("expected error, got nil") } + if !strings.Contains(err.Error(), "read error") { t.Errorf("error = %q, want to contain %q", err.Error(), "read error") } @@ -2482,18 +3399,25 @@ func TestConsumeStream_StreamError(t *testing.T) { func TestConsumeStream_OnChunkCallback(t *testing.T) { ch := make(chan protocoltypes.StreamEvent, 8) + go func() { ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} + ch <- protocoltypes.StreamEvent{ContentDelta: "world"} + ch <- protocoltypes.StreamEvent{ContentDelta: "!"} + ch <- protocoltypes.StreamEvent{FinishReason: "stop"} + close(ch) }() _, cancel := context.WithCancel(context.Background()) + defer cancel() var chunks []string + onChunk := func(accumulated, _ string) { chunks = append(chunks, accumulated) } @@ -2502,22 +3426,29 @@ func TestConsumeStream_OnChunkCallback(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } + if detected { t.Fatal("expected detected=false") } + if resp.Content != "Hello world!" { t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") } + // onChunk should be called once per content delta (3 times) + if len(chunks) != 3 { t.Fatalf("onChunk called %d times, want 3", len(chunks)) } + if chunks[0] != "Hello " { t.Errorf("chunks[0] = %q, want %q", chunks[0], "Hello ") } + if chunks[1] != "Hello world" { t.Errorf("chunks[1] = %q, want %q", chunks[1], "Hello world") } + if chunks[2] != "Hello world!" { t.Errorf("chunks[2] = %q, want %q", chunks[2], "Hello world!") } @@ -2525,26 +3456,33 @@ func TestConsumeStream_OnChunkCallback(t *testing.T) { func TestConsumeStream_OnChunkWithRepetitionDetection(t *testing.T) { ch := make(chan protocoltypes.StreamEvent, 64) + cancelCalled := false ctx, cancel := context.WithCancel(context.Background()) + wrappedCancel := func() { cancelCalled = true + cancel() } repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk + go func() { for i := 0; i < 6; i++ { ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} } + for i := 0; i < 10; i++ { ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} } + close(ch) }() var chunkCount int + onChunk := func(_, _ string) { chunkCount++ } @@ -2553,38 +3491,54 @@ func TestConsumeStream_OnChunkWithRepetitionDetection(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } + if !detected { t.Fatal("expected repetition detection to trigger") } + if !cancelCalled { t.Error("expected cancelFn to be called") } + // onChunk should have been called at least once before detection + if chunkCount == 0 { t.Error("expected onChunk to be called at least once") } + _ = ctx } // modelCapturingMockProvider records which model was passed to Chat. + type modelCapturingMockProvider struct { - mu sync.Mutex - models []string + mu sync.Mutex + + models []string + response string } func (m *modelCapturingMockProvider) Chat( ctx context.Context, + messages []providers.Message, + tools_ []providers.ToolDefinition, + model string, + opts map[string]any, ) (*providers.LLMResponse, error) { m.mu.Lock() + m.models = append(m.models, model) + m.mu.Unlock() + return &providers.LLMResponse{ - Content: m.response, + Content: m.response, + ToolCalls: []providers.ToolCall{}, }, nil } @@ -2598,43 +3552,61 @@ func TestAgentLoop_PlanModel_UsedDuringInterviewing(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "normal-model", - PlanModel: "plan-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "normal-model", + + PlanModel: "plan-model", + + MaxTokens: 4096, + MaxToolIterations: 2, }, }, } msgBus := bus.NewMessageBus() + provider := &modelCapturingMockProvider{response: "Plan interview response"} + al := NewAgentLoop(cfg, msgBus, provider) defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { t.Fatal("No default agent found") } // Write MEMORY.md with interviewing status to activate plan model + memoryDir := filepath.Join(tmpDir, "memory") + os.MkdirAll(memoryDir, 0o755) + memoryPath := filepath.Join(memoryDir, "MEMORY.md") + memoryContent := "# Active Plan\n\n> Task: Test plan model\n> Status: interviewing\n> Phase: 1\n" + if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { t.Fatalf("Failed to write MEMORY.md: %v", wErr) } _, err = al.ProcessDirectWithChannel( + context.Background(), + "Hello, plan model test", + "test-plan-session", + "test", + "test-chat", ) if err != nil { @@ -2642,12 +3614,15 @@ func TestAgentLoop_PlanModel_UsedDuringInterviewing(t *testing.T) { } provider.mu.Lock() + defer provider.mu.Unlock() if len(provider.models) == 0 { t.Fatal("Expected at least one Chat call") } + // The first call should use the plan model since we're in interviewing state + if provider.models[0] != "plan-model" { t.Errorf("Expected plan model 'plan-model' during interviewing, got %q", provider.models[0]) } @@ -2658,51 +3633,77 @@ func TestAgentLoop_PlanModel_NotUsedDuringExecuting(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "normal-model", - PlanModel: "plan-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "normal-model", + + PlanModel: "plan-model", + + MaxTokens: 4096, + MaxToolIterations: 2, }, }, } msgBus := bus.NewMessageBus() + provider := &modelCapturingMockProvider{response: "Executing response"} + al := NewAgentLoop(cfg, msgBus, provider) defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { t.Fatal("No default agent found") } // Write MEMORY.md with executing status - should use normal model + memoryDir := filepath.Join(tmpDir, "memory") + os.MkdirAll(memoryDir, 0o755) + memoryPath := filepath.Join(memoryDir, "MEMORY.md") + memoryContent := `# Active Plan + + > Task: Test plan model + > Status: executing + > Phase: 1 + + ## Phase 1: Build + - [ ] Run build + ` + if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { t.Fatalf("Failed to write MEMORY.md: %v", wErr) } _, err = al.ProcessDirectWithChannel( + context.Background(), + "Hello, executing test", + "test-exec-session", + "test", + "test-chat", ) if err != nil { @@ -2710,12 +3711,15 @@ func TestAgentLoop_PlanModel_NotUsedDuringExecuting(t *testing.T) { } provider.mu.Lock() + defer provider.mu.Unlock() if len(provider.models) == 0 { t.Fatal("Expected at least one Chat call") } + // During executing phase, should use normal model, not plan model + if provider.models[0] != "normal-model" { t.Errorf("Expected normal model 'normal-model' during executing, got %q", provider.models[0]) } @@ -2726,44 +3730,65 @@ func TestAgentLoop_PlanModel_ResolvesProviderForSingleCandidate(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "MiniMax-M2.5", - PlanModel: "openai/gpt-5.2", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "MiniMax-M2.5", + + PlanModel: "openai/gpt-5.2", + + MaxTokens: 4096, + MaxToolIterations: 2, }, }, } msgBus := bus.NewMessageBus() + // The main provider simulates the wrong provider (e.g. MiniMax). + mainProvider := &modelCapturingMockProvider{response: "wrong provider response"} + al := NewAgentLoop(cfg, msgBus, mainProvider) // Inject a mock provider into the cache so resolveProvider returns it + // for the "openai/gpt-5.2" candidate (provider="openai", model="gpt-5.2"). + resolvedProvider := &modelCapturingMockProvider{response: "correct provider response"} + al.providerCache["openai/gpt-5.2"] = resolvedProvider // Write MEMORY.md with interviewing status to activate plan model + memoryDir := filepath.Join(tmpDir, "memory") + os.MkdirAll(memoryDir, 0o755) + memoryPath := filepath.Join(memoryDir, "MEMORY.md") + memoryContent := "# Active Plan\n\n> Task: Test provider resolution\n> Status: interviewing\n> Phase: 1\n" + if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { t.Fatalf("Failed to write MEMORY.md: %v", wErr) } _, err = al.ProcessDirectWithChannel( + context.Background(), + "Hello, resolve provider test", + "test-resolve-session", + "test", + "test-chat", ) if err != nil { @@ -2771,74 +3796,103 @@ func TestAgentLoop_PlanModel_ResolvesProviderForSingleCandidate(t *testing.T) { } resolvedProvider.mu.Lock() + defer resolvedProvider.mu.Unlock() + mainProvider.mu.Lock() + defer mainProvider.mu.Unlock() // The resolved provider should have been called with the stripped model name + if len(resolvedProvider.models) == 0 { t.Fatal("Expected resolved provider to receive Chat call, but it got none") } + if resolvedProvider.models[0] != "gpt-5.2" { t.Errorf("Expected resolved provider to receive model 'gpt-5.2', got %q", resolvedProvider.models[0]) } // The main provider should NOT have been called for the LLM request + if len(mainProvider.models) > 0 { t.Errorf("Expected main provider to receive no Chat calls during plan model phase, got %d calls with models %v", + len(mainProvider.models), mainProvider.models) } } func TestPlanCommand_StartClear(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() // Create a plan in review status with phases + plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + _ = agent.ContextBuilder.WriteMemory(plan) // Seed session history so we can verify it gets cleared + agent.Sessions.AddMessage("test-session", "user", "hello") + agent.Sessions.AddMessage("test-session", "assistant", "world") + agent.Sessions.SetSummary("test-session", "some summary") // Approve with clear + response, handled := al.handleCommand(context.Background(), bus.InboundMessage{ - Content: "/plan start clear", + Content: "/plan start clear", + SessionKey: "test-session", }) + if !handled { t.Fatal("expected /plan start clear to be handled") } + if !strings.Contains(response, "clean history") { t.Errorf("expected 'clean history' in response, got %q", response) } + if !al.planStartPending { t.Error("expected planStartPending to be true") } + if !al.planClearHistory { t.Error("expected planClearHistory to be true") } // Simulate what Run() does when planStartPending is set + al.planStartPending = false + clearHistory := al.planClearHistory + al.planClearHistory = false + if clearHistory { agent.Sessions.SetHistory("test-session", nil) + agent.Sessions.SetSummary("test-session", "") + _ = agent.Sessions.Save("test-session") } // Verify history and summary are cleared + history := agent.Sessions.GetHistory("test-session") + if len(history) != 0 { t.Errorf("expected empty history after clear, got %d messages", len(history)) } + summary := agent.Sessions.GetSummary("test-session") + if summary != "" { t.Errorf("expected empty summary after clear, got %q", summary) } @@ -2846,37 +3900,51 @@ func TestPlanCommand_StartClear(t *testing.T) { func TestPlanCommand_StartWithoutClear_PreservesHistory(t *testing.T) { al, cleanup := newTestAgentLoop(t) + defer cleanup() agent := al.registry.GetDefaultAgent() // Create a plan in review status with phases + plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + _ = agent.ContextBuilder.WriteMemory(plan) // Seed session history + agent.Sessions.AddMessage("test-session", "user", "hello") + agent.Sessions.AddMessage("test-session", "assistant", "world") + agent.Sessions.SetSummary("test-session", "some summary") // Approve without clear + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{ - Content: "/plan start", + Content: "/plan start", + SessionKey: "test-session", }) + if strings.Contains(response, "clean history") { t.Errorf("did not expect 'clean history' in response, got %q", response) } + if al.planClearHistory { t.Error("planClearHistory should be false for /plan start without clear") } // Verify history is preserved + history := agent.Sessions.GetHistory("test-session") + if len(history) != 2 { t.Errorf("expected 2 history messages preserved, got %d", len(history)) } + summary := agent.Sessions.GetSummary("test-session") + if summary != "some summary" { t.Errorf("expected summary preserved, got %q", summary) } @@ -2885,41 +3953,63 @@ func TestPlanCommand_StartWithoutClear_PreservesHistory(t *testing.T) { func TestFilterInterviewTools(t *testing.T) { allDefs := []providers.ToolDefinition{ {Function: protocoltypes.ToolFunctionDefinition{Name: "read_file"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "list_dir"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "web_search"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "web_fetch"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "message"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "edit_file"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "append_file"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "write_file"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "exec"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "logs"}}, + // These should be filtered out: + {Function: protocoltypes.ToolFunctionDefinition{Name: "spawn_subagent"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "skills_search"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "skills_install"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "bg_monitor"}}, + {Function: protocoltypes.ToolFunctionDefinition{Name: "i2c_transfer"}}, } filtered := filterInterviewTools(allDefs) // Should keep exactly the 10 allowed tools + if len(filtered) != 10 { names := make([]string, len(filtered)) + for i, d := range filtered { names[i] = d.Function.Name } + t.Errorf("expected 10 allowed tools, got %d: %v", len(filtered), names) } // Verify none of the disallowed tools slipped through + disallowed := map[string]bool{ "spawnsubagent": true, "skillssearch": true, + "skillsinstall": true, "bgmonitor": true, "ictransfer": true, } + for _, d := range filtered { norm := tools.NormalizeToolName(d.Function.Name) + if disallowed[norm] { t.Errorf("disallowed tool %q should have been filtered out", d.Function.Name) } @@ -2928,13 +4018,17 @@ func TestFilterInterviewTools(t *testing.T) { func TestBuildStreamingDisplay_ContentOnly(t *testing.T) { display := buildStreamingDisplay("hello world", "") + if !strings.HasSuffix(display, " \u2589") { t.Error("expected cursor suffix") } + if strings.Contains(display, "\U0001f9e0") { t.Error("should not contain brain emoji when no reasoning") } + lines := strings.Count(display, "\n") + 1 + if lines != streamingDisplayLines+1 { // TailPad lines + cursor on last line t.Logf("display:\n%s", display) } @@ -2942,12 +4036,15 @@ func TestBuildStreamingDisplay_ContentOnly(t *testing.T) { func TestBuildStreamingDisplay_ReasoningOnly(t *testing.T) { display := buildStreamingDisplay("", "let me think about this") + if !strings.Contains(display, "\U0001f9e0") { t.Error("expected brain emoji for reasoning phase") } + if !strings.Contains(display, "Thinking...") { t.Error("expected Thinking... header") } + if !strings.HasSuffix(display, " \u2589") { t.Error("expected cursor suffix") } @@ -2955,12 +4052,15 @@ func TestBuildStreamingDisplay_ReasoningOnly(t *testing.T) { func TestBuildStreamingDisplay_Both(t *testing.T) { display := buildStreamingDisplay("the answer is 42", "first I considered...") + if !strings.Contains(display, "\U0001f9e0") { t.Error("expected brain emoji") } + if !strings.Contains(display, "responding") { t.Error("expected responding header when both present") } + if !strings.Contains(display, "the answer is 42") { t.Error("expected content in display") } @@ -2969,31 +4069,42 @@ func TestBuildStreamingDisplay_Both(t *testing.T) { func TestHandleReasoning(t *testing.T) { newLoop := func(t *testing.T) (*AgentLoop, *bus.MessageBus) { t.Helper() + tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + t.Cleanup(func() { _ = os.RemoveAll(tmpDir) }) + cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: "test-model", - MaxTokens: 4096, + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + MaxToolIterations: 10, }, }, } + msgBus := bus.NewMessageBus() + return NewAgentLoop(cfg, msgBus, &mockProvider{}), msgBus } t.Run("skips when any required field is empty", func(t *testing.T) { al, msgBus := newLoop(t) + al.handleReasoning(context.Background(), "reasoning", "telegram", "") ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + if msg, ok := msgBus.SubscribeOutbound(ctx); ok { t.Fatalf("expected no outbound message, got %+v", msg) } @@ -3001,14 +4112,19 @@ func TestHandleReasoning(t *testing.T) { t.Run("publishes one message for non telegram", func(t *testing.T) { al, msgBus := newLoop(t) + al.handleReasoning(context.Background(), "hello reasoning", "slack", "channel-1") ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + msg, ok := msgBus.SubscribeOutbound(ctx) + if !ok { t.Fatal("expected an outbound message") } + if msg.Channel != "slack" || msg.ChatID != "channel-1" || msg.Content != "hello reasoning" { t.Fatalf("unexpected outbound message: %+v", msg) } @@ -3016,12 +4132,17 @@ func TestHandleReasoning(t *testing.T) { t.Run("publishes one message for telegram", func(t *testing.T) { al, msgBus := newLoop(t) + reasoning := "hello telegram reasoning" + al.handleReasoning(context.Background(), reasoning, "telegram", "tg-chat") ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + msg, ok := msgBus.SubscribeOutbound(ctx) + if !ok { t.Fatal("expected outbound message") } @@ -3029,23 +4150,33 @@ func TestHandleReasoning(t *testing.T) { if msg.Channel != "telegram" { t.Fatalf("expected telegram channel message, got %+v", msg) } + if msg.ChatID != "tg-chat" { t.Fatalf("expected chatID tg-chat, got %+v", msg) } + if msg.Content != reasoning { t.Fatalf("content mismatch: got %q want %q", msg.Content, reasoning) } }) + t.Run("expired ctx", func(t *testing.T) { al, msgBus := newLoop(t) + reasoning := "hello telegram reasoning" + ctx, cancel := context.WithCancel(context.Background()) + cancel() + al.handleReasoning(ctx, reasoning, "telegram", "tg-chat") ctx, cancel = context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + msg, ok := msgBus.SubscribeOutbound(ctx) + if ok { t.Fatalf("expected no outbound message, got %+v", msg) } @@ -3055,50 +4186,73 @@ func TestHandleReasoning(t *testing.T) { al, msgBus := newLoop(t) // Fill the outbound bus buffer until a publish would block. + // Use a short timeout to detect when the buffer is full, + // rather than hardcoding the buffer size. + for i := 0; ; i++ { fillCtx, fillCancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + err := msgBus.PublishOutbound(fillCtx, bus.OutboundMessage{ Channel: "filler", - ChatID: "filler", + + ChatID: "filler", + Content: fmt.Sprintf("filler-%d", i), }) + fillCancel() + if err != nil { // Buffer is full (timed out trying to send). + break } } // Use a short-deadline parent context to bound the test. + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() start := time.Now() + al.handleReasoning(ctx, "should timeout", "slack", "channel-full") + elapsed := time.Since(start) // handleReasoning uses a 5s internal timeout, but the parent ctx + // expires in 500ms. It should return within ~500ms, not 5s. + if elapsed > 2*time.Second { t.Fatalf("handleReasoning blocked too long (%v); expected prompt return", elapsed) } // Drain the bus and verify the reasoning message was NOT published + // (it should have been dropped due to timeout). + drainCtx, drainCancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer drainCancel() + foundReasoning := false + for { msg, ok := msgBus.SubscribeOutbound(drainCtx) + if !ok { break } + if msg.Content == "should timeout" { foundReasoning = true } } + if foundReasoning { t.Fatal("expected reasoning message to be dropped when bus is full, but it was published") } @@ -3107,25 +4261,39 @@ func TestHandleReasoning(t *testing.T) { func TestFormatDurationMs(t *testing.T) { tests := []struct { - ms int64 + ms int64 + want string }{ {0, "0ms"}, + {500, "500ms"}, + {999, "999ms"}, + {1000, "1.0s"}, + {1200, "1.2s"}, + {3500, "3.5s"}, + {59900, "59.9s"}, + {60000, "1m"}, + {61000, "1m1s"}, + {65000, "1m5s"}, + {120000, "2m"}, + {3661000, "61m1s"}, } + for _, tt := range tests { t.Run(fmt.Sprintf("%dms", tt.ms), func(t *testing.T) { got := formatDurationMs(tt.ms) + if got != tt.want { t.Errorf("formatDurationMs(%d) = %q, want %q", tt.ms, got, tt.want) } @@ -3135,57 +4303,89 @@ func TestFormatDurationMs(t *testing.T) { func TestFormatSubagentCompletion(t *testing.T) { tests := []struct { - name string - label string + name string + + label string + metadata map[string]string - want string + + want string }{ { "no metadata", + "scout-1", + nil, + "📋 scout-1 completed.", }, + { "empty metadata", + "scout-1", + map[string]string{}, + "📋 scout-1 completed.", }, + { "duration and tool calls", + "scout-1", + map[string]string{"duration_ms": "3200", "tool_calls": "5"}, + "📋 scout-1 completed (3.2s, 5 tool calls).", }, + { "single tool call", + "coder-1", + map[string]string{"duration_ms": "1200", "tool_calls": "1"}, + "📋 coder-1 completed (1.2s, 1 tool call).", }, + { "duration only", + "scout-2", + map[string]string{"duration_ms": "65000", "tool_calls": "0"}, + "📋 scout-2 completed (1m5s).", }, + { "tool calls only", + "scout-3", + map[string]string{"duration_ms": "0", "tool_calls": "10"}, + "📋 scout-3 completed (10 tool calls).", }, + { "zero everything", + "scout-4", + map[string]string{"duration_ms": "0", "tool_calls": "0"}, + "📋 scout-4 completed.", }, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := formatSubagentCompletion(tt.label, tt.metadata) + if got != tt.want { t.Errorf("formatSubagentCompletion(%q, %v) = %q, want %q", tt.label, tt.metadata, got, tt.want) } diff --git a/pkg/agent/memory.go b/pkg/agent/memory.go index 34e77eff1..2d81bb398 100644 --- a/pkg/agent/memory.go +++ b/pkg/agent/memory.go @@ -1,7 +1,11 @@ // PicoClaw - Ultra-lightweight personal AI agent + // Inspired by and based on nanobot: https://github.com/HKUDS/nanobot + // License: MIT + // + // Copyright (c) 2026 PicoClaw contributors package agent @@ -20,73 +24,106 @@ import ( ) // MemoryStore manages persistent memory for the agent. + // - Long-term memory: memory/MEMORY.md + // - Daily notes: memory/YYYYMM/YYYYMMDD.md + type MemoryStore struct { - workspace string - memoryDir string + workspace string + + memoryDir string + memoryFile string - cacheMu sync.RWMutex - longTermCache longTermFileCache + cacheMu sync.RWMutex + + longTermCache longTermFileCache + parsedPlanCache parsedPlanStateCache } type longTermFileCache struct { - loaded bool - exists bool + loaded bool + + exists bool + modTime time.Time - size int64 + + size int64 + content string } type parsedPlanStateCache struct { - loaded bool + loaded bool + sourceContent string - state parsedPlanState + + state parsedPlanState } type parsedPlanState struct { - content string + content string + hasActivePlan bool - status string - currentPhase int - totalPhases int - workDir string - taskName string - phases []PlanPhase + + status string + + currentPhase int + + totalPhases int + + workDir string + + taskName string + + phases []PlanPhase } // NewMemoryStore creates a new MemoryStore with the given workspace path. + // It ensures the memory directory exists. + func NewMemoryStore(workspace string) *MemoryStore { memoryDir := filepath.Join(workspace, "memory") + memoryFile := filepath.Join(memoryDir, "MEMORY.md") // Ensure memory directory exists + os.MkdirAll(memoryDir, 0o755) return &MemoryStore{ - workspace: workspace, - memoryDir: memoryDir, + workspace: workspace, + + memoryDir: memoryDir, + memoryFile: memoryFile, } } // getTodayFile returns the path to today's daily note file (memory/YYYYMM/YYYYMMDD.md). + func (ms *MemoryStore) getTodayFile() string { today := time.Now().Format("20060102") // YYYYMMDD - monthDir := today[:6] // YYYYMM + + monthDir := today[:6] // YYYYMM + filePath := filepath.Join(ms.memoryDir, monthDir, today+".md") + return filePath } // InvalidateCache clears all in-memory caches for MEMORY.md content and parsed plan state. + func (ms *MemoryStore) InvalidateCache() { ms.cacheMu.Lock() + defer ms.cacheMu.Unlock() ms.longTermCache = longTermFileCache{} + ms.parsedPlanCache = parsedPlanStateCache{} } @@ -98,56 +135,83 @@ func (ms *MemoryStore) readLongTermCached() string { } ms.cacheMu.RLock() + cachedMissing := ms.longTermCache.loaded && !ms.longTermCache.exists + ms.cacheMu.RUnlock() + if cachedMissing { return "" } ms.cacheMu.Lock() + ms.longTermCache = longTermFileCache{loaded: true, exists: false} + ms.parsedPlanCache = parsedPlanStateCache{} + ms.cacheMu.Unlock() + return "" } modTime := info.ModTime() + size := info.Size() ms.cacheMu.RLock() + if ms.longTermCache.loaded && + ms.longTermCache.exists && + ms.longTermCache.modTime.Equal(modTime) && + ms.longTermCache.size == size { content := ms.longTermCache.content + ms.cacheMu.RUnlock() + return content } + ms.cacheMu.RUnlock() data, err := os.ReadFile(ms.memoryFile) if err != nil { if os.IsNotExist(err) { ms.cacheMu.Lock() + ms.longTermCache = longTermFileCache{loaded: true, exists: false} + ms.parsedPlanCache = parsedPlanStateCache{} + ms.cacheMu.Unlock() } + return "" } + content := string(data) ms.cacheMu.Lock() + ms.longTermCache = longTermFileCache{ - loaded: true, - exists: true, + loaded: true, + + exists: true, + modTime: modTime, - size: size, + + size: size, + content: content, } + if ms.parsedPlanCache.loaded && ms.parsedPlanCache.sourceContent != content { ms.parsedPlanCache = parsedPlanStateCache{} } + ms.cacheMu.Unlock() return content @@ -157,25 +221,33 @@ func (ms *MemoryStore) getParsedPlanState() parsedPlanState { content := ms.ReadLongTerm() ms.cacheMu.RLock() + if ms.parsedPlanCache.loaded && ms.parsedPlanCache.sourceContent == content { state := ms.parsedPlanCache.state + ms.cacheMu.RUnlock() + return state } + ms.cacheMu.RUnlock() state := ms.parsePlanState(content) ms.cacheMu.Lock() + if !ms.parsedPlanCache.loaded || ms.parsedPlanCache.sourceContent != content { ms.parsedPlanCache = parsedPlanStateCache{ - loaded: true, + loaded: true, + sourceContent: content, - state: state, + + state: state, } } else { state = ms.parsedPlanCache.state } + ms.cacheMu.Unlock() return state @@ -183,24 +255,31 @@ func (ms *MemoryStore) getParsedPlanState() parsedPlanState { func (ms *MemoryStore) parsePlanState(content string) parsedPlanState { state := parsedPlanState{content: content} + if content == "" || !reActivePlan.MatchString(content) { return state } state.hasActivePlan = true + if m := reStatus.FindStringSubmatch(content); len(m) >= 2 { state.status = strings.TrimSpace(m[1]) } + if m := rePhase.FindStringSubmatch(content); len(m) >= 2 { state.currentPhase, _ = strconv.Atoi(m[1]) } + state.totalPhases = maxPhaseNumber(content) + if m := reWorkDir.FindStringSubmatch(content); len(m) >= 2 { state.workDir = strings.TrimSpace(m[1]) } + if m := reTaskLine.FindStringSubmatch(content); len(m) >= 2 { state.taskName = strings.TrimSpace(m[1]) } + state.phases = ms.getPlanPhasesFrom(content) return state @@ -212,102 +291,139 @@ func clonePlanPhases(phases []PlanPhase) []PlanPhase { } result := make([]PlanPhase, 0, len(phases)) + for _, p := range phases { phase := PlanPhase{ Number: p.Number, - Title: p.Title, + + Title: p.Title, } + if len(p.Steps) > 0 { phase.Steps = append([]PlanStep(nil), p.Steps...) } + result = append(result, phase) } + return result } // ReadLongTerm reads the long-term memory (MEMORY.md). + // Returns empty string if the file doesn't exist. + func (ms *MemoryStore) ReadLongTerm() string { return ms.readLongTermCached() } // WriteLongTerm writes content to the long-term memory file (MEMORY.md). + func (ms *MemoryStore) WriteLongTerm(content string) error { // Use unified atomic write utility with explicit sync for flash storage reliability. + // Using 0o600 (owner read/write only) for secure default permissions. + if err := fileutil.WriteFileAtomic(ms.memoryFile, []byte(content), 0o600); err != nil { return err } + ms.InvalidateCache() + return nil } // ClearLongTerm removes the long-term memory file. + func (ms *MemoryStore) ClearLongTerm() error { if err := os.Remove(ms.memoryFile); err != nil && !os.IsNotExist(err) { return err } + ms.InvalidateCache() + return nil } // ReadToday reads today's daily note. + // Returns empty string if the file doesn't exist. + func (ms *MemoryStore) ReadToday() string { todayFile := ms.getTodayFile() + if data, err := os.ReadFile(todayFile); err == nil { return string(data) } + return "" } // AppendToday appends content to today's daily note. + // If the file doesn't exist, it creates a new file with a date header. + func (ms *MemoryStore) AppendToday(content string) error { todayFile := ms.getTodayFile() // Ensure month directory exists + monthDir := filepath.Dir(todayFile) + if err := os.MkdirAll(monthDir, 0o755); err != nil { return err } var existingContent string + if data, err := os.ReadFile(todayFile); err == nil { existingContent = string(data) } var newContent string + if existingContent == "" { // Add header for new day + header := fmt.Sprintf("# %s\n\n", time.Now().Format("2006-01-02")) + newContent = header + content } else { // Append to existing content + newContent = existingContent + "\n" + content } // Use unified atomic write utility with explicit sync for flash storage reliability. + return fileutil.WriteFileAtomic(todayFile, []byte(newContent), 0o600) } // GetRecentDailyNotes returns daily notes from the last N days. + // Contents are joined with "---" separator. + func (ms *MemoryStore) GetRecentDailyNotes(days int) string { var sb strings.Builder + first := true for i := range days { date := time.Now().AddDate(0, 0, -i) + dateStr := date.Format("20060102") // YYYYMMDD - monthDir := dateStr[:6] // YYYYMM + + monthDir := dateStr[:6] // YYYYMM + filePath := filepath.Join(ms.memoryDir, monthDir, dateStr+".md") if data, err := os.ReadFile(filePath); err == nil { if !first { sb.WriteString("\n\n---\n\n") } + sb.Write(data) + first = false } } @@ -318,77 +434,100 @@ func (ms *MemoryStore) GetRecentDailyNotes(days int) string { // ---------- Plan state query methods ---------- var ( - reActivePlan = regexp.MustCompile(`(?m)^# Active Plan`) - reStatus = regexp.MustCompile(`(?m)^> Status:\s*(.+)`) - rePhase = regexp.MustCompile(`(?m)^> Phase:\s*(\d+)`) + reActivePlan = regexp.MustCompile(`(?m)^# Active Plan`) + + reStatus = regexp.MustCompile(`(?m)^> Status:\s*(.+)`) + + rePhase = regexp.MustCompile(`(?m)^> Phase:\s*(\d+)`) + rePhaseHeader = regexp.MustCompile(`(?m)^## Phase (\d+):\s*(.*)`) - reWorkDir = regexp.MustCompile(`(?m)^> WorkDir:\s*(.+)`) + + reWorkDir = regexp.MustCompile(`(?m)^> WorkDir:\s*(.+)`) ) // HasActivePlan returns true if MEMORY.md contains an active plan. + func (ms *MemoryStore) HasActivePlan() bool { return ms.getParsedPlanState().hasActivePlan } // GetPlanStatus returns the plan status: "interviewing", "executing", or "". + func (ms *MemoryStore) GetPlanStatus() string { return ms.getParsedPlanState().status } // GetCurrentPhase returns the current phase number from "> Phase: N". + func (ms *MemoryStore) GetCurrentPhase() int { return ms.getParsedPlanState().currentPhase } // GetTotalPhases returns the total number of phases (max ## Phase N). + func (ms *MemoryStore) GetTotalPhases() int { return ms.getParsedPlanState().totalPhases } // IsPlanComplete returns true if all steps in all phases are [x]. + func (ms *MemoryStore) IsPlanComplete() bool { phases := ms.getParsedPlanState().phases + if len(phases) == 0 { return false } + hasSteps := false + for _, p := range phases { for _, s := range p.Steps { hasSteps = true + if !s.Done { return false } } } + return hasSteps } // IsCurrentPhaseComplete returns true if all steps in the current phase are [x]. + func (ms *MemoryStore) IsCurrentPhaseComplete() bool { state := ms.getParsedPlanState() + if state.currentPhase == 0 { return false } + for _, p := range state.phases { if p.Number == state.currentPhase { if len(p.Steps) == 0 { return false } + for _, s := range p.Steps { if !s.Done { return false } } + return true } } + return false } // extractPhaseContent returns the content of a specific phase section. + func (ms *MemoryStore) extractPhaseContent(content string, phase int) string { lines := strings.Split(content, "\n") + inPhase := false + var result []string phasePrefix := fmt.Sprintf("## Phase %d:", phase) @@ -396,34 +535,46 @@ func (ms *MemoryStore) extractPhaseContent(content string, phase int) string { for _, line := range lines { if strings.HasPrefix(line, phasePrefix) { inPhase = true + continue } + if inPhase { // Stop at next phase header or Context section + if strings.HasPrefix(line, "## Phase ") || strings.HasPrefix(line, "## Context") { break } + result = append(result, line) } } + return strings.Join(result, "\n") } // PlanPhase represents a phase with its steps, for structured API output. + type PlanPhase struct { - Number int `json:"number"` - Title string `json:"title"` - Steps []PlanStep `json:"steps"` + Number int `json:"number"` + + Title string `json:"title"` + + Steps []PlanStep `json:"steps"` } // PlanStep represents a single step within a phase. + type PlanStep struct { - Index int `json:"index"` // 1-based within the phase + Index int `json:"index"` // 1-based within the phase + Description string `json:"description"` - Done bool `json:"done"` + + Done bool `json:"done"` } // GetPlanPhases parses MEMORY.md and returns all phases with their steps. + func (ms *MemoryStore) GetPlanPhases() []PlanPhase { return clonePlanPhases(ms.getParsedPlanState().phases) } @@ -434,37 +585,50 @@ func (ms *MemoryStore) getPlanPhasesFrom(content string) []PlanPhase { } totalPhases := maxPhaseNumber(content) + phases := make([]PlanPhase, 0, totalPhases) for p := 1; p <= totalPhases; p++ { title := ms.getPhaseTitle(content, p) + phaseContent := ms.extractPhaseContent(content, p) var steps []PlanStep + stepIdx := 0 + for _, line := range strings.Split(phaseContent, "\n") { line = strings.TrimSpace(line) + if strings.HasPrefix(line, "- [x] ") { stepIdx++ + steps = append(steps, PlanStep{ - Index: stepIdx, + Index: stepIdx, + Description: line[6:], - Done: true, + + Done: true, }) } else if strings.HasPrefix(line, "- [ ] ") { stepIdx++ + steps = append(steps, PlanStep{ - Index: stepIdx, + Index: stepIdx, + Description: line[6:], - Done: false, + + Done: false, }) } } phases = append(phases, PlanPhase{ Number: p, - Title: title, - Steps: steps, + + Title: title, + + Steps: steps, }) } @@ -474,86 +638,122 @@ func (ms *MemoryStore) getPlanPhasesFrom(content string) []PlanPhase { // ---------- Plan mutation methods ---------- // SetStatus sets the plan status (interviewing or executing). + func (ms *MemoryStore) SetStatus(status string) error { content := ms.ReadLongTerm() + if m := reStatus.FindString(content); m != "" { content = strings.Replace(content, m, "> Status: "+status, 1) } + return ms.WriteLongTerm(content) } // AdvancePhase increments the current phase number by 1. + func (ms *MemoryStore) AdvancePhase() error { content := ms.ReadLongTerm() + m := rePhase.FindStringSubmatch(content) + if len(m) < 2 { return fmt.Errorf("no phase marker found") } + current, _ := strconv.Atoi(m[1]) + next := current + 1 + content = strings.Replace(content, m[0], fmt.Sprintf("> Phase: %d", next), 1) + return ms.WriteLongTerm(content) } // SetPhase sets the current phase number to n. + func (ms *MemoryStore) SetPhase(n int) error { content := ms.ReadLongTerm() + m := rePhase.FindString(content) + if m == "" { return fmt.Errorf("no phase marker found") } + content = strings.Replace(content, m, fmt.Sprintf("> Phase: %d", n), 1) + return ms.WriteLongTerm(content) } // MarkStep marks the nth step (1-based) in the given phase as done [x]. + func (ms *MemoryStore) MarkStep(phase, step int) error { content := ms.ReadLongTerm() + lines := strings.Split(content, "\n") + phasePrefix := fmt.Sprintf("## Phase %d:", phase) inPhase := false + stepCount := 0 + for i, line := range lines { if strings.HasPrefix(line, phasePrefix) { inPhase = true + continue } + if inPhase { if strings.HasPrefix(line, "## Phase ") || strings.HasPrefix(line, "## Context") { break } + if strings.HasPrefix(line, "- [ ] ") { stepCount++ + if stepCount == step { lines[i] = strings.Replace(line, "- [ ] ", "- [x] ", 1) + return ms.WriteLongTerm(strings.Join(lines, "\n")) } } } } + return fmt.Errorf("step %d not found in phase %d", step, phase) } // AddStep appends a new step to the given phase. + func (ms *MemoryStore) AddStep(phase int, desc string) error { content := ms.ReadLongTerm() + lines := strings.Split(content, "\n") + phasePrefix := fmt.Sprintf("## Phase %d:", phase) inPhase := false + insertIdx := -1 + for i, line := range lines { if strings.HasPrefix(line, phasePrefix) { inPhase = true + continue } + if inPhase { if strings.HasPrefix(line, "## Phase ") || strings.HasPrefix(line, "## Context") { insertIdx = i + break } + // Track last step line + if strings.HasPrefix(line, "- [") { insertIdx = i + 1 } @@ -562,6 +762,7 @@ func (ms *MemoryStore) AddStep(phase int, desc string) error { if insertIdx < 0 { // Phase not found or empty; append at end + if inPhase { insertIdx = len(lines) } else { @@ -570,40 +771,53 @@ func (ms *MemoryStore) AddStep(phase int, desc string) error { } newStep := "- [ ] " + desc + newLines := make([]string, 0, len(lines)+1) + newLines = append(newLines, lines[:insertIdx]...) + newLines = append(newLines, newStep) + newLines = append(newLines, lines[insertIdx:]...) return ms.WriteLongTerm(strings.Join(newLines, "\n")) } // ValidatePlanStructure checks that the plan has valid structure for + // transitioning out of the interview phase. Returns nil if valid, + // or an error describing the first problem found. + func (ms *MemoryStore) ValidatePlanStructure() error { content := ms.ReadLongTerm() // 1. Header: # Active Plan must exist + if !reActivePlan.MatchString(content) { return fmt.Errorf("missing '# Active Plan' header") } // 2. Required metadata lines + if !reStatus.MatchString(content) { return fmt.Errorf("missing '> Status:' line") } + if !rePhase.MatchString(content) { return fmt.Errorf("missing '> Phase:' line") } // 3. At least one phase header (## Phase N: title) + phases := ms.getPlanPhasesFrom(content) + if len(phases) == 0 { return fmt.Errorf("no '## Phase N:' sections found") } // 4. Every phase must have at least one checkbox step + for _, p := range phases { if len(p.Steps) == 0 { return fmt.Errorf("Phase %d has no checkbox steps (use '- [ ] ...')", p.Number) @@ -616,167 +830,257 @@ func (ms *MemoryStore) ValidatePlanStructure() error { // ---------- Selective injection methods ---------- // GetPlanWorkDir returns the WorkDir from the plan metadata, or "". + func (ms *MemoryStore) GetPlanWorkDir() string { return ms.getParsedPlanState().workDir } // reTaskLine extracts the task name from "> Task: <description>". + var reTaskLine = regexp.MustCompile(`(?m)^> Task:\s*(.+)`) // GetPlanTaskName returns the task description from the plan metadata, or "". + func (ms *MemoryStore) GetPlanTaskName() string { return ms.getParsedPlanState().taskName } // interviewSeed is the initial content written to MEMORY.md when /plan starts. + const interviewSeedTemplate = `# Active Plan + + > Task: %s + > WorkDir: %s + > Status: interviewing + > Phase: 1 + ` // BuildInterviewSeed creates the initial plan seed for a given task description. + func BuildInterviewSeed(task, workDir string) string { return fmt.Sprintf(interviewSeedTemplate, task, workDir) } // GetInterviewContext returns context for injection during the interviewing phase. + // Includes the full seed + interview guide + target format template. + func (ms *MemoryStore) GetInterviewContext() string { return ms.getInterviewContextFrom(ms.ReadLongTerm()) } func (ms *MemoryStore) getInterviewContextFrom(content string) string { var sb strings.Builder + sb.WriteString("## Active Plan (interviewing)\n\n") + sb.WriteString(content) + sb.WriteString("\n\n### Interview Guide\n") + sb.WriteString("Ask about:\n") + sb.WriteString("- Goals and success criteria\n") + sb.WriteString("- Constraints (time, budget, platform)\n") + sb.WriteString("- Environment (OS, language, runtime versions)\n") + sb.WriteString("- Tooling preferences (test framework, linter, formatter, CI)\n") + sb.WriteString("- Key commands the user already runs (build, test, deploy)\n") + sb.WriteString("\n### Rules\n") + sb.WriteString( "- NEVER remove or overwrite the header block (`# Active Plan`, `> Task:`, `> Status:`, `> Phase:` lines). The system parses these to track state.\n", ) + sb.WriteString( + "- After each answer, use edit_file to append findings to the ## Context section of memory/MEMORY.md.\n", ) + sb.WriteString( "- When you have enough information, use edit_file to add ## Phase, ## Commands, and ## Context sections BELOW the header block.\n", ) + sb.WriteString( + "- Each step MUST use checkbox syntax: `- [ ] description`. The system parses checkboxes to track progress.\n", ) + sb.WriteString("- Organize into 2-5 phases with 3-5 steps each.\n") + sb.WriteString( "- After writing Phases, change `> Status: interviewing` to `> Status: review` via edit_file. The user must approve with /plan start before execution begins.\n", ) + sb.WriteString("\n### Target Format (MANDATORY — system parses this exact structure)\n") + sb.WriteString("\n") + sb.WriteString("# Active Plan\n") + sb.WriteString("> Task: <description>\n") + sb.WriteString("> WorkDir: <path>\n") + sb.WriteString("> Status: interviewing\n") + sb.WriteString("> Phase: 1\n") + sb.WriteString("\n") + sb.WriteString("## Phase 1: <title>\n") + sb.WriteString("- [ ] Step description\n") + sb.WriteString("- [ ] Step description\n") + sb.WriteString("\n") + sb.WriteString("## Phase 2: <title>\n") + sb.WriteString("- [ ] Step description\n") + sb.WriteString("- [ ] Step description\n") + sb.WriteString("\n") + sb.WriteString("## Commands\n") + sb.WriteString("build: <project-specific build command>\n") + sb.WriteString("test: <project-specific test command>\n") + sb.WriteString("lint: <project-specific lint command>\n") + sb.WriteString("\n") + sb.WriteString("## Context\n") + sb.WriteString("<collected requirements, decisions, environment>\n") + return sb.String() } // GetReviewContext returns context for injection during the review phase. + // Shows the full plan and instructs the AI to wait for user approval. + func (ms *MemoryStore) GetReviewContext() string { return ms.getReviewContextFrom(ms.ReadLongTerm()) } func (ms *MemoryStore) getReviewContextFrom(content string) string { var sb strings.Builder + sb.WriteString("## Active Plan (awaiting approval)\n\n") + sb.WriteString(content) + sb.WriteString("\n\nThe plan is awaiting user approval.\n") + sb.WriteString("- If the user requests changes, update memory/MEMORY.md via edit_file.\n") + sb.WriteString("- Do NOT change Status yourself. The user will run /plan start to approve.\n") + return sb.String() } // GetPlanContext returns context for injection during the executing phase. + // Only the current phase is shown in detail; completed phases are compressed + // to one-line summaries; future phases are omitted. + func (ms *MemoryStore) GetPlanContext() string { return ms.getPlanContextFrom(ms.ReadLongTerm()) } func (ms *MemoryStore) getPlanContextFrom(content string) string { var currentPhase int + if m := rePhase.FindStringSubmatch(content); len(m) >= 2 { currentPhase, _ = strconv.Atoi(m[1]) } + totalPhases := maxPhaseNumber(content) taskLine := "" + if m := reTaskLine.FindStringSubmatch(content); len(m) >= 2 { taskLine = strings.TrimSpace(m[1]) } var sb strings.Builder + sb.WriteString("## Active Plan\n") + fmt.Fprintf(&sb, "Task: %s | Phase %d/%d\n", taskLine, currentPhase, totalPhases) // Completed phases: one-line summaries + for p := 1; p < currentPhase; p++ { title := ms.getPhaseTitle(content, p) + fmt.Fprintf(&sb, "Done: Phase %d (%s)\n", p, title) } // Current phase: full detail + if currentPhase > 0 { title := ms.getPhaseTitle(content, currentPhase) + fmt.Fprintf(&sb, "### Current: Phase %d — %s\n", currentPhase, title) + phaseContent := ms.extractPhaseContent(content, currentPhase) + sb.WriteString(strings.TrimSpace(phaseContent)) + sb.WriteString("\n") } // Commands section: always included if present + commandsContent := ms.extractCommandsSection(content) + if commandsContent != "" { sb.WriteString("### Commands\n") + sb.WriteString(commandsContent) + sb.WriteString("\n") } // Context section: always included + contextContent := ms.extractContextSection(content) + if contextContent != "" { sb.WriteString("### Context\n") + sb.WriteString(contextContent) + sb.WriteString("\n") } // Orchestration section: conductor's delegation tracking (Delegated/Findings/Decisions) + orchContent := ms.extractSection(content, "Orchestration") + if orchContent != "" { sb.WriteString("### Orchestration\n") + sb.WriteString(orchContent) + sb.WriteString("\n") } @@ -784,80 +1088,107 @@ func (ms *MemoryStore) getPlanContextFrom(content string) string { } // maxPhaseNumber returns the highest phase number found in content. + func maxPhaseNumber(content string) int { matches := rePhaseHeader.FindAllStringSubmatch(content, -1) + maxN := 0 + for _, m := range matches { if len(m) >= 2 { n, _ := strconv.Atoi(m[1]) + if n > maxN { maxN = n } } } + return maxN } // getPhaseTitle extracts the title of a phase from "## Phase N: Title". + func (ms *MemoryStore) getPhaseTitle(content string, phase int) string { matches := rePhaseHeader.FindAllStringSubmatch(content, -1) + for _, m := range matches { if len(m) >= 3 { n, _ := strconv.Atoi(m[1]) + if n == phase { return strings.TrimSpace(m[2]) } } } + return "" } // extractSection extracts a named ## section from the plan content. + // It returns everything between "## <name>" and the next "## " header. + func (ms *MemoryStore) extractSection(content, name string) string { lines := strings.Split(content, "\n") + prefix := "## " + name + inSection := false + var result []string + for _, line := range lines { if strings.HasPrefix(line, prefix) { inSection = true + continue } + if inSection { if strings.HasPrefix(line, "## ") { break } + result = append(result, line) } } + return strings.TrimSpace(strings.Join(result, "\n")) } // extractContextSection extracts the ## Context section from the plan. + func (ms *MemoryStore) extractContextSection(content string) string { return ms.extractSection(content, "Context") } // extractCommandsSection extracts the ## Commands section from the plan. + func (ms *MemoryStore) extractCommandsSection(content string) string { return ms.extractSection(content, "Commands") } // FormatPlanDisplay returns a user-facing display of the full plan with emoji indicators. + func (ms *MemoryStore) FormatPlanDisplay() string { state := ms.getParsedPlanState() + if !state.hasActivePlan { return "No active plan." } var sb strings.Builder + sb.WriteString(fmt.Sprintf("Plan: %s\n", state.taskName)) + sb.WriteString(fmt.Sprintf("Status: %s | Phase %d/%d\n\n", state.status, state.currentPhase, len(state.phases))) for _, p := range state.phases { // Determine phase emoji + var emoji string + if p.Number < state.currentPhase { emoji = "\u2705" // checkmark } else if p.Number == state.currentPhase { @@ -869,6 +1200,7 @@ func (ms *MemoryStore) FormatPlanDisplay() string { sb.WriteString(fmt.Sprintf("%s Phase %d: %s\n", emoji, p.Number, p.Title)) // Show steps for current and completed phases + if p.Number <= state.currentPhase { for _, s := range p.Steps { if s.Done { @@ -881,10 +1213,13 @@ func (ms *MemoryStore) FormatPlanDisplay() string { } commandsContent := ms.extractCommandsSection(state.content) + if commandsContent != "" { sb.WriteString("\nCommands:\n") + for _, line := range strings.Split(commandsContent, "\n") { line = strings.TrimSpace(line) + if line != "" { sb.WriteString(" " + line + "\n") } @@ -892,6 +1227,7 @@ func (ms *MemoryStore) FormatPlanDisplay() string { } contextContent := ms.extractContextSection(state.content) + if contextContent != "" { sb.WriteString("\nContext: " + contextContent + "\n") } @@ -902,24 +1238,35 @@ func (ms *MemoryStore) FormatPlanDisplay() string { // ---------- GetMemoryContext (plan-aware) ---------- // GetMemoryContext returns formatted memory context for the agent prompt. + // When an active plan exists, it uses selective injection based on plan status. + // During interviewing: full seed + interview guide. + // During executing: current phase only with compressed completed phases. + // When plan is active, daily notes injection is suppressed to save context. + func (ms *MemoryStore) GetMemoryContext() string { var parts []string state := ms.getParsedPlanState() + longTerm := state.content if longTerm != "" { if state.hasActivePlan { switch state.status { case "interviewing": + parts = append(parts, ms.getInterviewContextFrom(longTerm)) + case "review": + parts = append(parts, ms.getReviewContextFrom(longTerm)) + default: + parts = append(parts, ms.getPlanContextFrom(longTerm)) } } else { @@ -928,8 +1275,10 @@ func (ms *MemoryStore) GetMemoryContext() string { } // Suppress daily notes when a plan is active to save context + if !state.hasActivePlan { recentNotes := ms.GetRecentDailyNotes(3) + if recentNotes != "" { parts = append(parts, "## Recent Daily Notes\n\n"+recentNotes) } @@ -938,5 +1287,6 @@ func (ms *MemoryStore) GetMemoryContext() string { if len(parts) == 0 { return "" } + return strings.Join(parts, "\n\n---\n\n") } diff --git a/pkg/agent/memory_test.go b/pkg/agent/memory_test.go index 13633e442..e3e8971ce 100644 --- a/pkg/agent/memory_test.go +++ b/pkg/agent/memory_test.go @@ -9,101 +9,172 @@ import ( func newTestMemoryStore(t *testing.T) (*MemoryStore, func()) { t.Helper() + tmpDir, err := os.MkdirTemp("", "memory-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + ms := NewMemoryStore(tmpDir) + return ms, func() { os.RemoveAll(tmpDir) } } const testPlanInterviewing = `# Active Plan + + > Task: Set up server monitoring + > Status: interviewing + > Phase: 1 + ` const testPlanExecuting = `# Active Plan + + > Task: Set up server monitoring + > Status: executing + > Phase: 2 + + ## Phase 1: Prometheus Install + - [x] Install Prometheus + - [x] Configure node_exporter + + ## Phase 2: Grafana Setup + - [ ] Install Grafana + - [ ] Create dashboard + + ## Phase 3: Alert Configuration + - [ ] Set up alert rules + - [ ] Configure Telegram notifications + + ## Commands + build: go build ./... + test: go test ./pkg/... -count=1 + lint: golangci-lint run + + ## Context + Pi: Debian Bookworm arm64, ports: 3000/9090 + ` const testPlanPhase1Complete = `# Active Plan + + > Task: Set up server monitoring + > Status: executing + > Phase: 1 + + ## Phase 1: Prometheus Install + - [x] Install Prometheus + - [x] Configure node_exporter + + ## Phase 2: Grafana Setup + - [ ] Install Grafana + - [ ] Create dashboard + + ## Context + Pi: Debian Bookworm arm64 + ` const testPlanAllComplete = `# Active Plan + + > Task: Set up server monitoring + > Status: executing + > Phase: 2 + + ## Phase 1: Prometheus Install + - [x] Install Prometheus + - [x] Configure node_exporter + + ## Phase 2: Grafana Setup + - [x] Install Grafana + - [x] Create dashboard + + ## Context + Pi: Debian Bookworm arm64 + ` func TestHasActivePlan(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // No plan + if ms.HasActivePlan() { t.Error("expected no active plan for empty memory") } // With regular content + ms.WriteLongTerm("Some random notes") + if ms.HasActivePlan() { t.Error("expected no active plan for regular content") } // With active plan + ms.WriteLongTerm(testPlanExecuting) + if !ms.HasActivePlan() { t.Error("expected active plan to be detected") } @@ -111,21 +182,27 @@ func TestHasActivePlan(t *testing.T) { func TestGetPlanStatus(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // No plan + if status := ms.GetPlanStatus(); status != "" { t.Errorf("expected empty status, got %q", status) } // Interviewing + ms.WriteLongTerm(testPlanInterviewing) + if status := ms.GetPlanStatus(); status != "interviewing" { t.Errorf("expected 'interviewing', got %q", status) } // Executing + ms.WriteLongTerm(testPlanExecuting) + if status := ms.GetPlanStatus(); status != "executing" { t.Errorf("expected 'executing', got %q", status) } @@ -133,15 +210,19 @@ func TestGetPlanStatus(t *testing.T) { func TestGetCurrentPhase(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // No plan + if phase := ms.GetCurrentPhase(); phase != 0 { t.Errorf("expected phase 0, got %d", phase) } // Phase 2 + ms.WriteLongTerm(testPlanExecuting) + if phase := ms.GetCurrentPhase(); phase != 2 { t.Errorf("expected phase 2, got %d", phase) } @@ -149,15 +230,19 @@ func TestGetCurrentPhase(t *testing.T) { func TestGetTotalPhases(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // No plan + if total := ms.GetTotalPhases(); total != 0 { t.Errorf("expected 0 phases, got %d", total) } // 3 phases + ms.WriteLongTerm(testPlanExecuting) + if total := ms.GetTotalPhases(); total != 3 { t.Errorf("expected 3 phases, got %d", total) } @@ -165,22 +250,29 @@ func TestGetTotalPhases(t *testing.T) { func TestIsPlanComplete(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // Not complete + ms.WriteLongTerm(testPlanExecuting) + if ms.IsPlanComplete() { t.Error("expected plan to be incomplete") } // All complete + ms.WriteLongTerm(testPlanAllComplete) + if !ms.IsPlanComplete() { t.Error("expected plan to be complete") } // No plan + ms.ClearLongTerm() + if ms.IsPlanComplete() { t.Error("expected false when no plan exists") } @@ -188,16 +280,21 @@ func TestIsPlanComplete(t *testing.T) { func TestIsCurrentPhaseComplete(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // Phase 2 not complete + ms.WriteLongTerm(testPlanExecuting) + if ms.IsCurrentPhaseComplete() { t.Error("expected current phase to be incomplete") } // Phase 1 complete (current=1) + ms.WriteLongTerm(testPlanPhase1Complete) + if !ms.IsCurrentPhaseComplete() { t.Error("expected phase 1 to be complete") } @@ -205,12 +302,15 @@ func TestIsCurrentPhaseComplete(t *testing.T) { func TestSetStatus(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanInterviewing) + if err := ms.SetStatus("executing"); err != nil { t.Fatalf("SetStatus failed: %v", err) } + if status := ms.GetPlanStatus(); status != "executing" { t.Errorf("expected 'executing', got %q", status) } @@ -218,12 +318,15 @@ func TestSetStatus(t *testing.T) { func TestAdvancePhase(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanPhase1Complete) + if err := ms.AdvancePhase(); err != nil { t.Fatalf("AdvancePhase failed: %v", err) } + if phase := ms.GetCurrentPhase(); phase != 2 { t.Errorf("expected phase 2 after advance, got %d", phase) } @@ -231,37 +334,49 @@ func TestAdvancePhase(t *testing.T) { func TestMarkStep(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanExecuting) // Mark step 1 in phase 2 + if err := ms.MarkStep(2, 1); err != nil { t.Fatalf("MarkStep failed: %v", err) } content := ms.ReadLongTerm() + // Phase 2 should have first step checked + lines := strings.Split(content, "\n") + foundChecked := false + inPhase2 := false + for _, line := range lines { if strings.HasPrefix(line, "## Phase 2:") { inPhase2 = true + continue } + if inPhase2 && strings.HasPrefix(line, "## ") { break } + if inPhase2 && strings.HasPrefix(line, "- [x] Install Grafana") { foundChecked = true } } + if !foundChecked { t.Error("expected 'Install Grafana' to be marked [x]") } // Error case: invalid step + if err := ms.MarkStep(2, 99); err == nil { t.Error("expected error for invalid step number") } @@ -269,23 +384,29 @@ func TestMarkStep(t *testing.T) { func TestAddStep(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanExecuting) // Add step to phase 2 + if err := ms.AddStep(2, "Test dashboard"); err != nil { t.Fatalf("AddStep failed: %v", err) } content := ms.ReadLongTerm() + if !strings.Contains(content, "- [ ] Test dashboard") { t.Error("expected new step to be added") } // Verify it's in the right place (before Phase 3) + idx := strings.Index(content, "- [ ] Test dashboard") + phase3Idx := strings.Index(content, "## Phase 3:") + if idx > phase3Idx { t.Error("expected new step to be before Phase 3") } @@ -293,17 +414,21 @@ func TestAddStep(t *testing.T) { func TestClearLongTerm(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanExecuting) + if err := ms.ClearLongTerm(); err != nil { t.Fatalf("ClearLongTerm failed: %v", err) } + if content := ms.ReadLongTerm(); content != "" { t.Errorf("expected empty memory after clear, got %q", content) } // Clearing again should not error + if err := ms.ClearLongTerm(); err != nil { t.Fatalf("ClearLongTerm (idempotent) failed: %v", err) } @@ -311,34 +436,45 @@ func TestClearLongTerm(t *testing.T) { func TestGetInterviewContext(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanInterviewing) + ctx := ms.GetInterviewContext() if !strings.Contains(ctx, "Active Plan (interviewing)") { t.Error("expected 'Active Plan (interviewing)' header") } + if !strings.Contains(ctx, "Interview Guide") { t.Error("expected 'Interview Guide' section") } + if !strings.Contains(ctx, "Target Format") { t.Error("expected 'Target Format' section") } + if !strings.Contains(ctx, "Set up server monitoring") { t.Error("expected task description in context") } + // Should guide AI to ask about tooling + if !strings.Contains(ctx, "test framework") || !strings.Contains(ctx, "linter") { t.Error("expected interview guide to mention test framework and linter") } + // Target format should include Commands section example + if !strings.Contains(ctx, "## Commands") { t.Error("expected target format to include ## Commands section") } + if !strings.Contains(ctx, "project-specific test command") { t.Error("expected target format Commands to include test command placeholder") } + if !strings.Contains(ctx, "project-specific lint command") { t.Error("expected target format Commands to include lint command placeholder") } @@ -346,46 +482,57 @@ func TestGetInterviewContext(t *testing.T) { func TestGetPlanContext(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanExecuting) + ctx := ms.GetPlanContext() // Should have task summary + if !strings.Contains(ctx, "Phase 2/3") { t.Error("expected 'Phase 2/3' in plan context") } // Completed phase should be summarized + if !strings.Contains(ctx, "Done: Phase 1") { t.Error("expected completed phase summary") } // Current phase should have full detail + if !strings.Contains(ctx, "Current: Phase 2") { t.Error("expected current phase detail") } + if !strings.Contains(ctx, "Install Grafana") { t.Error("expected current phase steps") } // Future phases should NOT appear + if strings.Contains(ctx, "Phase 3") { t.Error("expected future phases to be omitted") } // Commands should be included + if !strings.Contains(ctx, "### Commands") { t.Error("expected Commands section in plan context") } + if !strings.Contains(ctx, "go test") { t.Error("expected test command in Commands section") } + if !strings.Contains(ctx, "golangci-lint") { t.Error("expected lint command in Commands section") } // Context should be included + if !strings.Contains(ctx, "Debian Bookworm") { t.Error("expected Context section") } @@ -393,20 +540,27 @@ func TestGetPlanContext(t *testing.T) { func TestGetMemoryContext_PlanActive_SuppressesDailyNotes(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // Write a daily note + ms.AppendToday("Today's note") // Without plan, daily notes should appear + ctx := ms.GetMemoryContext() + if !strings.Contains(ctx, "Recent Daily Notes") { t.Error("expected daily notes when no plan active") } // With plan, daily notes should be suppressed + ms.WriteLongTerm(testPlanExecuting) + ctx = ms.GetMemoryContext() + if strings.Contains(ctx, "Recent Daily Notes") { t.Error("expected daily notes to be suppressed when plan is active") } @@ -414,14 +568,17 @@ func TestGetMemoryContext_PlanActive_SuppressesDailyNotes(t *testing.T) { func TestGetMemoryContext_InterviewingMode(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanInterviewing) + ctx := ms.GetMemoryContext() if !strings.Contains(ctx, "interviewing") { t.Error("expected interviewing context") } + if !strings.Contains(ctx, "Interview Guide") { t.Error("expected interview guide in context") } @@ -429,14 +586,17 @@ func TestGetMemoryContext_InterviewingMode(t *testing.T) { func TestGetMemoryContext_ExecutingMode(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(testPlanExecuting) + ctx := ms.GetMemoryContext() if !strings.Contains(ctx, "Active Plan") { t.Error("expected active plan in context") } + if !strings.Contains(ctx, "Current: Phase 2") { t.Error("expected current phase in context") } @@ -444,14 +604,17 @@ func TestGetMemoryContext_ExecutingMode(t *testing.T) { func TestGetMemoryContext_RegularMemory(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm("Some notes about projects") + ctx := ms.GetMemoryContext() if !strings.Contains(ctx, "Long-term Memory") { t.Error("expected regular long-term memory section") } + if !strings.Contains(ctx, "Some notes about projects") { t.Error("expected memory content") } @@ -463,12 +626,15 @@ func TestBuildInterviewSeed(t *testing.T) { if !strings.Contains(seed, "# Active Plan") { t.Error("expected '# Active Plan' header") } + if !strings.Contains(seed, "Deploy monitoring stack") { t.Error("expected task description") } + if !strings.Contains(seed, "interviewing") { t.Error("expected interviewing status") } + if !strings.Contains(seed, "> Phase: 1") { t.Error("expected Phase: 1") } @@ -476,28 +642,37 @@ func TestBuildInterviewSeed(t *testing.T) { func TestFormatPlanDisplay(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() // No plan + display := ms.FormatPlanDisplay() + if display != "No active plan." { t.Errorf("expected 'No active plan.', got %q", display) } // With plan + ms.WriteLongTerm(testPlanExecuting) + display = ms.FormatPlanDisplay() if !strings.Contains(display, "Set up server monitoring") { t.Error("expected task name in display") } + if !strings.Contains(display, "Phase 2/3") { t.Error("expected phase count in display") } + // Commands section should be visible + if !strings.Contains(display, "Commands:") { t.Error("expected Commands section in display") } + if !strings.Contains(display, "go test") { t.Error("expected test command in display") } @@ -505,118 +680,211 @@ func TestFormatPlanDisplay(t *testing.T) { func TestValidatePlanStructure(t *testing.T) { tests := []struct { - name string + name string + content string + wantErr string // "" means nil error expected }{ { name: "valid plan with 1 phase and 1 step", + content: `# Active Plan + + > Task: Do something + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [ ] Install deps + `, + wantErr: "", }, + { - name: "missing Active Plan header", + name: "missing Active Plan header", + content: `> Status: executing`, + wantErr: "missing '# Active Plan' header", }, + { name: "missing Status line", + content: `# Active Plan + + > Phase: 1 + + ## Phase 1: Setup + - [ ] Install deps + `, + wantErr: "missing '> Status:' line", }, + { name: "missing Phase line", + content: `# Active Plan + + > Status: executing + + ## Phase 1: Setup + - [ ] Install deps + `, + wantErr: "missing '> Phase:' line", }, + { name: "no Phase sections", + content: `# Active Plan + + > Task: Do something + > Status: executing + > Phase: 1 + `, + wantErr: "no '## Phase N:' sections found", }, + { name: "phase with no checkbox steps", + content: `# Active Plan + + > Task: Do something + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + Some description without checkboxes + `, + wantErr: "Phase 1 has no checkbox steps", }, + { name: "all steps done is valid", + content: `# Active Plan + + > Task: Do something + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [x] Install deps + - [x] Configure + `, + wantErr: "", }, + { name: "multi-phase valid", + content: `# Active Plan + + > Task: Do something + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [ ] Install deps + + ## Phase 2: Build + - [ ] Compile + - [ ] Test + `, + wantErr: "", }, + { name: "second phase empty steps", + content: `# Active Plan + + > Task: Do something + > Status: executing + > Phase: 1 + + ## Phase 1: Setup + - [ ] Install deps + + ## Phase 2: Build + No checkboxes here + `, + wantErr: "Phase 2 has no checkbox steps", }, } @@ -624,9 +892,11 @@ No checkboxes here for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { ms, cleanup := newTestMemoryStore(t) + defer cleanup() ms.WriteLongTerm(tt.content) + err := ms.ValidatePlanStructure() if tt.wantErr == "" { @@ -649,18 +919,23 @@ func TestMemoryStoreCreation(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } + defer os.RemoveAll(tmpDir) ms := NewMemoryStore(tmpDir) // Verify memory directory was created + memoryDir := filepath.Join(tmpDir, "memory") + if _, err := os.Stat(memoryDir); os.IsNotExist(err) { t.Error("expected memory directory to be created") } // Verify memory file path + expectedFile := filepath.Join(memoryDir, "MEMORY.md") + if ms.memoryFile != expectedFile { t.Errorf("expected memory file %q, got %q", expectedFile, ms.memoryFile) } diff --git a/pkg/agent/mock_provider_test.go b/pkg/agent/mock_provider_test.go index 4962810dc..f4042fd01 100644 --- a/pkg/agent/mock_provider_test.go +++ b/pkg/agent/mock_provider_test.go @@ -10,13 +10,18 @@ type mockProvider struct{} func (m *mockProvider) Chat( ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, ) (*providers.LLMResponse, error) { return &providers.LLMResponse{ - Content: "Mock response", + Content: "Mock response", + ToolCalls: []providers.ToolCall{}, }, nil } diff --git a/pkg/agent/registry.go b/pkg/agent/registry.go index 77b846832..d511da80d 100644 --- a/pkg/agent/registry.go +++ b/pkg/agent/registry.go @@ -10,43 +10,62 @@ import ( ) // AgentRegistry manages multiple agent instances and routes messages to them. + type AgentRegistry struct { - agents map[string]*AgentInstance + agents map[string]*AgentInstance + resolver *routing.RouteResolver - mu sync.RWMutex + + mu sync.RWMutex } // NewAgentRegistry creates a registry from config, instantiating all agents. + func NewAgentRegistry( cfg *config.Config, + provider providers.LLMProvider, ) *AgentRegistry { registry := &AgentRegistry{ - agents: make(map[string]*AgentInstance), + agents: make(map[string]*AgentInstance), + resolver: routing.NewRouteResolver(cfg), } agentConfigs := cfg.Agents.List + if len(agentConfigs) == 0 { implicitAgent := &config.AgentConfig{ - ID: "main", + ID: "main", + Default: true, } + instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider) + registry.agents["main"] = instance + logger.InfoCF("agent", "Created implicit main agent (no agents.list configured)", nil) } else { for i := range agentConfigs { ac := &agentConfigs[i] + id := routing.NormalizeAgentID(ac.ID) + instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider) + registry.agents[id] = instance + logger.InfoCF("agent", "Registered agent", + map[string]any{ - "agent_id": id, - "name": ac.Name, + "agent_id": id, + + "name": ac.Name, + "workspace": instance.Workspace, - "model": instance.Model, + + "model": instance.Model, }) } } @@ -55,60 +74,83 @@ func NewAgentRegistry( } // GetAgent returns the agent instance for a given ID. + func (r *AgentRegistry) GetAgent(agentID string) (*AgentInstance, bool) { r.mu.RLock() + defer r.mu.RUnlock() + id := routing.NormalizeAgentID(agentID) + agent, ok := r.agents[id] + return agent, ok } // ResolveRoute determines which agent handles the message. + func (r *AgentRegistry) ResolveRoute(input routing.RouteInput) routing.ResolvedRoute { return r.resolver.ResolveRoute(input) } // ListAgentIDs returns all registered agent IDs. + func (r *AgentRegistry) ListAgentIDs() []string { r.mu.RLock() + defer r.mu.RUnlock() + ids := make([]string, 0, len(r.agents)) + for id := range r.agents { ids = append(ids, id) } + return ids } // CanSpawnSubagent checks if parentAgentID is allowed to spawn targetAgentID. + func (r *AgentRegistry) CanSpawnSubagent(parentAgentID, targetAgentID string) bool { parent, ok := r.GetAgent(parentAgentID) + if !ok { return false } + if parent.Subagents == nil || parent.Subagents.AllowAgents == nil { return false } + targetNorm := routing.NormalizeAgentID(targetAgentID) + for _, allowed := range parent.Subagents.AllowAgents { if allowed == "*" { return true } + if routing.NormalizeAgentID(allowed) == targetNorm { return true } } + return false } // GetDefaultAgent returns the default agent instance. + func (r *AgentRegistry) GetDefaultAgent() *AgentInstance { r.mu.RLock() + defer r.mu.RUnlock() + if agent, ok := r.agents["main"]; ok { return agent } + for _, agent := range r.agents { return agent } + return nil } diff --git a/pkg/agent/registry_test.go b/pkg/agent/registry_test.go index 518bb441f..5a53f92e6 100644 --- a/pkg/agent/registry_test.go +++ b/pkg/agent/registry_test.go @@ -12,9 +12,13 @@ type mockRegistryProvider struct{} func (m *mockRegistryProvider) Chat( ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, ) (*providers.LLMResponse, error) { return &providers.LLMResponse{Content: "mock", FinishReason: "stop"}, nil @@ -28,11 +32,15 @@ func testCfg(agents []config.AgentConfig) *config.Config { return &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: "/tmp/picoclaw-test-registry", - Model: "gpt-4", - MaxTokens: 8192, + Workspace: "/tmp/picoclaw-test-registry", + + Model: "gpt-4", + + MaxTokens: 8192, + MaxToolIterations: 10, }, + List: agents, }, } @@ -40,17 +48,21 @@ func testCfg(agents []config.AgentConfig) *config.Config { func TestNewAgentRegistry_ImplicitMain(t *testing.T) { cfg := testCfg(nil) + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) ids := registry.ListAgentIDs() + if len(ids) != 1 || ids[0] != "main" { t.Errorf("expected implicit main agent, got %v", ids) } agent, ok := registry.GetAgent("main") + if !ok || agent == nil { t.Fatal("expected to find 'main' agent") } + if agent.ID != "main" { t.Errorf("agent.ID = %q, want 'main'", agent.ID) } @@ -59,24 +71,30 @@ func TestNewAgentRegistry_ImplicitMain(t *testing.T) { func TestNewAgentRegistry_ExplicitAgents(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "sales", Default: true, Name: "Sales Bot"}, + {ID: "support", Name: "Support Bot"}, }) + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) ids := registry.ListAgentIDs() + if len(ids) != 2 { t.Fatalf("expected 2 agents, got %d: %v", len(ids), ids) } sales, ok := registry.GetAgent("sales") + if !ok || sales == nil { t.Fatal("expected to find 'sales' agent") } + if sales.Name != "Sales Bot" { t.Errorf("sales.Name = %q, want 'Sales Bot'", sales.Name) } support, ok := registry.GetAgent("support") + if !ok || support == nil { t.Fatal("expected to find 'support' agent") } @@ -86,12 +104,15 @@ func TestAgentRegistry_GetAgent_Normalize(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "my-agent", Default: true}, }) + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, ok := registry.GetAgent("My-Agent") + if !ok || agent == nil { t.Fatal("expected to find agent with normalized ID") } + if agent.ID != "my-agent" { t.Errorf("agent.ID = %q, want 'my-agent'", agent.ID) } @@ -100,12 +121,16 @@ func TestAgentRegistry_GetAgent_Normalize(t *testing.T) { func TestAgentRegistry_GetDefaultAgent(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "alpha"}, + {ID: "beta", Default: true}, }) + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) // GetDefaultAgent first checks for "main", then returns any + agent := registry.GetDefaultAgent() + if agent == nil { t.Fatal("expected a default agent") } @@ -114,27 +139,36 @@ func TestAgentRegistry_GetDefaultAgent(t *testing.T) { func TestAgentRegistry_CanSpawnSubagent(t *testing.T) { cfg := testCfg([]config.AgentConfig{ { - ID: "parent", + ID: "parent", + Default: true, + Subagents: &config.SubagentsConfig{ AllowAgents: []string{"child1", "child2"}, }, }, + {ID: "child1"}, + {ID: "child2"}, + {ID: "restricted"}, }) + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) if !registry.CanSpawnSubagent("parent", "child1") { t.Error("expected parent to be allowed to spawn child1") } + if !registry.CanSpawnSubagent("parent", "child2") { t.Error("expected parent to be allowed to spawn child2") } + if registry.CanSpawnSubagent("parent", "restricted") { t.Error("expected parent to NOT be allowed to spawn restricted") } + if registry.CanSpawnSubagent("child1", "child2") { t.Error("expected child1 to NOT be allowed to spawn (no subagents config)") } @@ -143,19 +177,24 @@ func TestAgentRegistry_CanSpawnSubagent(t *testing.T) { func TestAgentRegistry_CanSpawnSubagent_Wildcard(t *testing.T) { cfg := testCfg([]config.AgentConfig{ { - ID: "admin", + ID: "admin", + Default: true, + Subagents: &config.SubagentsConfig{ AllowAgents: []string{"*"}, }, }, + {ID: "any-agent"}, }) + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) if !registry.CanSpawnSubagent("admin", "any-agent") { t.Error("expected wildcard to allow spawning any agent") } + if !registry.CanSpawnSubagent("admin", "nonexistent") { t.Error("expected wildcard to allow spawning even nonexistent agents") } @@ -163,12 +202,15 @@ func TestAgentRegistry_CanSpawnSubagent_Wildcard(t *testing.T) { func TestAgentInstance_Model(t *testing.T) { model := &config.AgentModelConfig{Primary: "claude-opus"} + cfg := testCfg([]config.AgentConfig{ {ID: "custom", Default: true, Model: model}, }) + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, _ := registry.GetAgent("custom") + if agent.Model != "claude-opus" { t.Errorf("agent.Model = %q, want 'claude-opus'", agent.Model) } @@ -178,10 +220,13 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "inherit", Default: true}, }) + cfg.Agents.Defaults.ModelFallbacks = []string{"openai/gpt-4o-mini", "anthropic/haiku"} + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, _ := registry.GetAgent("inherit") + if len(agent.Fallbacks) != 2 { t.Errorf("expected 2 fallbacks inherited from defaults, got %d", len(agent.Fallbacks)) } @@ -189,16 +234,22 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) { func TestAgentInstance_FallbackExplicitEmpty(t *testing.T) { model := &config.AgentModelConfig{ - Primary: "gpt-4", + Primary: "gpt-4", + Fallbacks: []string{}, // explicitly empty = disable + } + cfg := testCfg([]config.AgentConfig{ {ID: "no-fallback", Default: true, Model: model}, }) + cfg.Agents.Defaults.ModelFallbacks = []string{"should-not-inherit"} + registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, _ := registry.GetAgent("no-fallback") + if len(agent.Fallbacks) != 0 { t.Errorf("expected 0 fallbacks (explicit empty), got %d: %v", len(agent.Fallbacks), agent.Fallbacks) } diff --git a/pkg/agent/session_tracker.go b/pkg/agent/session_tracker.go index b842aa9fa..d76298146 100644 --- a/pkg/agent/session_tracker.go +++ b/pkg/agent/session_tracker.go @@ -8,38 +8,55 @@ import ( ) // SessionEntry represents an active or recently-active session. + type SessionEntry struct { - SessionKey string `json:"session_key"` - Channel string `json:"channel"` - ChatID string `json:"chat_id"` - TouchDir string `json:"touch_dir"` - ProjectPath string `json:"project_path,omitempty"` // canonical project path - Purpose string `json:"purpose,omitempty"` // 1-line task description - Branch string `json:"branch,omitempty"` // git branch name - LastSeenAt time.Time `json:"last_seen_at"` + SessionKey string `json:"session_key"` + + Channel string `json:"channel"` + + ChatID string `json:"chat_id"` + + TouchDir string `json:"touch_dir"` + + ProjectPath string `json:"project_path,omitempty"` // canonical project path + + Purpose string `json:"purpose,omitempty"` // 1-line task description + + Branch string `json:"branch,omitempty"` // git branch name + + LastSeenAt time.Time `json:"last_seen_at"` } // TouchMeta carries optional metadata for Touch calls. + type TouchMeta struct { ProjectPath string // canonical project path (always original workspace-relative) - Purpose string // 1-line task description - Branch string // git branch name + + Purpose string // 1-line task description + + Branch string // git branch name } // PeerInfo is the minimal info shared between sessions on the same project. + type PeerInfo struct { SessionKey string - Purpose string - Branch string + + Purpose string + + Branch string } // SessionTracker tracks per-session tool-call activity. + // Thread-safe; used by AgentLoop for plan coordination and by the mini app API for observability. + type SessionTracker struct { entries sync.Map // sessionKey → *SessionEntry } // NewSessionTracker creates a new tracker. + func NewSessionTracker() *SessionTracker { return &SessionTracker{} } @@ -47,121 +64,177 @@ func NewSessionTracker() *SessionTracker { const sessionActivityTimeout = 15 * time.Minute // Touch records a tool-call activity for a session. + // dir is the workspace-relative directory the tool call targeted. + // If dir is empty, only LastSeenAt is updated. + // meta is optional and carries project coordination metadata. + func (st *SessionTracker) Touch(sessionKey, channel, chatID, dir string, meta *TouchMeta) { now := time.Now() + val, loaded := st.entries.Load(sessionKey) + if loaded { entry := val.(*SessionEntry) + entry.LastSeenAt = now + if dir != "" { entry.TouchDir = dir } + if channel != "" { entry.Channel = channel } + if chatID != "" { entry.ChatID = chatID } + if meta != nil { if meta.ProjectPath != "" { entry.ProjectPath = meta.ProjectPath } + if meta.Purpose != "" { entry.Purpose = meta.Purpose } + if meta.Branch != "" { entry.Branch = meta.Branch } } + return } + entry := &SessionEntry{ SessionKey: sessionKey, - Channel: channel, - ChatID: chatID, - TouchDir: dir, + + Channel: channel, + + ChatID: chatID, + + TouchDir: dir, + LastSeenAt: now, } + if meta != nil { entry.ProjectPath = meta.ProjectPath + entry.Purpose = meta.Purpose + entry.Branch = meta.Branch } + st.entries.Store(sessionKey, entry) } // IsActiveInDir returns true if any session (excluding those matching excludeKey) + // has touched a directory overlapping with dir within sessionActivityTimeout. + // Overlap = either is a prefix of the other (parent/child relationship). + func (st *SessionTracker) IsActiveInDir(dir, excludeKey string) bool { cutoff := time.Now().Add(-sessionActivityTimeout) + active := false + st.entries.Range(func(key, val any) bool { if key.(string) == excludeKey { return true } + entry := val.(*SessionEntry) + if entry.LastSeenAt.After(cutoff) && entry.TouchDir != "" && + (strings.HasPrefix(entry.TouchDir, dir) || strings.HasPrefix(dir, entry.TouchDir)) { active = true + return false } + return true }) + return active } // ListActive returns all sessions seen within sessionActivityTimeout, + // sorted by LastSeenAt descending (most recent first). + func (st *SessionTracker) ListActive() []SessionEntry { cutoff := time.Now().Add(-sessionActivityTimeout) + var result []SessionEntry + st.entries.Range(func(key, val any) bool { entry := val.(*SessionEntry) + if entry.LastSeenAt.After(cutoff) { result = append(result, *entry) // copy } + return true }) + sort.Slice(result, func(i, j int) bool { return result[i].LastSeenAt.After(result[j].LastSeenAt) }) + return result } // GetTouchDir returns the TouchDir for a given session key, or "" if not found. + func (st *SessionTracker) GetTouchDir(sessionKey string) string { val, ok := st.entries.Load(sessionKey) + if !ok { return "" } + return val.(*SessionEntry).TouchDir } // GetPeerPurposes returns purposes of other active sessions targeting the same project. + // Used for lightweight coordination without context pollution. + func (st *SessionTracker) GetPeerPurposes(sessionKey, projectPath string) []PeerInfo { if projectPath == "" { return nil } + cutoff := time.Now().Add(-sessionActivityTimeout) + var result []PeerInfo + st.entries.Range(func(key, val any) bool { if key.(string) == sessionKey { return true } + entry := val.(*SessionEntry) + if entry.LastSeenAt.After(cutoff) && entry.ProjectPath == projectPath { result = append(result, PeerInfo{ SessionKey: entry.SessionKey, - Purpose: entry.Purpose, - Branch: entry.Branch, + + Purpose: entry.Purpose, + + Branch: entry.Branch, }) } + return true }) + return result } diff --git a/pkg/agent/session_tracker_test.go b/pkg/agent/session_tracker_test.go index 61f19064e..f69a7fdd3 100644 --- a/pkg/agent/session_tracker_test.go +++ b/pkg/agent/session_tracker_test.go @@ -9,38 +9,53 @@ func TestTouch(t *testing.T) { st := NewSessionTracker() // Basic touch creates entry + st.Touch("sess1", "telegram", "123", "projects/myapp", nil) + entries := st.ListActive() + if len(entries) != 1 { t.Fatalf("expected 1 entry, got %d", len(entries)) } + if entries[0].SessionKey != "sess1" { t.Errorf("expected session_key=sess1, got %s", entries[0].SessionKey) } + if entries[0].Channel != "telegram" { t.Errorf("expected channel=telegram, got %s", entries[0].Channel) } + if entries[0].TouchDir != "projects/myapp" { t.Errorf("expected touch_dir=projects/myapp, got %s", entries[0].TouchDir) } // Touch again with new dir overwrites TouchDir + st.Touch("sess1", "", "", "projects/other", nil) + entries = st.ListActive() + if len(entries) != 1 { t.Fatalf("expected 1 entry, got %d", len(entries)) } + if entries[0].TouchDir != "projects/other" { t.Errorf("expected touch_dir=projects/other, got %s", entries[0].TouchDir) } + // Channel should remain from first touch + if entries[0].Channel != "telegram" { t.Errorf("expected channel=telegram (unchanged), got %s", entries[0].Channel) } // Touch with empty dir does not overwrite TouchDir + st.Touch("sess1", "", "", "", nil) + entries = st.ListActive() + if entries[0].TouchDir != "projects/other" { t.Errorf("expected touch_dir unchanged, got %s", entries[0].TouchDir) } @@ -50,36 +65,45 @@ func TestIsActiveInDir(t *testing.T) { st := NewSessionTracker() // Setup: sess1 touches "projects/myapp" + st.Touch("sess1", "telegram", "123", "projects/myapp", nil) // Same dir, excluding sess1 → false + if st.IsActiveInDir("projects/myapp", "sess1") { t.Error("expected false when excluding the only active session") } // Same dir, excluding different key → true + if !st.IsActiveInDir("projects/myapp", "heartbeat") { t.Error("expected true for exact dir match") } // Parent dir match: "projects" is prefix of "projects/myapp" + if !st.IsActiveInDir("projects", "heartbeat") { t.Error("expected true for parent dir match") } // Child dir match: "projects/myapp/src" has prefix "projects/myapp" + if !st.IsActiveInDir("projects/myapp/src", "heartbeat") { t.Error("expected true for child dir match") } // Unrelated dir → false + if st.IsActiveInDir("other/stuff", "heartbeat") { t.Error("expected false for unrelated dir") } // Stale entry (manually set LastSeenAt to past) + val, _ := st.entries.Load("sess1") + entry := val.(*SessionEntry) + entry.LastSeenAt = time.Now().Add(-sessionActivityTimeout - time.Minute) if st.IsActiveInDir("projects/myapp", "heartbeat") { @@ -91,32 +115,43 @@ func TestListActive(t *testing.T) { st := NewSessionTracker() // Add two sessions + st.Touch("sess1", "telegram", "123", "projects/a", nil) + time.Sleep(5 * time.Millisecond) // ensure different timestamps + st.Touch("sess2", "discord", "456", "projects/b", nil) entries := st.ListActive() + if len(entries) != 2 { t.Fatalf("expected 2 entries, got %d", len(entries)) } // Most recent first + if entries[0].SessionKey != "sess2" { t.Errorf("expected sess2 first (most recent), got %s", entries[0].SessionKey) } + if entries[1].SessionKey != "sess1" { t.Errorf("expected sess1 second, got %s", entries[1].SessionKey) } // Make sess1 stale + val, _ := st.entries.Load("sess1") + entry := val.(*SessionEntry) + entry.LastSeenAt = time.Now().Add(-sessionActivityTimeout - time.Minute) entries = st.ListActive() + if len(entries) != 1 { t.Fatalf("expected 1 active entry after stale, got %d", len(entries)) } + if entries[0].SessionKey != "sess2" { t.Errorf("expected only sess2, got %s", entries[0].SessionKey) } diff --git a/pkg/session/graph.go b/pkg/session/graph.go index 117a4aa53..6566bb1dc 100644 --- a/pkg/session/graph.go +++ b/pkg/session/graph.go @@ -8,96 +8,136 @@ import ( ) // SessionGraph is a thin wrapper around SessionStore that provides + // structured turn-writing via BeginTurn/TurnWriter. + // It does NOT replace LegacyAdapter — existing call sites remain unchanged. + // Future phases will migrate callers to use SessionGraph directly. + type SessionGraph struct { store SessionStore } // NewSessionGraph creates a SessionGraph backed by the given store. + func NewSessionGraph(store SessionStore) *SessionGraph { return &SessionGraph{store: store} } // Messages returns all messages for the session by reading turns from the store. + func (g *SessionGraph) Messages(sessionKey string) ([]providers.Message, error) { turns, err := g.store.Turns(sessionKey, 0) if err != nil { return nil, err } + var msgs []providers.Message + for _, t := range turns { msgs = append(msgs, t.Messages...) } + if msgs == nil { msgs = []providers.Message{} } + return msgs, nil } // BeginTurn starts a new turn that can be built up incrementally + // and committed atomically. + func (g *SessionGraph) BeginTurn(sessionKey string, kind TurnKind) *TurnWriter { return &TurnWriter{ - store: g.store, + store: g.store, + sessionKey: sessionKey, + turn: Turn{ SessionKey: sessionKey, - Kind: kind, + + Kind: kind, }, } } // TurnWriter accumulates messages for a single turn and commits them atomically. + type TurnWriter struct { - mu sync.Mutex - store SessionStore + mu sync.Mutex + + store SessionStore + sessionKey string - turn Turn - committed bool - discarded bool + + turn Turn + + committed bool + + discarded bool } // Add appends a message to the pending turn. + func (tw *TurnWriter) Add(msg providers.Message) { tw.mu.Lock() + defer tw.mu.Unlock() + tw.turn.Messages = append(tw.turn.Messages, msg) } // SetOrigin sets the origin session key for this turn (e.g. subagent source). + func (tw *TurnWriter) SetOrigin(sessionKey string) { tw.mu.Lock() + defer tw.mu.Unlock() + tw.turn.OriginKey = sessionKey } // SetAuthor sets the author field for this turn. + func (tw *TurnWriter) SetAuthor(author string) { tw.mu.Lock() + defer tw.mu.Unlock() + tw.turn.Author = author } // Commit writes the accumulated turn to the store. + // Returns an error if already committed or discarded. + func (tw *TurnWriter) Commit() error { tw.mu.Lock() + defer tw.mu.Unlock() + if tw.committed { return errors.New("turn already committed") } + if tw.discarded { return errors.New("turn already discarded") } + tw.committed = true + return tw.store.Append(tw.sessionKey, &tw.turn) } // Discard marks the turn as abandoned — nothing is written. + func (tw *TurnWriter) Discard() { tw.mu.Lock() + defer tw.mu.Unlock() + tw.discarded = true } diff --git a/pkg/session/graph_test.go b/pkg/session/graph_test.go index de29616f2..c6aa51ae4 100644 --- a/pkg/session/graph_test.go +++ b/pkg/session/graph_test.go @@ -8,19 +8,25 @@ import ( func TestSessionGraph_Messages(t *testing.T) { store := newTestStore(t) + if err := store.Create("g1", nil); err != nil { t.Fatal(err) } + if err := store.Append("g1", &Turn{ - Kind: TurnNormal, + Kind: TurnNormal, + Messages: []providers.Message{{Role: "user", Content: "hello"}}, }); err != nil { t.Fatal(err) } + if err := store.Append("g1", &Turn{ Kind: TurnNormal, + Messages: []providers.Message{ {Role: "assistant", Content: "hi"}, + {Role: "user", Content: "how are you"}, }, }); err != nil { @@ -28,13 +34,16 @@ func TestSessionGraph_Messages(t *testing.T) { } g := NewSessionGraph(store) + msgs, err := g.Messages("g1") if err != nil { t.Fatal(err) } + if len(msgs) != 3 { t.Fatalf("expected 3 messages, got %d", len(msgs)) } + if msgs[0].Content != "hello" || msgs[1].Content != "hi" || msgs[2].Content != "how are you" { t.Errorf("unexpected messages: %+v", msgs) } @@ -42,14 +51,18 @@ func TestSessionGraph_Messages(t *testing.T) { func TestSessionGraph_Messages_Empty(t *testing.T) { store := newTestStore(t) + if err := store.Create("empty", nil); err != nil { t.Fatal(err) } + g := NewSessionGraph(store) + msgs, err := g.Messages("empty") if err != nil { t.Fatal(err) } + if msgs == nil || len(msgs) != 0 { t.Errorf("expected empty slice, got %v", msgs) } @@ -57,15 +70,21 @@ func TestSessionGraph_Messages_Empty(t *testing.T) { func TestTurnWriter_Commit(t *testing.T) { store := newTestStore(t) + if err := store.Create("tw1", nil); err != nil { t.Fatal(err) } g := NewSessionGraph(store) + tw := g.BeginTurn("tw1", TurnNormal) + tw.Add(providers.Message{Role: "user", Content: "msg1"}) + tw.Add(providers.Message{Role: "assistant", Content: "msg2"}) + tw.SetOrigin("parent-key") + tw.SetAuthor("agent-1") if err := tw.Commit(); err != nil { @@ -76,15 +95,19 @@ func TestTurnWriter_Commit(t *testing.T) { if err != nil { t.Fatal(err) } + if len(turns) != 1 { t.Fatalf("expected 1 turn, got %d", len(turns)) } + if len(turns[0].Messages) != 2 { t.Fatalf("expected 2 messages, got %d", len(turns[0].Messages)) } + if turns[0].OriginKey != "parent-key" { t.Errorf("expected origin 'parent-key', got %q", turns[0].OriginKey) } + if turns[0].Author != "agent-1" { t.Errorf("expected author 'agent-1', got %q", turns[0].Author) } @@ -92,19 +115,24 @@ func TestTurnWriter_Commit(t *testing.T) { func TestTurnWriter_Discard(t *testing.T) { store := newTestStore(t) + if err := store.Create("tw2", nil); err != nil { t.Fatal(err) } g := NewSessionGraph(store) + tw := g.BeginTurn("tw2", TurnNormal) + tw.Add(providers.Message{Role: "user", Content: "should not persist"}) + tw.Discard() turns, err := store.Turns("tw2", 0) if err != nil { t.Fatal(err) } + if len(turns) != 0 { t.Errorf("expected 0 turns after discard, got %d", len(turns)) } @@ -112,17 +140,21 @@ func TestTurnWriter_Discard(t *testing.T) { func TestTurnWriter_DoubleCommit(t *testing.T) { store := newTestStore(t) + if err := store.Create("tw3", nil); err != nil { t.Fatal(err) } g := NewSessionGraph(store) + tw := g.BeginTurn("tw3", TurnNormal) + tw.Add(providers.Message{Role: "user", Content: "once"}) if err := tw.Commit(); err != nil { t.Fatal(err) } + if err := tw.Commit(); err == nil { t.Error("expected error on double commit") } @@ -130,13 +162,17 @@ func TestTurnWriter_DoubleCommit(t *testing.T) { func TestTurnWriter_CommitAfterDiscard(t *testing.T) { store := newTestStore(t) + if err := store.Create("tw4", nil); err != nil { t.Fatal(err) } g := NewSessionGraph(store) + tw := g.BeginTurn("tw4", TurnNormal) + tw.Add(providers.Message{Role: "user", Content: "x"}) + tw.Discard() if err := tw.Commit(); err == nil { diff --git a/pkg/session/legacy_adapter.go b/pkg/session/legacy_adapter.go index 79e72924c..29e7b4719 100644 --- a/pkg/session/legacy_adapter.go +++ b/pkg/session/legacy_adapter.go @@ -460,94 +460,137 @@ func (la *LegacyAdapter) Save(key string) error { } // DefaultPruneTTL is the default time-to-live for session pruning. + const DefaultPruneTTL = 7 * 24 * time.Hour // CompactOldTurns flushes pending writes, then compacts SQLite turns + // keeping only the last keepLast messages. Sets session summary to the given value. + func (la *LegacyAdapter) CompactOldTurns(key string, keepLast int, summary string) error { // 1. Flush pending messages to SQLite + if err := la.Save(key); err != nil { return err } + // 2. Query all turns + turns, err := la.store.Turns(key, 0) if err != nil { return err } + // 3. Count total messages, find cut point + totalMsgs := 0 + for _, t := range turns { totalMsgs += len(t.Messages) } + if keepLast >= totalMsgs { // Nothing to compact, just update summary + if err := la.store.SetSummary(key, summary); err != nil { return err } + la.mu.Lock() + if c, ok := la.cache[key]; ok { c.summary = summary } + la.mu.Unlock() + return nil } + dropCount := totalMsgs - keepLast + accumulated := 0 + cutSeq := 0 + for _, t := range turns { accumulated += len(t.Messages) + if accumulated <= dropCount { cutSeq = t.Seq } else { break } } + if cutSeq == 0 { if err := la.store.SetSummary(key, summary); err != nil { return err } + la.mu.Lock() + if c, ok := la.cache[key]; ok { c.summary = summary } + la.mu.Unlock() + return nil } + // 4. Compact in SQLite + if err := la.store.Compact(key, cutSeq, summary); err != nil { return err } + // 5. Update in-memory cache + la.mu.Lock() + defer la.mu.Unlock() + if c, ok := la.cache[key]; ok { if keepLast < len(c.messages) { c.messages = c.messages[len(c.messages)-keepLast:] } + c.stored = len(c.messages) + c.replaced = false + c.dirty = false + c.summary = summary } + return nil } // Store returns the underlying SessionStore for direct DAG operations. + func (la *LegacyAdapter) Store() SessionStore { return la.store } // Graph returns a SessionGraph backed by the underlying store. + func (la *LegacyAdapter) Graph() *SessionGraph { return NewSessionGraph(la.store) } // AdvanceStored increments the stored counter for a session by delta, + // preventing the flush loop from re-persisting messages already written + // directly to the store (e.g. TurnReport). + func (la *LegacyAdapter) AdvanceStored(key string, delta int) { la.mu.Lock() + defer la.mu.Unlock() + if c, ok := la.cache[key]; ok { c.stored += delta } @@ -573,16 +616,25 @@ func (la *LegacyAdapter) Close() { func (la *LegacyAdapter) flushLoop() { flushTicker := time.NewTicker(5 * time.Minute) + pruneTicker := time.NewTicker(6 * time.Hour) + defer flushTicker.Stop() + defer pruneTicker.Stop() + for { select { case <-flushTicker.C: + la.FlushDirty() + case <-pruneTicker.C: + _, _ = la.store.Prune(DefaultPruneTTL) + case <-la.done: + return } } diff --git a/pkg/session/legacy_adapter_test.go b/pkg/session/legacy_adapter_test.go index 4e48c0044..b313140cb 100644 --- a/pkg/session/legacy_adapter_test.go +++ b/pkg/session/legacy_adapter_test.go @@ -126,7 +126,11 @@ func TestBackend_AddFullMessage(t *testing.T) { Content: "sure", ToolCalls: []providers.ToolCall{ - {ID: "call_1", Type: "function", Function: &providers.FunctionCall{Name: "exec", Arguments: map[string]any{}}}, + { + ID: "call_1", + Type: "function", + Function: &providers.FunctionCall{Name: "exec", Arguments: map[string]any{}}, + }, }, }) @@ -485,49 +489,72 @@ func TestBackend_IncrementalSave(t *testing.T) { func TestCompactOldTurns(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "test.db") + store, err := OpenSQLiteStore(dbPath) if err != nil { t.Fatal(err) } + la := NewLegacyAdapter(store) + defer la.Close() la.GetOrCreate("k1") + // Turn 1: 2 messages + la.AddMessage("k1", "user", "a") + la.AddMessage("k1", "assistant", "b") + la.Save("k1") + // Turn 2: 3 messages + la.AddMessage("k1", "user", "c") + la.AddMessage("k1", "assistant", "d") + la.AddMessage("k1", "user", "e") + la.Save("k1") + // Turn 3: 2 messages + la.AddMessage("k1", "user", "f") + la.AddMessage("k1", "assistant", "g") + la.Save("k1") // Total: 7 messages across 3 turns. keepLast=2 → drop 5 → compact turns 1+2 (5 msgs) + if err := la.CompactOldTurns("k1", 2, "test summary"); err != nil { t.Fatalf("CompactOldTurns: %v", err) } h := la.GetHistory("k1") + if len(h) != 2 { t.Fatalf("expected 2 messages in cache, got %d", len(h)) } + if h[0].Content != "f" || h[1].Content != "g" { t.Errorf("unexpected messages: %+v", h) } + if s := la.GetSummary("k1"); s != "test summary" { t.Errorf("expected summary 'test summary', got %q", s) } // Verify in SQLite: only turn 3 remains + turns, _ := store.Turns("k1", 0) + if len(turns) != 1 { t.Fatalf("expected 1 turn in SQLite, got %d", len(turns)) } + if len(turns[0].Messages) != 2 { t.Errorf("expected 2 messages in remaining turn, got %d", len(turns[0].Messages)) } @@ -535,27 +562,36 @@ func TestCompactOldTurns(t *testing.T) { func TestCompactOldTurns_NothingToCompact(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "test.db") + store, err := OpenSQLiteStore(dbPath) if err != nil { t.Fatal(err) } + la := NewLegacyAdapter(store) + defer la.Close() la.GetOrCreate("k1") + la.AddMessage("k1", "user", "a") + la.AddMessage("k1", "assistant", "b") + la.Save("k1") // keepLast=10 >= total 2 → nothing compacted, summary still updated + if err := la.CompactOldTurns("k1", 10, "new summary"); err != nil { t.Fatalf("CompactOldTurns: %v", err) } h := la.GetHistory("k1") + if len(h) != 2 { t.Fatalf("expected 2 messages, got %d", len(h)) } + if s := la.GetSummary("k1"); s != "new summary" { t.Errorf("expected 'new summary', got %q", s) } @@ -563,29 +599,40 @@ func TestCompactOldTurns_NothingToCompact(t *testing.T) { func TestCompactOldTurns_SingleTurn(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "test.db") + store, err := OpenSQLiteStore(dbPath) if err != nil { t.Fatal(err) } + la := NewLegacyAdapter(store) + defer la.Close() la.GetOrCreate("k1") + la.AddMessage("k1", "user", "a") + la.AddMessage("k1", "assistant", "b") + la.AddMessage("k1", "user", "c") + la.Save("k1") // Single turn with 3 messages, keepLast=2 → dropCount=1, but first turn has 3 msgs + // accumulated(3) > dropCount(1) on first turn → cutSeq=0 → no compaction + if err := la.CompactOldTurns("k1", 2, "sum"); err != nil { t.Fatalf("CompactOldTurns: %v", err) } h := la.GetHistory("k1") + if len(h) != 3 { t.Fatalf("expected 3 messages (no compaction), got %d", len(h)) } + if s := la.GetSummary("k1"); s != "sum" { t.Errorf("expected 'sum', got %q", s) } @@ -593,22 +640,29 @@ func TestCompactOldTurns_SingleTurn(t *testing.T) { func TestCompactOldTurns_Graph(t *testing.T) { dbPath := filepath.Join(t.TempDir(), "test.db") + store, err := OpenSQLiteStore(dbPath) if err != nil { t.Fatal(err) } + la := NewLegacyAdapter(store) + defer la.Close() la.GetOrCreate("k1") + la.AddMessage("k1", "user", "hello") + la.Save("k1") g := la.Graph() + msgs, err := g.Messages("k1") if err != nil { t.Fatal(err) } + if len(msgs) != 1 || msgs[0].Content != "hello" { t.Errorf("unexpected graph messages: %+v", msgs) } diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 2d6578895..0b6b7b3b6 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -12,56 +12,76 @@ import ( ) type Session struct { - Key string `json:"key"` + Key string `json:"key"` + Messages []providers.Message `json:"messages"` - Summary string `json:"summary,omitempty"` - Created time.Time `json:"created"` - Updated time.Time `json:"updated"` + + Summary string `json:"summary,omitempty"` + + Created time.Time `json:"created"` + + Updated time.Time `json:"updated"` } type SessionManager struct { sessions map[string]*Session - mu sync.RWMutex - storage string + + mu sync.RWMutex + + storage string // Write-behind: dirty keys are flushed periodically to reduce disk writes. - dirtyMu sync.Mutex + + dirtyMu sync.Mutex + dirtyKeys map[string]bool - done chan struct{} + + done chan struct{} } func NewSessionManager(storage string) *SessionManager { sm := &SessionManager{ - sessions: make(map[string]*Session), - storage: storage, + sessions: make(map[string]*Session), + + storage: storage, + dirtyKeys: make(map[string]bool), - done: make(chan struct{}), + + done: make(chan struct{}), } if storage != "" { os.MkdirAll(storage, 0o755) + sm.loadSessions() } go sm.flushLoop() + return sm } func (sm *SessionManager) GetOrCreate(key string) *Session { sm.mu.Lock() + defer sm.mu.Unlock() session, ok := sm.sessions[key] + if ok { return session } session = &Session{ - Key: key, + Key: key, + Messages: []providers.Message{}, - Created: time.Now(), - Updated: time.Now(), + + Created: time.Now(), + + Updated: time.Now(), } + sm.sessions[key] = session return session @@ -69,79 +89,102 @@ func (sm *SessionManager) GetOrCreate(key string) *Session { func (sm *SessionManager) AddMessage(sessionKey, role, content string) { sm.AddFullMessage(sessionKey, providers.Message{ - Role: role, + Role: role, + Content: content, }) } // AddFullMessage adds a complete message with tool calls and tool call ID to the session. + // This is used to save the full conversation flow including tool calls and tool results. + func (sm *SessionManager) AddFullMessage(sessionKey string, msg providers.Message) { sm.mu.Lock() + defer sm.mu.Unlock() session, ok := sm.sessions[sessionKey] + if !ok { session = &Session{ - Key: sessionKey, + Key: sessionKey, + Messages: []providers.Message{}, - Created: time.Now(), + + Created: time.Now(), } + sm.sessions[sessionKey] = session } session.Messages = append(session.Messages, msg) + session.Updated = time.Now() } func (sm *SessionManager) GetHistory(key string) []providers.Message { sm.mu.RLock() + defer sm.mu.RUnlock() session, ok := sm.sessions[key] + if !ok { return []providers.Message{} } history := make([]providers.Message, len(session.Messages)) + copy(history, session.Messages) + return history } func (sm *SessionManager) GetSummary(key string) string { sm.mu.RLock() + defer sm.mu.RUnlock() session, ok := sm.sessions[key] + if !ok { return "" } + return session.Summary } func (sm *SessionManager) SetSummary(key string, summary string) { sm.mu.Lock() + defer sm.mu.Unlock() session, ok := sm.sessions[key] + if ok { session.Summary = summary + session.Updated = time.Now() } } func (sm *SessionManager) TruncateHistory(key string, keepLast int) { sm.mu.Lock() + defer sm.mu.Unlock() session, ok := sm.sessions[key] + if !ok { return } if keepLast <= 0 { session.Messages = []providers.Message{} + session.Updated = time.Now() + return } @@ -150,14 +193,20 @@ func (sm *SessionManager) TruncateHistory(key string, keepLast int) { } session.Messages = session.Messages[len(session.Messages)-keepLast:] + session.Updated = time.Now() } // sanitizeFilename converts a session key into a cross-platform safe filename. + // Session keys use "channel:chatID" (e.g. "telegram:123456") but ':' is the + // volume separator on Windows, so filepath.Base would misinterpret the key. + // We replace it with '_'. The original key is preserved inside the JSON file, + // so loadSessions still maps back to the right in-memory key. + func sanitizeFilename(key string) string { return strings.ReplaceAll(key, ":", "_") } @@ -170,33 +219,47 @@ func (sm *SessionManager) Save(key string) error { filename := sanitizeFilename(key) // filepath.IsLocal rejects empty names, "..", absolute paths, and + // OS-reserved device names (NUL, COM1 … on Windows). + // The extra checks reject "." and any directory separators so that + // the session file is always written directly inside sm.storage. + if filename == "." || !filepath.IsLocal(filename) || strings.ContainsAny(filename, `/\`) { return os.ErrInvalid } // Snapshot under read lock, then perform slow file I/O after unlock. + sm.mu.RLock() + stored, ok := sm.sessions[key] + if !ok { sm.mu.RUnlock() + return nil } snapshot := Session{ - Key: stored.Key, + Key: stored.Key, + Summary: stored.Summary, + Created: stored.Created, + Updated: stored.Updated, } + if len(stored.Messages) > 0 { snapshot.Messages = make([]providers.Message, len(stored.Messages)) + copy(snapshot.Messages, stored.Messages) } else { snapshot.Messages = []providers.Message{} } + sm.mu.RUnlock() data, err := json.MarshalIndent(snapshot, "", " ") @@ -205,13 +268,16 @@ func (sm *SessionManager) Save(key string) error { } sessionPath := filepath.Join(sm.storage, filename+".json") + tmpFile, err := os.CreateTemp(sm.storage, "session-*.tmp") if err != nil { return err } tmpPath := tmpFile.Name() + cleanup := true + defer func() { if cleanup { _ = os.Remove(tmpPath) @@ -220,16 +286,22 @@ func (sm *SessionManager) Save(key string) error { if _, err := tmpFile.Write(data); err != nil { _ = tmpFile.Close() + return err } + if err := tmpFile.Chmod(0o644); err != nil { _ = tmpFile.Close() + return err } + if err := tmpFile.Sync(); err != nil { _ = tmpFile.Close() + return err } + if err := tmpFile.Close(); err != nil { return err } @@ -237,7 +309,9 @@ func (sm *SessionManager) Save(key string) error { if err := os.Rename(tmpPath, sessionPath); err != nil { return err } + cleanup = false + return nil } @@ -257,12 +331,14 @@ func (sm *SessionManager) loadSessions() error { } sessionPath := filepath.Join(sm.storage, file.Name()) + data, err := os.ReadFile(sessionPath) if err != nil { continue } var session Session + if err := json.Unmarshal(data, &session); err != nil { continue } @@ -274,57 +350,84 @@ func (sm *SessionManager) loadSessions() error { } // SanitizeHistory rebuilds session history to ensure valid tool-call ordering. + // LLM APIs require that every assistant message with ToolCalls is immediately + // followed by exactly the matching tool-result messages (role="tool"), with no + // other messages in between. Violations can happen from session collisions or + // mid-execution crashes. + // + // The function walks the full history and copies only well-formed groups: + // - user/system messages are always kept + // - assistant messages without tool calls are always kept + // - assistant messages WITH tool calls are kept only if the immediately + // following messages are the complete set of matching tool results + // + // Returns the sanitized history and the number of messages removed. + func SanitizeHistory(history []providers.Message) ([]providers.Message, int) { if len(history) == 0 { return history, 0 } result := make([]providers.Message, 0, len(history)) + i := 0 for i < len(history) { msg := history[i] // Non-assistant messages or assistant without tool calls: keep + if msg.Role != "assistant" || len(msg.ToolCalls) == 0 { // Skip stray tool results not preceded by their assistant + if msg.Role == "tool" { i++ + continue } + result = append(result, msg) + i++ + continue } // Assistant with tool calls: validate the immediately following messages + expectedIDs := make(map[string]bool, len(msg.ToolCalls)) + for _, tc := range msg.ToolCalls { expectedIDs[tc.ID] = true } + needed := len(expectedIDs) // Peek ahead: the next `needed` messages must all be tool results with matching IDs + groupOK := true + if i+needed >= len(history) { groupOK = false } else { for j := 0; j < needed; j++ { next := history[i+1+j] + if next.Role != "tool" || !expectedIDs[next.ToolCallID] { groupOK = false + break } } @@ -332,14 +435,19 @@ func SanitizeHistory(history []providers.Message) ([]providers.Message, int) { if groupOK { // Copy assistant + all tool results + result = append(result, msg) + for j := 0; j < needed; j++ { result = append(result, history[i+1+j]) } + i += 1 + needed } else { // Skip the broken assistant message; tool results will be skipped + // individually when encountered (the "stray tool result" check above) + i++ } } @@ -348,37 +456,54 @@ func SanitizeHistory(history []providers.Message) ([]providers.Message, int) { } // SetHistory updates the messages of a session. + func (sm *SessionManager) SetHistory(key string, history []providers.Message) { sm.mu.Lock() + defer sm.mu.Unlock() session, ok := sm.sessions[key] + if ok { // Create a deep copy to strictly isolate internal state + // from the caller's slice. + msgs := make([]providers.Message, len(history)) + copy(msgs, history) + session.Messages = msgs + session.Updated = time.Now() } } // MarkDirty marks a session key for deferred persistence. + // The session will be written to disk on the next periodic flush or on Close(). + func (sm *SessionManager) MarkDirty(key string) { sm.dirtyMu.Lock() + sm.dirtyKeys[key] = true + sm.dirtyMu.Unlock() } // FlushDirty writes all dirty sessions to disk. + func (sm *SessionManager) FlushDirty() { sm.dirtyMu.Lock() + keys := make([]string, 0, len(sm.dirtyKeys)) + for k := range sm.dirtyKeys { keys = append(keys, k) } + sm.dirtyKeys = make(map[string]bool) + sm.dirtyMu.Unlock() for _, k := range keys { @@ -387,24 +512,34 @@ func (sm *SessionManager) FlushDirty() { } // Close stops the background flush goroutine and writes all dirty sessions. + func (sm *SessionManager) Close() { select { case <-sm.done: + return // already closed + default: } + close(sm.done) + sm.FlushDirty() } func (sm *SessionManager) flushLoop() { ticker := time.NewTicker(5 * time.Minute) + defer ticker.Stop() + for { select { case <-ticker.C: + sm.FlushDirty() + case <-sm.done: + return } } diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go index 3e7dac9e7..05b150015 100644 --- a/pkg/session/manager_test.go +++ b/pkg/session/manager_test.go @@ -10,20 +10,27 @@ import ( func TestSanitizeFilename(t *testing.T) { tests := []struct { - input string + input string + expected string }{ {"simple", "simple"}, + {"telegram:123456", "telegram_123456"}, + {"discord:987654321", "discord_987654321"}, + {"slack:C01234", "slack_C01234"}, + {"no-colons-here", "no-colons-here"}, + {"multiple:colons:here", "multiple_colons_here"}, } for _, tt := range tests { t.Run(tt.input, func(t *testing.T) { got := sanitizeFilename(tt.input) + if got != tt.expected { t.Errorf("sanitizeFilename(%q) = %q, want %q", tt.input, got, tt.expected) } @@ -33,30 +40,41 @@ func TestSanitizeFilename(t *testing.T) { func TestSave_WithColonInKey(t *testing.T) { tmpDir := t.TempDir() + sm := NewSessionManager(tmpDir) // Create a session with a key containing colon (typical channel session key). + key := "telegram:123456" + sm.GetOrCreate(key) + sm.AddMessage(key, "user", "hello") // Save should succeed even though the key contains ':' + if err := sm.Save(key); err != nil { t.Fatalf("Save(%q) failed: %v", key, err) } // The file on disk should use sanitized name. + expectedFile := filepath.Join(tmpDir, "telegram_123456.json") + if _, err := os.Stat(expectedFile); os.IsNotExist(err) { t.Fatalf("expected session file %s to exist", expectedFile) } // Load into a fresh manager and verify the session round-trips. + sm2 := NewSessionManager(tmpDir) + history := sm2.GetHistory(key) + if len(history) != 1 { t.Fatalf("expected 1 message after reload, got %d", len(history)) } + if history[0].Content != "hello" { t.Errorf("expected message content %q, got %q", "hello", history[0].Content) } @@ -65,19 +83,27 @@ func TestSave_WithColonInKey(t *testing.T) { func TestSanitizeHistory_OrphanedToolCall(t *testing.T) { history := []providers.Message{ {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ {ID: "call_1", Name: "exec"}, + {ID: "call_2", Name: "list_dir"}, }}, + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, + // Missing tool result for call_2 → orphaned + } sanitized, removed := SanitizeHistory(history) + if removed == 0 { t.Fatal("expected orphaned messages to be removed") } + // After sanitization, only the user message should remain + if len(sanitized) != 1 || sanitized[0].Role != "user" { t.Errorf("expected [user], got %d messages", len(sanitized)) } @@ -85,25 +111,36 @@ func TestSanitizeHistory_OrphanedToolCall(t *testing.T) { func TestSanitizeHistory_InterleavedMessages(t *testing.T) { // Simulates session collision: a user message got interleaved between + // an assistant tool call and its tool result + history := []providers.Message{ {Role: "user", Content: "first"}, + {Role: "assistant", Content: "ok", ToolCalls: []providers.ToolCall{ {ID: "call_1", Name: "exec"}, }}, - {Role: "user", Content: "collision!"}, // ← interleaved from other session + + {Role: "user", Content: "collision!"}, // ← interleaved from other session + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, // ← out of order + {Role: "assistant", Content: "done"}, } sanitized, removed := SanitizeHistory(history) + if removed == 0 { t.Fatal("expected interleaved messages to be removed") } + // Should keep: user("first"), user("collision!"), assistant("done") + // Should remove: assistant(call_1), tool(call_1) + if len(sanitized) != 3 { t.Errorf("expected 3 messages, got %d", len(sanitized)) + for i, m := range sanitized { t.Logf(" [%d] role=%s content=%q", i, m.Role, m.Content) } @@ -113,17 +150,22 @@ func TestSanitizeHistory_InterleavedMessages(t *testing.T) { func TestSanitizeHistory_CleanHistory(t *testing.T) { history := []providers.Message{ {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ {ID: "call_1", Name: "exec"}, }}, + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, + {Role: "assistant", Content: "done"}, } sanitized, removed := SanitizeHistory(history) + if removed != 0 { t.Errorf("expected 0 removed, got %d", removed) } + if len(sanitized) != 4 { t.Errorf("expected 4 messages, got %d", len(sanitized)) } @@ -132,19 +174,26 @@ func TestSanitizeHistory_CleanHistory(t *testing.T) { func TestSanitizeHistory_MultipleToolCalls(t *testing.T) { history := []providers.Message{ {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ {ID: "call_1", Name: "exec"}, + {ID: "call_2", Name: "read_file"}, }}, + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, + {Role: "tool", Content: "content", ToolCallID: "call_2"}, + {Role: "assistant", Content: "all done"}, } sanitized, removed := SanitizeHistory(history) + if removed != 0 { t.Errorf("expected 0 removed, got %d", removed) } + if len(sanitized) != 5 { t.Errorf("expected 5 messages, got %d", len(sanitized)) } @@ -152,6 +201,7 @@ func TestSanitizeHistory_MultipleToolCalls(t *testing.T) { func TestSanitizeHistory_Empty(t *testing.T) { sanitized, removed := SanitizeHistory(nil) + if removed != 0 || sanitized != nil { t.Errorf("expected nil/0, got %v/%d", sanitized, removed) } @@ -159,11 +209,14 @@ func TestSanitizeHistory_Empty(t *testing.T) { func TestSave_RejectsPathTraversal(t *testing.T) { tmpDir := t.TempDir() + sm := NewSessionManager(tmpDir) badKeys := []string{"", ".", "..", "foo/bar", "foo\\bar"} + for _, key := range badKeys { sm.GetOrCreate(key) + if err := sm.Save(key); err == nil { t.Errorf("Save(%q) should have failed but didn't", key) } diff --git a/pkg/session/sqlite.go b/pkg/session/sqlite.go index 9af6ce909..2a8636214 100644 --- a/pkg/session/sqlite.go +++ b/pkg/session/sqlite.go @@ -14,54 +14,104 @@ const sqliteDriver = "sqlite" const schema = ` + + CREATE TABLE IF NOT EXISTS sessions ( + + key TEXT PRIMARY KEY, + + parent_key TEXT NOT NULL DEFAULT '', + + fork_turn_id TEXT NOT NULL DEFAULT '', + + status TEXT NOT NULL DEFAULT 'active', + + label TEXT NOT NULL DEFAULT '', + + summary TEXT NOT NULL DEFAULT '', + + created_at TEXT NOT NULL, + + updated_at TEXT NOT NULL + + ); + + CREATE TABLE IF NOT EXISTS turns ( + + id TEXT PRIMARY KEY, + + session_key TEXT NOT NULL REFERENCES sessions(key) ON DELETE CASCADE, + + seq INTEGER NOT NULL, + + kind INTEGER NOT NULL DEFAULT 0, + + messages TEXT NOT NULL DEFAULT '[]', + + origin_key TEXT NOT NULL DEFAULT '', + + summary TEXT NOT NULL DEFAULT '', + + author TEXT NOT NULL DEFAULT '', + + created_at TEXT NOT NULL, + + meta TEXT NOT NULL DEFAULT '{}' + + ); + + CREATE INDEX IF NOT EXISTS idx_turns_session_seq ON turns(session_key, seq); + + CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_key); + + ` // SQLiteStore implements SessionStore backed by a single SQLite file. @@ -130,6 +180,8 @@ func (s *SQLiteStore) Create(key string, opts *CreateOpts) error { `INSERT INTO sessions (key, parent_key, fork_turn_id, label, created_at, updated_at) + + VALUES (?, ?, ?, ?, ?, ?)`, key, parentKey, forkTurnID, label, now, now, @@ -143,6 +195,8 @@ func (s *SQLiteStore) Get(key string) (*SessionInfo, error) { `SELECT key, parent_key, fork_turn_id, status, label, summary, created_at, updated_at + + FROM sessions WHERE key = ?`, key, ) @@ -291,6 +345,8 @@ func (s *SQLiteStore) Append(sessionKey string, turn *Turn) error { `INSERT INTO turns (id, session_key, seq, kind, messages, origin_key, summary, author, created_at, meta) + + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, turn.ID, sessionKey, turn.Seq, int(turn.Kind), @@ -315,6 +371,8 @@ func (s *SQLiteStore) Turns(sessionKey string, sinceSeq int) ([]*Turn, error) { `SELECT id, session_key, seq, kind, messages, origin_key, summary, author, created_at, meta + + FROM turns WHERE session_key = ? AND seq > ? ORDER BY seq`, sessionKey, sinceSeq, @@ -379,6 +437,8 @@ func (s *SQLiteStore) LastTurn(sessionKey string) (*Turn, error) { `SELECT id, session_key, seq, kind, messages, origin_key, summary, author, created_at, meta + + FROM turns WHERE session_key = ? ORDER BY seq DESC LIMIT 1`, sessionKey, diff --git a/pkg/tools/bg_monitor.go b/pkg/tools/bg_monitor.go index 4f9bb9dd1..43080adf2 100644 --- a/pkg/tools/bg_monitor.go +++ b/pkg/tools/bg_monitor.go @@ -10,17 +10,21 @@ import ( ) const ( - bgWatchPollInterval = 100 * time.Millisecond + bgWatchPollInterval = 100 * time.Millisecond + bgWatchDefaultTimeout = 30 * time.Second - bgTailDefaultLines = 20 + + bgTailDefaultLines = 20 ) // BgMonitorTool monitors and inspects background processes managed by ExecTool. + type BgMonitorTool struct { exec *ExecTool } // NewBgMonitorTool creates a new BgMonitorTool that accesses bg processes from the given ExecTool. + func NewBgMonitorTool(exec *ExecTool) *BgMonitorTool { return &BgMonitorTool{exec: exec} } @@ -36,95 +40,130 @@ func (t *BgMonitorTool) Description() string { func (t *BgMonitorTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "action": map[string]any{ - "type": "string", - "enum": []string{"list", "watch", "tail"}, + "type": "string", + + "enum": []string{"list", "watch", "tail"}, + "description": "Action: 'list' all bg processes, 'watch' for a pattern in output, 'tail' recent output lines.", }, + "bg_id": map[string]any{ - "type": "string", + "type": "string", + "description": "Background process ID (e.g. 'bg-1'). Required for watch and tail.", }, + "pattern": map[string]any{ - "type": "string", + "type": "string", + "description": "Regex pattern to watch for in output (used with action='watch').", }, + "lines": map[string]any{ - "type": "number", + "type": "number", + "description": "Number of recent lines to return (used with action='tail', default 20).", }, + "watch_timeout": map[string]any{ - "type": "number", + "type": "number", + "description": "Timeout in seconds for watch action (default 30).", }, }, + "required": []string{"action"}, } } func (t *BgMonitorTool) Execute(ctx context.Context, args map[string]any) *ToolResult { action, _ := args["action"].(string) + switch action { case "list": + return t.actionList() + case "watch": + return t.actionWatch(ctx, args) + case "tail": + return t.actionTail(args) + default: + return ErrorResult(fmt.Sprintf("unknown action %q (use 'list', 'watch', or 'tail')", action)) } } func (t *BgMonitorTool) actionList() *ToolResult { procs := t.exec.BgProcesses() + if len(procs) == 0 { return &ToolResult{ - ForLLM: "No background processes.", + ForLLM: "No background processes.", + ForUser: "No background processes.", } } ids := make([]string, 0, len(procs)) + for id := range procs { ids = append(ids, id) } + sort.Strings(ids) var sb strings.Builder + sb.WriteString("Background Processes:\n\n") + for _, id := range ids { bp := procs[id] + if bp.isRunning() { uptime := time.Since(bp.startedAt).Truncate(time.Second) + fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n", + id, bp.pid, uptime, getBgMaxLifetime(), bp.command) } else { ran := time.Since(bp.startedAt).Truncate(time.Second) + if bp.exitErr != nil { fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n", + id, bp.pid, ran, bp.command) } else { fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n", + id, bp.pid, ran, bp.command) } } } return &ToolResult{ - ForLLM: sb.String(), + ForLLM: sb.String(), + ForUser: sb.String(), } } func (t *BgMonitorTool) actionWatch(ctx context.Context, args map[string]any) *ToolResult { bgID, _ := args["bg_id"].(string) + if bgID == "" { return ErrorResult("bg_id is required for watch action") } patternStr, _ := args["pattern"].(string) + if patternStr == "" { return ErrorResult("pattern is required for watch action") } @@ -135,64 +174,93 @@ func (t *BgMonitorTool) actionWatch(ctx context.Context, args map[string]any) *T } timeout := bgWatchDefaultTimeout + if t, ok := args["watch_timeout"].(float64); ok && t > 0 { timeout = time.Duration(t) * time.Second } procs := t.exec.BgProcesses() + bp, ok := procs[bgID] + if !ok { return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) } deadline := time.After(timeout) + ticker := time.NewTicker(bgWatchPollInterval) + defer ticker.Stop() for { // Check for pattern match + if match := bp.output.Match(pattern); match != "" { return &ToolResult{ - ForLLM: fmt.Sprintf("Match found in [%s]: %s", bgID, match), + ForLLM: fmt.Sprintf("Match found in [%s]: %s", bgID, match), + ForUser: fmt.Sprintf("Match found in [%s]: %s", bgID, match), } } // Check if process exited + if !bp.isRunning() { output := bp.output.String() + tail := lastNLines(output, 10) + var sb strings.Builder + fmt.Fprintf(&sb, "Process %s exited before pattern matched.\n", bgID) + if bp.exitErr != nil { fmt.Fprintf(&sb, "Exit: %v\n", bp.exitErr) } else { fmt.Fprintf(&sb, "Exit: 0\n") } + fmt.Fprintf(&sb, "\nLast output:\n%s", tail) + return &ToolResult{ - ForLLM: sb.String(), + ForLLM: sb.String(), + ForUser: sb.String(), + IsError: true, } } select { case <-deadline: + // Timeout + output := bp.output.String() + tail := lastNLines(output, 10) + var sb strings.Builder + fmt.Fprintf(&sb, "Watch timed out after %s waiting for pattern %q in [%s].\n", timeout, patternStr, bgID) + fmt.Fprintf(&sb, "\nLast output:\n%s", tail) + return &ToolResult{ - ForLLM: sb.String(), + ForLLM: sb.String(), + ForUser: sb.String(), + IsError: true, } + case <-ctx.Done(): + return ErrorResult("watch canceled") + case <-ticker.C: + // Continue polling } } @@ -200,17 +268,21 @@ func (t *BgMonitorTool) actionWatch(ctx context.Context, args map[string]any) *T func (t *BgMonitorTool) actionTail(args map[string]any) *ToolResult { bgID, _ := args["bg_id"].(string) + if bgID == "" { return ErrorResult("bg_id is required for tail action") } n := bgTailDefaultLines + if lines, ok := args["lines"].(float64); ok && lines > 0 { n = int(lines) } procs := t.exec.BgProcesses() + bp, ok := procs[bgID] + if !ok { return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) } @@ -218,7 +290,9 @@ func (t *BgMonitorTool) actionTail(args map[string]any) *ToolResult { lines := bp.output.Lines(n) var sb strings.Builder + fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command) + if bp.isRunning() { fmt.Fprintf(&sb, "Status: running\n") } else { @@ -228,7 +302,9 @@ func (t *BgMonitorTool) actionTail(args map[string]any) *ToolResult { fmt.Fprintf(&sb, "Status: exited=0\n") } } + fmt.Fprintf(&sb, "\nLast %d lines:\n", n) + for _, line := range lines { fmt.Fprintf(&sb, "%s\n", line) } @@ -238,19 +314,24 @@ func (t *BgMonitorTool) actionTail(args map[string]any) *ToolResult { } return &ToolResult{ - ForLLM: sb.String(), + ForLLM: sb.String(), + ForUser: sb.String(), } } // lastNLines returns the last n lines from a string. + func lastNLines(s string, n int) string { lines := strings.Split(s, "\n") + if len(lines) > 0 && lines[len(lines)-1] == "" { lines = lines[:len(lines)-1] } + if n >= len(lines) { return strings.Join(lines, "\n") } + return strings.Join(lines[len(lines)-n:], "\n") } diff --git a/pkg/tools/bg_monitor_test.go b/pkg/tools/bg_monitor_test.go index 0aca596cf..a30a8971b 100644 --- a/pkg/tools/bg_monitor_test.go +++ b/pkg/tools/bg_monitor_test.go @@ -10,64 +10,83 @@ import ( func TestBgMonitor_List(t *testing.T) { tool, _ := NewExecTool("", false) + monitor := NewBgMonitorTool(tool) // List with no processes + result := monitor.Execute(context.Background(), map[string]any{"action": "list"}) + if result.IsError { t.Fatalf("unexpected error: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "No background") { t.Errorf("expected 'No background' message, got: %s", result.ForLLM) } // Start two bg processes + var cmd1, cmd2 string + if runtime.GOOS == "windows" { cmd1 = "Start-Sleep -Seconds 30" + cmd2 = "Start-Sleep -Seconds 30" } else { cmd1 = "sleep 30" + cmd2 = "sleep 30" } r1 := tool.Execute(context.Background(), map[string]any{ - "command": cmd1, + "command": cmd1, + "background": true, }) + if r1.IsError { t.Fatalf("failed to start bg-1: %s", r1.ForLLM) } r2 := tool.Execute(context.Background(), map[string]any{ - "command": cmd2, + "command": cmd2, + "background": true, }) + if r2.IsError { t.Fatalf("failed to start bg-2: %s", r2.ForLLM) } // List should show both + result = monitor.Execute(context.Background(), map[string]any{"action": "list"}) + if result.IsError { t.Fatalf("unexpected error: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "bg-1") { t.Errorf("expected bg-1 in list, got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "bg-2") { t.Errorf("expected bg-2 in list, got: %s", result.ForLLM) } // Cleanup + tool.Shutdown() } func TestBgMonitor_Watch_Match(t *testing.T) { tool, _ := NewExecTool("", false) + monitor := NewBgMonitorTool(tool) var cmd string + if runtime.GOOS == "windows" { cmd = "Write-Output 'Server ready on port 3000'; Start-Sleep -Seconds 30" } else { @@ -75,26 +94,35 @@ func TestBgMonitor_Watch_Match(t *testing.T) { } r := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if r.IsError { t.Fatalf("failed to start bg: %s", r.ForLLM) } // Watch for "ready" pattern — should match quickly + result := monitor.Execute(context.Background(), map[string]any{ - "action": "watch", - "bg_id": "bg-1", - "pattern": "ready", + "action": "watch", + + "bg_id": "bg-1", + + "pattern": "ready", + "watch_timeout": float64(10), }) + if result.IsError { t.Fatalf("expected watch to match, got error: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "Match found") { t.Errorf("expected 'Match found' message, got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "ready") { t.Errorf("expected match to contain 'ready', got: %s", result.ForLLM) } @@ -104,9 +132,11 @@ func TestBgMonitor_Watch_Match(t *testing.T) { func TestBgMonitor_Watch_Timeout(t *testing.T) { tool, _ := NewExecTool("", false) + monitor := NewBgMonitorTool(tool) var cmd string + if runtime.GOOS == "windows" { cmd = "Start-Sleep -Seconds 30" } else { @@ -114,23 +144,31 @@ func TestBgMonitor_Watch_Timeout(t *testing.T) { } r := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if r.IsError { t.Fatalf("failed to start bg: %s", r.ForLLM) } // Watch for a pattern that won't appear, with short timeout + result := monitor.Execute(context.Background(), map[string]any{ - "action": "watch", - "bg_id": "bg-1", - "pattern": "never_going_to_match", + "action": "watch", + + "bg_id": "bg-1", + + "pattern": "never_going_to_match", + "watch_timeout": float64(1), }) + if !result.IsError { t.Fatalf("expected watch to timeout with error, got success: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "timed out") { t.Errorf("expected 'timed out' message, got: %s", result.ForLLM) } @@ -140,9 +178,11 @@ func TestBgMonitor_Watch_Timeout(t *testing.T) { func TestBgMonitor_Watch_ProcessExit(t *testing.T) { tool, _ := NewExecTool("", false) + monitor := NewBgMonitorTool(tool) var cmd string + if runtime.GOOS == "windows" { cmd = "Write-Output 'done quickly'" } else { @@ -150,26 +190,35 @@ func TestBgMonitor_Watch_ProcessExit(t *testing.T) { } r := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if r.IsError { t.Fatalf("failed to start bg: %s", r.ForLLM) } // Wait a bit for the process to exit + time.Sleep(4 * time.Second) // Watch for a pattern that doesn't match — process should have exited + result := monitor.Execute(context.Background(), map[string]any{ - "action": "watch", - "bg_id": "bg-1", - "pattern": "never_match", + "action": "watch", + + "bg_id": "bg-1", + + "pattern": "never_match", + "watch_timeout": float64(5), }) + if !result.IsError { t.Fatalf("expected error when process exits, got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "exited") { t.Errorf("expected 'exited' message, got: %s", result.ForLLM) } @@ -179,9 +228,11 @@ func TestBgMonitor_Watch_ProcessExit(t *testing.T) { func TestBgMonitor_Tail(t *testing.T) { tool, _ := NewExecTool("", false) + monitor := NewBgMonitorTool(tool) var cmd string + if runtime.GOOS == "windows" { cmd = "1..5 | ForEach-Object { Write-Output \"line $_\" }; Start-Sleep -Seconds 30" } else { @@ -189,25 +240,33 @@ func TestBgMonitor_Tail(t *testing.T) { } r := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if r.IsError { t.Fatalf("failed to start bg: %s", r.ForLLM) } // Wait for initial output to be captured + time.Sleep(4 * time.Second) // Tail last 3 lines + result := monitor.Execute(context.Background(), map[string]any{ "action": "tail", - "bg_id": "bg-1", - "lines": float64(3), + + "bg_id": "bg-1", + + "lines": float64(3), }) + if result.IsError { t.Fatalf("unexpected error: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "line 5") { t.Errorf("expected tail to contain 'line 5', got: %s", result.ForLLM) } @@ -217,12 +276,15 @@ func TestBgMonitor_Tail(t *testing.T) { func TestBgMonitor_InvalidAction(t *testing.T) { tool, _ := NewExecTool("", false) + monitor := NewBgMonitorTool(tool) result := monitor.Execute(context.Background(), map[string]any{"action": "invalid"}) + if !result.IsError { t.Fatalf("expected error for invalid action") } + if !strings.Contains(result.ForLLM, "unknown action") { t.Errorf("expected 'unknown action' message, got: %s", result.ForLLM) } diff --git a/pkg/tools/createpr.go b/pkg/tools/createpr.go index c0f351fb3..c630779b2 100644 --- a/pkg/tools/createpr.go +++ b/pkg/tools/createpr.go @@ -10,27 +10,42 @@ import ( const ( ciPollInterval = 30 * time.Second - ciPollTimeout = 15 * time.Minute + + ciPollTimeout = 15 * time.Minute ) // CreatePRTool creates a GitHub pull request from the current worktree branch. + // + // Safety invariants: + // - Only works inside a worktree (WorktreeInfo must be in context) + // - Base branch is auto-detected from WorktreeInfo.BaseBranch + // - Requires the branch to be already pushed (use git_push first) + // - Checks for merge conflicts with base before creating + // - Uses `gh pr create` under the hood + // + // Async behavior: + // - PR creation itself is synchronous and returns immediately with the PR URL + // - If CI runs are triggered, a background goroutine polls `gh pr checks` + // and calls the AsyncCallback when CI completes (pass or fail) + type CreatePRTool struct { callback AsyncCallback } // NewCreatePRTool creates a CreatePRTool. + func NewCreatePRTool() *CreatePRTool { return &CreatePRTool{} } @@ -38,122 +53,187 @@ func NewCreatePRTool() *CreatePRTool { func (t *CreatePRTool) Name() string { return "create_pr" } // SetCallback implements AsyncTool for CI completion notification. + func (t *CreatePRTool) SetCallback(cb AsyncCallback) { t.callback = cb } func (t *CreatePRTool) Description() string { return "Create a GitHub pull request from the current worktree branch. " + + "The base branch is auto-detected from the worktree's parent branch. " + + "The branch must be pushed to origin first (use git_push). " + + "Checks for merge conflicts with the base branch before creating. " + + "After PR creation, polls CI status in the background and notifies when complete. " + + "Requires the `gh` CLI to be installed and authenticated." } func (t *CreatePRTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "title": map[string]any{ - "type": "string", + "type": "string", + "description": "Pull request title", }, + "body": map[string]any{ - "type": "string", + "type": "string", + "description": "Pull request body/description (supports markdown)", }, + "draft": map[string]any{ - "type": "boolean", + "type": "boolean", + "description": "Create as draft PR (default: false)", }, }, + "required": []string{"title"}, } } func (t *CreatePRTool) Execute(ctx context.Context, args map[string]any) *ToolResult { wt := WorktreeInfoFromCtx(ctx) + if wt == nil { return ErrorResult( + "create_pr requires an active worktree.\n" + + "This tool can only be used during worktree-based sessions " + + "(e.g., heartbeat tasks or plan executing phase).\n" + + "The worktree provides the branch name and base branch for the PR.") } branch := wt.Branch + if branch == "" { return ErrorResult( + "worktree has no branch name.\n" + + "The WorktreeInfo was set but Branch is empty. " + + "This is an internal error — the worktree may not have been created correctly.") } baseBranch := wt.BaseBranch + if baseBranch == "" { baseBranch = "main" } title, ok := args["title"].(string) + if !ok || strings.TrimSpace(title) == "" { return ErrorResult( + "title is required.\n" + + "Provide a concise PR title describing the change (e.g., \"Add rate limiter to API endpoints\").") } // Verify the branch has been pushed by checking if the remote ref exists + checkCtx, checkCancel := context.WithTimeout(ctx, 15*time.Second) + defer checkCancel() + checkCmd := exec.CommandContext(checkCtx, "git", "ls-remote", "--exit-code", "origin", branch) + checkCmd.Dir = wt.Path + if err := checkCmd.Run(); err != nil { return ErrorResult(fmt.Sprintf( + "branch %q not found on origin.\n"+ + "The branch must be pushed before creating a PR. Use the git_push tool first.\n"+ + "git_push will auto-commit uncommitted changes and push the worktree branch to origin.", + branch)) } // Fetch latest base branch and check for merge conflicts + fetchCtx, fetchCancel := context.WithTimeout(ctx, 30*time.Second) + defer fetchCancel() + fetchCmd := exec.CommandContext(fetchCtx, "git", "fetch", "origin", baseBranch) + fetchCmd.Dir = wt.Path + if out, err := fetchCmd.CombinedOutput(); err != nil { return ErrorResult(fmt.Sprintf( + "failed to fetch origin/%s: %s\n%s\n"+ + "Cannot verify merge compatibility without the latest base branch. "+ + "Check network connectivity and that the base branch %q exists on origin.", + baseBranch, err, strings.TrimSpace(string(out)), baseBranch)) } // Try a merge dry-run to detect conflicts. + // merge-tree --write-tree is a plumbing command (Git 2.38+) that performs a + // three-way merge entirely in-memory without touching the working tree. + // Exit code 0 = clean merge, non-zero = conflicts detected. + mergeCtx, mergeCancel := context.WithTimeout(ctx, 30*time.Second) + defer mergeCancel() + mergeCmd := exec.CommandContext(mergeCtx, "git", "merge-tree", + "--write-tree", "--no-messages", + branch, "origin/"+baseBranch) + mergeCmd.Dir = wt.RepoRoot + mergeOut, mergeErr := mergeCmd.CombinedOutput() + if mergeErr != nil { conflictInfo := strings.TrimSpace(string(mergeOut)) + return ErrorResult(fmt.Sprintf( + "merge conflict detected between %q and %s.\n"+ + "The PR cannot be created cleanly. Resolve the conflicts in the worktree first, "+ + "then use git_push to push the resolution before retrying create_pr.\n"+ + "Conflict details:\n%s", + branch, baseBranch, conflictInfo)) } // Build gh pr create command + ghArgs := []string{ "pr", "create", + "--base", baseBranch, + "--head", branch, + "--title", title, } @@ -168,92 +248,143 @@ func (t *CreatePRTool) Execute(ctx context.Context, args map[string]any) *ToolRe } prCtx, prCancel := context.WithTimeout(ctx, 30*time.Second) + defer prCancel() cmd := exec.CommandContext(prCtx, "gh", ghArgs...) + cmd.Dir = wt.RepoRoot + out, err := cmd.CombinedOutput() + output := strings.TrimSpace(string(out)) if err != nil { return ErrorResult(fmt.Sprintf( + "gh pr create failed: %s\n%s\n"+ + "Possible causes:\n"+ + "- gh CLI not installed or not authenticated (run `gh auth login`)\n"+ + "- A PR already exists for branch %q (check with `gh pr list`)\n"+ + "- Repository not configured as a GitHub remote", + err, output, branch)) } prURL := output // gh pr create outputs the PR URL // Start background CI polling if callback is set + if t.callback != nil && prURL != "" { cb := t.callback + repoRoot := wt.RepoRoot + go pollCIStatus(repoRoot, prURL, cb) } return AsyncResult(fmt.Sprintf( + "Pull request created: %s\n"+ + "Branch: %s -> %s\n"+ + "CI status will be reported asynchronously when checks complete.", + prURL, branch, baseBranch)) } // pollCIStatus polls `gh pr checks` in the background until all checks + // pass, fail, or the timeout is reached. Reports back via AsyncCallback. + func pollCIStatus(repoRoot, prURL string, callback AsyncCallback) { // Detached context with hard timeout — this goroutine outlives the tool call. + ctx, cancel := context.WithTimeout(context.Background(), ciPollTimeout) + defer cancel() // Initial wait: CI runs take a few seconds to register after PR creation + select { case <-time.After(10 * time.Second): + case <-ctx.Done(): + return } ticker := time.NewTicker(ciPollInterval) + defer ticker.Stop() for { status, detail := checkPRChecks(ctx, repoRoot, prURL) + switch status { case ciStatusPass: + callback(ctx, NewToolResult(fmt.Sprintf( + "CI passed for %s\n%s", + prURL, detail))) + return + case ciStatusFail: + callback(ctx, &ToolResult{ ForLLM: fmt.Sprintf( + "CI failed for %s\n%s\n"+ + "Run `gh run view` for detailed logs.", + prURL, detail), + IsError: true, }) + return + case ciStatusNone: + callback(ctx, NewToolResult(fmt.Sprintf( + "No CI checks configured for %s. PR is ready for review.", + prURL))) + return + case ciStatusPending: + // Still running, continue polling } select { case <-ticker.C: + case <-ctx.Done(): + callback(ctx, &ToolResult{ ForLLM: fmt.Sprintf( + "CI polling timed out after %s for %s.\n"+ + "Checks may still be running. Run `gh pr checks %s` to check.", + ciPollTimeout, prURL, prURL), + IsError: true, }) + return } } @@ -263,36 +394,51 @@ type ciStatus int const ( ciStatusPending ciStatus = iota + ciStatusPass + ciStatusFail + ciStatusNone ) // checkPRChecks runs `gh pr checks` and parses the result. + // Returns the aggregate status and raw output for the caller to include. + func checkPRChecks(ctx context.Context, repoRoot, prURL string) (ciStatus, string) { checkCtx, cancel := context.WithTimeout(ctx, 15*time.Second) + defer cancel() cmd := exec.CommandContext(checkCtx, "gh", "pr", "checks", prURL) + cmd.Dir = repoRoot + out, err := cmd.CombinedOutput() + output := strings.TrimSpace(string(out)) if err != nil { // gh pr checks exits 1 when any check has failed + if strings.Contains(output, "fail") || strings.Contains(output, "X ") { return ciStatusFail, output } + // "no checks" case + if strings.Contains(output, "no checks") || output == "" { return ciStatusNone, "" } + // Transient error or still pending — keep polling + return ciStatusPending, output } // Exit 0: all checks completed. Check for pending. + if strings.Contains(output, "pending") || strings.Contains(output, "- ") { return ciStatusPending, output } diff --git a/pkg/tools/createpr_test.go b/pkg/tools/createpr_test.go index 26c330b7c..99c46cbb4 100644 --- a/pkg/tools/createpr_test.go +++ b/pkg/tools/createpr_test.go @@ -9,191 +9,259 @@ import ( ) // TestCreatePRTool_NoWorktree verifies that create_pr fails without worktree context. + func TestCreatePRTool_NoWorktree(t *testing.T) { tool := NewCreatePRTool() result := tool.Execute(context.Background(), map[string]any{ "title": "Test PR", }) + if !result.IsError { t.Fatal("expected error when no worktree in context") } + assertContains(t, result.ForLLM, "worktree") + assertContains(t, result.ForLLM, "heartbeat") } // TestCreatePRTool_EmptyBranch verifies that empty branch name is rejected. + func TestCreatePRTool_EmptyBranch(t *testing.T) { tool := NewCreatePRTool() ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ - Branch: "", + Branch: "", + BaseBranch: "main", - Path: t.TempDir(), - RepoRoot: t.TempDir(), + + Path: t.TempDir(), + + RepoRoot: t.TempDir(), }) + result := tool.Execute(ctx, map[string]any{ "title": "Test PR", }) + if !result.IsError { t.Fatal("expected error for empty branch") } + assertContains(t, result.ForLLM, "no branch name") } // TestCreatePRTool_MissingTitle verifies that missing title is rejected. + func TestCreatePRTool_MissingTitle(t *testing.T) { tool := NewCreatePRTool() ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ - Branch: "plan/test", + Branch: "plan/test", + BaseBranch: "main", - Path: t.TempDir(), - RepoRoot: t.TempDir(), + + Path: t.TempDir(), + + RepoRoot: t.TempDir(), }) tests := []struct { name string + args map[string]any }{ {"no title key", map[string]any{}}, + {"empty title", map[string]any{"title": ""}}, + {"whitespace title", map[string]any{"title": " "}}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := tool.Execute(ctx, tt.args) + if !result.IsError { t.Fatal("expected error for missing/empty title") } + assertContains(t, result.ForLLM, "title is required") }) } } // TestCreatePRTool_BranchNotPushed verifies the tool checks for remote branch existence. + func TestCreatePRTool_BranchNotPushed(t *testing.T) { tool := NewCreatePRTool() ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ - Branch: "plan/not-pushed", + Branch: "plan/not-pushed", + BaseBranch: "main", - Path: t.TempDir(), - RepoRoot: t.TempDir(), + + Path: t.TempDir(), + + RepoRoot: t.TempDir(), }) + result := tool.Execute(ctx, map[string]any{ "title": "Test PR", }) + if !result.IsError { t.Fatal("expected error for unpushed branch") } + // Should mention git_push as the remedy + assertContains(t, result.ForLLM, "git_push") } // TestCreatePRTool_DefaultBaseBranch verifies fallback to "main" when BaseBranch is empty. + func TestCreatePRTool_DefaultBaseBranch(t *testing.T) { tool := NewCreatePRTool() // With empty BaseBranch, tool should default to "main" + ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ - Branch: "plan/test", + Branch: "plan/test", + BaseBranch: "", - Path: t.TempDir(), - RepoRoot: t.TempDir(), + + Path: t.TempDir(), + + RepoRoot: t.TempDir(), }) + result := tool.Execute(ctx, map[string]any{ "title": "Test PR", }) + // Will fail at ls-remote (no real repo), but should not fail at baseBranch validation + if result.IsError && strings.Contains(result.ForLLM, "base branch") { t.Fatal("should not fail on base branch when defaulting to main") } } // TestCreatePRTool_Interface verifies the tool satisfies both Tool and AsyncTool interfaces. + func TestCreatePRTool_Interface(t *testing.T) { var _ Tool = (*CreatePRTool)(nil) + var _ AsyncTool = (*CreatePRTool)(nil) tool := NewCreatePRTool() + if tool.Name() != "create_pr" { t.Errorf("Name: got %q, want %q", tool.Name(), "create_pr") } + if tool.Description() == "" { t.Error("Description should not be empty") } + params := tool.Parameters() + if params == nil { t.Fatal("Parameters should not be nil") } // Verify "title" is required + required, ok := params["required"].([]string) + if !ok { t.Fatal("required should be []string") } + foundTitle := false + for _, r := range required { if r == "title" { foundTitle = true } } + if !foundTitle { t.Error("title should be in required parameters") } } // TestCreatePRTool_SetCallback verifies callback is stored. + func TestCreatePRTool_SetCallback(t *testing.T) { tool := NewCreatePRTool() + if tool.callback != nil { t.Fatal("callback should be nil initially") } called := false + tool.SetCallback(func(ctx context.Context, result *ToolResult) { called = true }) + if tool.callback == nil { t.Fatal("callback should be set after SetCallback") } + // Verify it's callable (doesn't panic) + tool.callback(context.Background(), NewToolResult("test")) + if !called { t.Fatal("callback was not invoked") } } // TestCheckPRChecks_ParseResults tests CI status parsing logic. + func TestCheckPRChecks_ParseResults(t *testing.T) { // This tests the parsing logic conceptually — actual `gh` calls + // would need integration tests. We verify the status constants exist + // and the type is usable. + if ciStatusPending != 0 { t.Error("ciStatusPending should be 0 (default)") } + if ciStatusPass == ciStatusFail { t.Error("ciStatusPass and ciStatusFail should differ") } + if ciStatusNone == ciStatusPending { t.Error("ciStatusNone and ciStatusPending should differ") } } // TestAllowedToolsForPreset_GitTools checks git tools are correctly assigned to presets. + func TestAllowedToolsForPreset_GitTools(t *testing.T) { tests := []struct { - name string - preset Preset - wantGitPush bool + name string + + preset Preset + + wantGitPush bool + wantCreatePR bool }{ {"scout", PresetScout, false, false}, + {"analyst", PresetAnalyst, false, false}, + {"coder", PresetCoder, true, false}, + {"worker", PresetWorker, true, true}, + {"coordinator", PresetCoordinator, true, true}, } @@ -204,6 +272,7 @@ func TestAllowedToolsForPreset_GitTools(t *testing.T) { if got := allowed["git_push"]; got != tt.wantGitPush { t.Errorf("git_push: got %v, want %v", got, tt.wantGitPush) } + if got := allowed["create_pr"]; got != tt.wantCreatePR { t.Errorf("create_pr: got %v, want %v", got, tt.wantCreatePR) } diff --git a/pkg/tools/cron.go b/pkg/tools/cron.go index e7a380a2c..6489cba6e 100644 --- a/pkg/tools/cron.go +++ b/pkg/tools/cron.go @@ -14,25 +14,36 @@ import ( ) // JobExecutor is the interface for executing cron jobs through the agent + type JobExecutor interface { ProcessDirectWithChannel(ctx context.Context, content, sessionKey, channel, chatID string) (string, error) } // CronTool provides scheduling capabilities for the agent + type CronTool struct { cronService *cron.CronService - executor JobExecutor - msgBus *bus.MessageBus - execTool *ExecTool - channel string - chatID string - mu sync.RWMutex + + executor JobExecutor + + msgBus *bus.MessageBus + + execTool *ExecTool + + channel string + + chatID string + + mu sync.RWMutex } // NewCronTool creates a new CronTool + // execTimeout: 0 means no timeout, >0 sets the timeout duration + func NewCronTool( cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool, + execTimeout time.Duration, config *config.Config, ) (*CronTool, error) { execTool, err := NewExecToolWithConfig(workspace, restrict, config) @@ -41,102 +52,147 @@ func NewCronTool( } execTool.SetTimeout(execTimeout) + return &CronTool{ cronService: cronService, - executor: executor, - msgBus: msgBus, - execTool: execTool, + + executor: executor, + + msgBus: msgBus, + + execTool: execTool, }, nil } // Name returns the tool name + func (t *CronTool) Name() string { return "cron" } // Description returns the tool description + func (t *CronTool) Description() string { return "Schedule reminders, tasks, or system commands. IMPORTANT: When user asks to be reminded or scheduled, you MUST call this tool. Use 'at_seconds' for one-time reminders (e.g., 'remind me in 10 minutes' → at_seconds=600). Use 'every_seconds' ONLY for recurring tasks (e.g., 'every 2 hours' → every_seconds=7200). Use 'cron_expr' for complex recurring schedules. Use 'command' to execute shell commands directly." } // Parameters returns the tool parameters schema + func (t *CronTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "action": map[string]any{ - "type": "string", - "enum": []string{"add", "list", "remove", "enable", "disable"}, + "type": "string", + + "enum": []string{"add", "list", "remove", "enable", "disable"}, + "description": "Action to perform. Use 'add' when user wants to schedule a reminder or task.", }, + "message": map[string]any{ - "type": "string", + "type": "string", + "description": "The reminder/task message to display when triggered. If 'command' is used, this describes what the command does.", }, + "command": map[string]any{ - "type": "string", + "type": "string", + "description": "Optional: Shell command to execute directly (e.g., 'df -h'). If set, the agent will run this command and report output instead of just showing the message. 'deliver' will be forced to false for commands.", }, + "at_seconds": map[string]any{ - "type": "integer", + "type": "integer", + "description": "One-time reminder: seconds from now when to trigger (e.g., 600 for 10 minutes later). Use this for one-time reminders like 'remind me in 10 minutes'.", }, + "every_seconds": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Recurring interval in seconds (e.g., 3600 for every hour). Use this ONLY for recurring tasks like 'every 2 hours' or 'daily reminder'.", }, + "cron_expr": map[string]any{ - "type": "string", + "type": "string", + "description": "Cron expression for complex recurring schedules (e.g., '0 9 * * *' for daily at 9am). Use this for complex recurring schedules.", }, + "job_id": map[string]any{ - "type": "string", + "type": "string", + "description": "Job ID (for remove/enable/disable)", }, + "deliver": map[string]any{ - "type": "boolean", + "type": "boolean", + "description": "If true, send message directly to channel. If false, let agent process message (for complex tasks). Default: true", }, }, + "required": []string{"action"}, } } // SetContext sets the current session context for job creation + func (t *CronTool) SetContext(channel, chatID string) { t.mu.Lock() + defer t.mu.Unlock() + t.channel = channel + t.chatID = chatID } // Execute runs the tool with the given arguments + func (t *CronTool) Execute(ctx context.Context, args map[string]any) *ToolResult { action, ok := args["action"].(string) + if !ok { return ErrorResult("action is required") } switch action { case "add": + return t.addJob(args) + case "list": + return t.listJobs() + case "remove": + return t.removeJob(args) + case "enable": + return t.enableJob(args, true) + case "disable": + return t.enableJob(args, false) + default: + return ErrorResult(fmt.Sprintf("unknown action: %s", action)) } } func (t *CronTool) addJob(args map[string]any) *ToolResult { t.mu.RLock() + channel := t.channel + chatID := t.chatID + t.mu.RUnlock() if channel == "" || chatID == "" { @@ -144,6 +200,7 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult { } message, ok := args["message"].(string) + if !ok || message == "" { return ErrorResult("message is required for add") } @@ -151,26 +208,35 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult { var schedule cron.CronSchedule // Check for at_seconds (one-time), every_seconds (recurring), or cron_expr + atSeconds, hasAt := args["at_seconds"].(float64) + everySeconds, hasEvery := args["every_seconds"].(float64) + cronExpr, hasCron := args["cron_expr"].(string) // Priority: at_seconds > every_seconds > cron_expr + if hasAt { atMS := time.Now().UnixMilli() + int64(atSeconds)*1000 + schedule = cron.CronSchedule{ Kind: "at", + AtMS: &atMS, } } else if hasEvery { everyMS := int64(everySeconds) * 1000 + schedule = cron.CronSchedule{ - Kind: "every", + Kind: "every", + EveryMS: &everyMS, } } else if hasCron { schedule = cron.CronSchedule{ Kind: "cron", + Expr: cronExpr, } } else { @@ -178,29 +244,43 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult { } // Read deliver parameter, default to true + deliver := true + if d, ok := args["deliver"].(bool); ok { deliver = d } command, _ := args["command"].(string) + if command != "" { // Commands must be processed by agent/exec tool, so deliver must be false (or handled specifically) + // Actually, let's keep deliver=false to let the system know it's not a simple chat message + // But for our new logic in ExecuteJob, we can handle it regardless of deliver flag if Payload.Command is set. + // However, logically, it's not "delivered" to chat directly as is. + deliver = false } // Truncate message for job name (max 30 chars) + messagePreview := utils.Truncate(message, 30) job, err := t.cronService.AddJob( + messagePreview, + schedule, + message, + deliver, + channel, + chatID, ) if err != nil { @@ -209,7 +289,9 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult { if command != "" { job.Payload.Command = command + // Need to save the updated payload + t.cronService.UpdateJob(job) } @@ -224,9 +306,12 @@ func (t *CronTool) listJobs() *ToolResult { } var sb strings.Builder + sb.WriteString("Scheduled jobs:\n") + for _, j := range jobs { var scheduleInfo string + if j.Schedule.Kind == "every" && j.Schedule.EveryMS != nil { scheduleInfo = fmt.Sprintf("every %ds", *j.Schedule.EveryMS/1000) } else if j.Schedule.Kind == "cron" { @@ -236,6 +321,7 @@ func (t *CronTool) listJobs() *ToolResult { } else { scheduleInfo = "unknown" } + fmt.Fprintf(&sb, "- %s (id: %s, %s)\n", j.Name, j.ID, scheduleInfo) } @@ -244,6 +330,7 @@ func (t *CronTool) listJobs() *ToolResult { func (t *CronTool) removeJob(args map[string]any) *ToolResult { jobID, ok := args["job_id"].(string) + if !ok || jobID == "" { return ErrorResult("job_id is required for remove") } @@ -251,49 +338,62 @@ func (t *CronTool) removeJob(args map[string]any) *ToolResult { if t.cronService.RemoveJob(jobID) { return SilentResult(fmt.Sprintf("Cron job removed: %s", jobID)) } + return ErrorResult(fmt.Sprintf("Job %s not found", jobID)) } func (t *CronTool) enableJob(args map[string]any, enable bool) *ToolResult { jobID, ok := args["job_id"].(string) + if !ok || jobID == "" { return ErrorResult("job_id is required for enable/disable") } job := t.cronService.EnableJob(jobID, enable) + if job == nil { return ErrorResult(fmt.Sprintf("Job %s not found", jobID)) } status := "enabled" + if !enable { status = "disabled" } + return SilentResult(fmt.Sprintf("Cron job '%s' %s", job.Name, status)) } // ExecuteJob executes a cron job through the agent + func (t *CronTool) ExecuteJob(ctx context.Context, job *cron.CronJob) string { // Get channel/chatID from job payload + channel := job.Payload.Channel + chatID := job.Payload.To // Default values if not set + if channel == "" { channel = "cli" } + if chatID == "" { chatID = "direct" } // Execute command if present + if job.Payload.Command != "" { args := map[string]any{ "command": job.Payload.Command, } result := t.execTool.Execute(ctx, args) + var output string + if result.IsError { output = fmt.Sprintf("Error executing scheduled command: %s", result.ForLLM) } else { @@ -301,36 +401,54 @@ func (t *CronTool) ExecuteJob(ctx context.Context, job *cron.CronJob) string { } pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + t.msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{ Channel: channel, - ChatID: chatID, + + ChatID: chatID, + Content: output, }) + return "ok" } // If deliver=true, send message directly without agent processing + if job.Payload.Deliver { pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer pubCancel() + t.msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{ Channel: channel, - ChatID: chatID, + + ChatID: chatID, + Content: job.Payload.Message, }) + return "ok" } // For deliver=false, process through agent (for complex tasks) + sessionKey := fmt.Sprintf("cron-%s", job.ID) // Call agent with job's message + response, err := t.executor.ProcessDirectWithChannel( + ctx, + job.Payload.Message, + sessionKey, + channel, + chatID, ) if err != nil { @@ -338,6 +456,8 @@ func (t *CronTool) ExecuteJob(ctx context.Context, job *cron.CronJob) string { } // Response is automatically sent via MessageBus by AgentLoop + _ = response // Will be sent by AgentLoop + return "ok" } diff --git a/pkg/tools/dev_preview.go b/pkg/tools/dev_preview.go index 0ee3658ca..0e4fa0fde 100644 --- a/pkg/tools/dev_preview.go +++ b/pkg/tools/dev_preview.go @@ -10,11 +10,13 @@ import ( ) // DevPreviewTool allows the agent to control the Mini App dev reverse proxy. + type DevPreviewTool struct { manager miniapp.DevTargetManager } // NewDevPreviewTool creates a new DevPreviewTool. + func NewDevPreviewTool(manager miniapp.DevTargetManager) *DevPreviewTool { return &DevPreviewTool{manager: manager} } @@ -28,113 +30,157 @@ func (t *DevPreviewTool) Description() string { func (t *DevPreviewTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "action": map[string]any{ - "type": "string", - "enum": []string{"start", "stop", "unregister", "status"}, + "type": "string", + + "enum": []string{"start", "stop", "unregister", "status"}, + "description": "Action to perform: start (register + activate target), stop (deactivate proxy), unregister (remove a registered target), status (list all targets).", }, + "target": map[string]any{ - "type": "string", + "type": "string", + "description": "Target URL for the dev server (e.g. http://localhost:3000). Required for 'start' action. Must be a localhost URL.", }, + "name": map[string]any{ - "type": "string", + "type": "string", + "description": "Display name for the target (e.g. 'frontend'). Optional for 'start' action; auto-generated from host:port if omitted.", }, + "id": map[string]any{ - "type": "string", + "type": "string", + "description": "Target ID. Required for 'unregister' action.", }, }, + "required": []string{"action"}, } } func (t *DevPreviewTool) Execute(ctx context.Context, args map[string]any) *ToolResult { action, ok := args["action"].(string) + if !ok { return ErrorResult("action is required") } switch action { case "start": + target, _ := args["target"].(string) + if target == "" { return ErrorResult("target is required for start action") } + name, _ := args["name"].(string) + if name == "" { name = inferName(target) } + id, err := t.manager.RegisterDevTarget(name, target) if err != nil { return ErrorResult(fmt.Sprintf("failed to register dev target: %v", err)) } + if err := t.manager.ActivateDevTarget(id); err != nil { return ErrorResult(fmt.Sprintf("failed to activate dev target: %v", err)) } + return SilentResult( + fmt.Sprintf( + "Dev preview started (id=%s, name=%s). Target: %s\nUsers can view it in the Mini App Dev tab.", + id, + name, + target, ), ) case "stop": + if err := t.manager.DeactivateDevTarget(); err != nil { return ErrorResult(fmt.Sprintf("failed to stop dev preview: %v", err)) } + return SilentResult("Dev preview stopped.") case "unregister": + id, _ := args["id"].(string) + if id == "" { return ErrorResult("id is required for unregister action") } + if err := t.manager.UnregisterDevTarget(id); err != nil { return ErrorResult(fmt.Sprintf("failed to unregister target: %v", err)) } + return SilentResult(fmt.Sprintf("Dev target %s unregistered.", id)) case "status": + targets := t.manager.ListDevTargets() + active := t.manager.GetDevTarget() + if len(targets) == 0 { if active == "" { return SilentResult("Dev preview is not active. No targets registered.") } + return SilentResult(fmt.Sprintf("Dev preview is active. Target: %s\nNo registered targets.", active)) } + var sb strings.Builder + if active != "" { sb.WriteString(fmt.Sprintf("Dev preview is active. Target: %s\n", active)) } else { sb.WriteString("Dev preview is not active.\n") } + sb.WriteString("Registered targets:\n") + for _, dt := range targets { sb.WriteString(fmt.Sprintf(" [%s] %s → %s\n", dt.ID, dt.Name, dt.Target)) } + return SilentResult(sb.String()) default: + return ErrorResult(fmt.Sprintf("unknown action: %s", action)) } } // inferName generates a display name from a target URL (e.g. "localhost:3000"). + func inferName(target string) string { u, err := url.Parse(target) if err != nil { return target } + host := u.Hostname() + port := u.Port() + if port != "" { return host + ":" + port } + return host } diff --git a/pkg/tools/dev_preview_test.go b/pkg/tools/dev_preview_test.go index 01c884777..3db4c37a6 100644 --- a/pkg/tools/dev_preview_test.go +++ b/pkg/tools/dev_preview_test.go @@ -10,12 +10,17 @@ import ( ) // mockDevTargetManager implements miniapp.DevTargetManager for testing. + type mockDevTargetManager struct { - targets map[string]*miniapp.DevTarget - nextID int + targets map[string]*miniapp.DevTarget + + nextID int + activeID string - active string // active target URL - regErr error + + active string // active target URL + + regErr error } func newMockManager() *mockDevTargetManager { @@ -26,9 +31,13 @@ func (m *mockDevTargetManager) RegisterDevTarget(name, target string) (string, e if m.regErr != nil { return "", m.regErr } + m.nextID++ + id := fmt.Sprintf("%d", m.nextID) + m.targets[id] = &miniapp.DevTarget{ID: id, Name: name, Target: target} + return id, nil } @@ -36,27 +45,37 @@ func (m *mockDevTargetManager) UnregisterDevTarget(id string) error { if _, ok := m.targets[id]; !ok { return fmt.Errorf("target %q not found", id) } + delete(m.targets, id) + if m.activeID == id { m.activeID = "" + m.active = "" } + return nil } func (m *mockDevTargetManager) ActivateDevTarget(id string) error { dt, ok := m.targets[id] + if !ok { return fmt.Errorf("target %q not found", id) } + m.activeID = id + m.active = dt.Target + return nil } func (m *mockDevTargetManager) DeactivateDevTarget() error { m.activeID = "" + m.active = "" + return nil } @@ -66,34 +85,43 @@ func (m *mockDevTargetManager) GetDevTarget() string { func (m *mockDevTargetManager) ListDevTargets() []miniapp.DevTarget { out := make([]miniapp.DevTarget, 0, len(m.targets)) + for _, dt := range m.targets { out = append(out, *dt) } + return out } func TestDevPreviewTool_Start(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://localhost:3000", - "name": "frontend", + + "name": "frontend", }) if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + if len(mgr.targets) != 1 { t.Errorf("expected 1 registered target, got %d", len(mgr.targets)) } + if mgr.active != "http://localhost:3000" { t.Errorf("expected active target http://localhost:3000, got %q", mgr.active) } + if !strings.Contains(result.ForLLM, "started") { t.Errorf("expected result to contain 'started', got %q", result.ForLLM) } + if !strings.Contains(result.ForLLM, "frontend") { t.Errorf("expected result to contain 'frontend', got %q", result.ForLLM) } @@ -101,17 +129,21 @@ func TestDevPreviewTool_Start(t *testing.T) { func TestDevPreviewTool_StartAutoName(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://localhost:3000", }) if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + // Auto-generated name should be "localhost:3000" + for _, dt := range mgr.targets { if dt.Name != "localhost:3000" { t.Errorf("expected auto-name 'localhost:3000', got %q", dt.Name) @@ -121,6 +153,7 @@ func TestDevPreviewTool_StartAutoName(t *testing.T) { func TestDevPreviewTool_StartMissingTarget(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ @@ -134,11 +167,14 @@ func TestDevPreviewTool_StartMissingTarget(t *testing.T) { func TestDevPreviewTool_StartError(t *testing.T) { mgr := newMockManager() + mgr.regErr = fmt.Errorf("only localhost") + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://example.com:3000", }) @@ -149,7 +185,9 @@ func TestDevPreviewTool_StartError(t *testing.T) { func TestDevPreviewTool_Stop(t *testing.T) { mgr := newMockManager() + mgr.active = "http://localhost:3000" + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ @@ -159,6 +197,7 @@ func TestDevPreviewTool_Stop(t *testing.T) { if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + if mgr.active != "" { t.Errorf("expected empty active target after stop, got %q", mgr.active) } @@ -166,19 +205,23 @@ func TestDevPreviewTool_Stop(t *testing.T) { func TestDevPreviewTool_Unregister(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) // Register a target first + id, _ := mgr.RegisterDevTarget("frontend", "http://localhost:3000") result := tool.Execute(context.Background(), map[string]any{ "action": "unregister", - "id": id, + + "id": id, }) if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + if len(mgr.targets) != 0 { t.Errorf("expected 0 targets after unregister, got %d", len(mgr.targets)) } @@ -186,6 +229,7 @@ func TestDevPreviewTool_Unregister(t *testing.T) { func TestDevPreviewTool_UnregisterMissingID(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ @@ -199,11 +243,13 @@ func TestDevPreviewTool_UnregisterMissingID(t *testing.T) { func TestDevPreviewTool_UnregisterNotFound(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ "action": "unregister", - "id": "999", + + "id": "999", }) if !result.IsError { @@ -213,10 +259,13 @@ func TestDevPreviewTool_UnregisterNotFound(t *testing.T) { func TestDevPreviewTool_Status(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) mgr.RegisterDevTarget("api", "http://localhost:8080") + mgr.RegisterDevTarget("frontend", "http://localhost:3000") + mgr.active = "http://localhost:8080" result := tool.Execute(context.Background(), map[string]any{ @@ -226,15 +275,19 @@ func TestDevPreviewTool_Status(t *testing.T) { if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "active") { t.Errorf("expected 'active' in result, got %q", result.ForLLM) } + if !strings.Contains(result.ForLLM, "http://localhost:8080") { t.Errorf("expected target URL in result, got %q", result.ForLLM) } + if !strings.Contains(result.ForLLM, "api") { t.Errorf("expected 'api' in result, got %q", result.ForLLM) } + if !strings.Contains(result.ForLLM, "frontend") { t.Errorf("expected 'frontend' in result, got %q", result.ForLLM) } @@ -242,6 +295,7 @@ func TestDevPreviewTool_Status(t *testing.T) { func TestDevPreviewTool_StatusInactive(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ @@ -251,6 +305,7 @@ func TestDevPreviewTool_StatusInactive(t *testing.T) { if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "not active") { t.Errorf("expected 'not active' in result, got %q", result.ForLLM) } @@ -258,6 +313,7 @@ func TestDevPreviewTool_StatusInactive(t *testing.T) { func TestDevPreviewTool_UnknownAction(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ @@ -271,6 +327,7 @@ func TestDevPreviewTool_UnknownAction(t *testing.T) { func TestDevPreviewTool_MissingAction(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{}) @@ -282,15 +339,19 @@ func TestDevPreviewTool_MissingAction(t *testing.T) { func TestDevPreviewTool_NameAndSchema(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) if tool.Name() != "dev_preview" { t.Errorf("expected name dev_preview, got %q", tool.Name()) } + if tool.Description() == "" { t.Error("expected non-empty description") } + params := tool.Parameters() + if params == nil { t.Fatal("expected non-nil parameters") } @@ -300,26 +361,35 @@ func TestDevPreviewTool_NameAndSchema(t *testing.T) { func TestDevPreviewTool_StartMultipleTargets(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) r1 := tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://localhost:8080", - "name": "api", + + "name": "api", }) + r2 := tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://localhost:3000", - "name": "frontend", + + "name": "frontend", }) if r1.IsError || r2.IsError { t.Fatalf("expected both starts to succeed, got err1=%v err2=%v", r1.IsError, r2.IsError) } + if len(mgr.targets) != 2 { t.Errorf("expected 2 registered targets, got %d", len(mgr.targets)) } + // The second start should make the frontend active + if mgr.active != "http://localhost:3000" { t.Errorf("expected last started target to be active, got %q", mgr.active) } @@ -327,12 +397,15 @@ func TestDevPreviewTool_StartMultipleTargets(t *testing.T) { func TestDevPreviewTool_StopPreservesRegistrations(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://localhost:3000", - "name": "frontend", + + "name": "frontend", }) result := tool.Execute(context.Background(), map[string]any{ @@ -342,11 +415,15 @@ func TestDevPreviewTool_StopPreservesRegistrations(t *testing.T) { if result.IsError { t.Fatalf("stop failed: %s", result.ForLLM) } + // Registration should still be there + if len(mgr.targets) != 1 { t.Errorf("expected 1 registered target after stop, got %d", len(mgr.targets)) } + // But active should be cleared + if mgr.active != "" { t.Errorf("expected inactive after stop, got %q", mgr.active) } @@ -354,9 +431,11 @@ func TestDevPreviewTool_StopPreservesRegistrations(t *testing.T) { func TestDevPreviewTool_StatusWithTargetsButInactive(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) mgr.RegisterDevTarget("api", "http://localhost:8080") + // active remains empty result := tool.Execute(context.Background(), map[string]any{ @@ -366,9 +445,11 @@ func TestDevPreviewTool_StatusWithTargetsButInactive(t *testing.T) { if result.IsError { t.Fatalf("status failed: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "not active") { t.Errorf("expected 'not active' in status, got %q", result.ForLLM) } + if !strings.Contains(result.ForLLM, "api") { t.Errorf("expected 'api' listed in status, got %q", result.ForLLM) } @@ -376,22 +457,29 @@ func TestDevPreviewTool_StatusWithTargetsButInactive(t *testing.T) { func TestDevPreviewTool_ResultIsSilent(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) cases := []struct { name string + args map[string]any }{ {"start", map[string]any{"action": "start", "target": "http://localhost:3000"}}, + {"stop", map[string]any{"action": "stop"}}, + {"status", map[string]any{"action": "status"}}, } + for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { result := tool.Execute(context.Background(), tc.args) + if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + if result.Silent != true { t.Errorf("expected SilentResult (IsSilent=true), got IsSilent=%v", result.Silent) } @@ -401,11 +489,13 @@ func TestDevPreviewTool_ResultIsSilent(t *testing.T) { func TestDevPreviewTool_ActionTypeNotString(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ "action": 123, }) + if !result.IsError { t.Error("expected error for non-string action") } @@ -414,17 +504,26 @@ func TestDevPreviewTool_ActionTypeNotString(t *testing.T) { func TestDevPreviewTool_InferName(t *testing.T) { cases := []struct { target string - want string + + want string }{ {"http://localhost:3000", "localhost:3000"}, + {"http://localhost:8080", "localhost:8080"}, + {"http://127.0.0.1:9000", "127.0.0.1:9000"}, + {"http://localhost", "localhost"}, + {"http://[::1]:5000", "::1:5000"}, + {"not-a-url", ""}, // url.Parse succeeds but Hostname() is empty + } + for _, tc := range cases { got := inferName(tc.target) + if got != tc.want { t.Errorf("inferName(%q) = %q, want %q", tc.target, got, tc.want) } @@ -433,18 +532,23 @@ func TestDevPreviewTool_InferName(t *testing.T) { func TestDevPreviewTool_StartEmptyName(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) // Explicitly pass empty name — should auto-infer + result := tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://localhost:5000", - "name": "", + + "name": "", }) if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + for _, dt := range mgr.targets { if dt.Name != "localhost:5000" { t.Errorf("expected auto-name 'localhost:5000', got %q", dt.Name) @@ -454,32 +558,41 @@ func TestDevPreviewTool_StartEmptyName(t *testing.T) { func TestDevPreviewTool_UnregisterActiveTarget(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) // Register and activate + tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "http://localhost:3000", - "name": "frontend", + + "name": "frontend", }) // Find the registered ID + var id string + for k := range mgr.targets { id = k } result := tool.Execute(context.Background(), map[string]any{ "action": "unregister", - "id": id, + + "id": id, }) if result.IsError { t.Fatalf("unregister failed: %s", result.ForLLM) } + if len(mgr.targets) != 0 { t.Errorf("expected 0 targets, got %d", len(mgr.targets)) } + if mgr.active != "" { t.Errorf("expected no active target, got %q", mgr.active) } @@ -487,10 +600,12 @@ func TestDevPreviewTool_UnregisterActiveTarget(t *testing.T) { func TestDevPreviewTool_StartTargetEmptyString(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ "action": "start", + "target": "", }) @@ -501,8 +616,11 @@ func TestDevPreviewTool_StartTargetEmptyString(t *testing.T) { func TestDevPreviewTool_StatusActiveNoTargets(t *testing.T) { // Edge case: active proxy but no registered targets (shouldn't normally happen) + mgr := newMockManager() + mgr.active = "http://localhost:9999" // active but targets map is empty + tool := NewDevPreviewTool(mgr) result := tool.Execute(context.Background(), map[string]any{ @@ -512,12 +630,15 @@ func TestDevPreviewTool_StatusActiveNoTargets(t *testing.T) { if result.IsError { t.Fatalf("expected success, got error: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "active") { t.Errorf("expected 'active' in result, got %q", result.ForLLM) } + if !strings.Contains(result.ForLLM, "http://localhost:9999") { t.Errorf("expected target URL in result, got %q", result.ForLLM) } + if !strings.Contains(result.ForLLM, "No registered targets") { t.Errorf("expected 'No registered targets' in result, got %q", result.ForLLM) } @@ -525,10 +646,13 @@ func TestDevPreviewTool_StatusActiveNoTargets(t *testing.T) { func TestDevPreviewTool_StatusOutputFormat(t *testing.T) { mgr := newMockManager() + tool := NewDevPreviewTool(mgr) id1, _ := mgr.RegisterDevTarget("api", "http://localhost:8080") + mgr.RegisterDevTarget("frontend", "http://localhost:3000") + mgr.ActivateDevTarget(id1) result := tool.Execute(context.Background(), map[string]any{ @@ -538,15 +662,21 @@ func TestDevPreviewTool_StatusOutputFormat(t *testing.T) { if result.IsError { t.Fatalf("status failed: %s", result.ForLLM) } + // Should contain IDs in bracket format + if !strings.Contains(result.ForLLM, "["+id1+"]") { t.Errorf("expected [%s] in output, got %q", id1, result.ForLLM) } + // Should contain the arrow + if !strings.Contains(result.ForLLM, "→") { t.Errorf("expected arrow in output, got %q", result.ForLLM) } + // Should contain "Registered targets:" + if !strings.Contains(result.ForLLM, "Registered targets:") { t.Errorf("expected 'Registered targets:' header, got %q", result.ForLLM) } diff --git a/pkg/tools/edit.go b/pkg/tools/edit.go index 3447d6e96..7946fd2fa 100644 --- a/pkg/tools/edit.go +++ b/pkg/tools/edit.go @@ -9,19 +9,24 @@ import ( ) // EditFileTool edits a file by replacing old_text with new_text. + // The old_text must exist exactly in the file. + type EditFileTool struct { fs fileSystem } // NewEditFileTool creates a new EditFileTool with optional directory restriction. + func NewEditFileTool(workspace string, restrict bool) *EditFileTool { var fs fileSystem + if restrict { fs = &sandboxFs{workspace: workspace} } else { fs = &hostFs{} } + return &EditFileTool{fs: fs} } @@ -36,36 +41,46 @@ func (t *EditFileTool) Description() string { func (t *EditFileTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "path": map[string]any{ - "type": "string", + "type": "string", + "description": "The file path to edit", }, + "old_text": map[string]any{ - "type": "string", + "type": "string", + "description": "The exact text to find and replace", }, + "new_text": map[string]any{ - "type": "string", + "type": "string", + "description": "The text to replace with", }, }, + "required": []string{"path", "old_text", "new_text"}, } } func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { path, ok := args["path"].(string) + if !ok { return ErrorResult("path is required") } oldText, ok := args["old_text"].(string) + if !ok { return ErrorResult("old_text is required") } newText, ok := args["new_text"].(string) + if !ok { return ErrorResult("new_text is required") } @@ -73,6 +88,7 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe if err := editFile(resolveFS(ctx, t.fs, path), path, oldText, newText); err != nil { return ErrorResult(err.Error()) } + return SilentResult(fmt.Sprintf("File edited: %s", path)) } @@ -82,11 +98,13 @@ type AppendFileTool struct { func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool { var fs fileSystem + if restrict { fs = &sandboxFs{workspace: workspace} } else { fs = &hostFs{} } + return &AppendFileTool{fs: fs} } @@ -101,27 +119,34 @@ func (t *AppendFileTool) Description() string { func (t *AppendFileTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "path": map[string]any{ - "type": "string", + "type": "string", + "description": "The file path to append to", }, + "content": map[string]any{ - "type": "string", + "type": "string", + "description": "The content to append", }, }, + "required": []string{"path", "content"}, } } func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { path, ok := args["path"].(string) + if !ok { return ErrorResult("path is required") } content, ok := args["content"].(string) + if !ok { return ErrorResult("content is required") } @@ -129,11 +154,14 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool if err := appendFile(resolveFS(ctx, t.fs, path), path, content); err != nil { return ErrorResult(err.Error()) } + return SilentResult(fmt.Sprintf("Appended to %s", path)) } // editFile reads the file via sysFs, performs the replacement, and writes back. + // It uses a fileSystem interface, allowing the same logic for both restricted and unrestricted modes. + func editFile(sysFs fileSystem, path, oldText, newText string) error { content, err := sysFs.ReadFile(path) if err != nil { @@ -149,17 +177,21 @@ func editFile(sysFs fileSystem, path, oldText, newText string) error { } // appendFile reads the existing content (if any) via sysFs, appends new content, and writes back. + func appendFile(sysFs fileSystem, path, appendContent string) error { content, err := sysFs.ReadFile(path) + if err != nil && !errors.Is(err, fs.ErrNotExist) { return err } newContent := append(content, []byte(appendContent)...) + return sysFs.WriteFile(path, newContent) } // replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText. + func replaceEditContent(content []byte, oldText, newText string) ([]byte, error) { contentStr := string(content) @@ -168,10 +200,12 @@ func replaceEditContent(content []byte, oldText, newText string) ([]byte, error) } count := strings.Count(contentStr, oldText) + if count > 1 { return nil, fmt.Errorf("old_text appears %d times. Please provide more context to make it unique", count) } newContent := strings.Replace(contentStr, oldText, newText, 1) + return []byte(newContent), nil } diff --git a/pkg/tools/edit_test.go b/pkg/tools/edit_test.go index 65781ed1d..ccad894eb 100644 --- a/pkg/tools/edit_test.go +++ b/pkg/tools/edit_test.go @@ -11,261 +11,349 @@ import ( ) // TestEditTool_EditFile_Success verifies successful file editing + func TestEditTool_EditFile_Success(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "test.txt") + os.WriteFile(testFile, []byte("Hello World\nThis is a test"), 0o644) tool := NewEditFileTool(tmpDir, true) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "old_text": "World", + "new_text": "Universe", } result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // Should return SilentResult + if !result.Silent { t.Errorf("Expected Silent=true for EditFile, got false") } // ForUser should be empty (silent result) + if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) } // Verify file was actually edited + content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read edited file: %v", err) } + contentStr := string(content) + if !strings.Contains(contentStr, "Hello Universe") { t.Errorf("Expected file to contain 'Hello Universe', got: %s", contentStr) } + if strings.Contains(contentStr, "Hello World") { t.Errorf("Expected 'Hello World' to be replaced, got: %s", contentStr) } } // TestEditTool_EditFile_NotFound verifies error handling for non-existent file + func TestEditTool_EditFile_NotFound(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "nonexistent.txt") tool := NewEditFileTool(tmpDir, true) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "old_text": "old", + "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error for non-existent file") } // Should mention file not found + if !strings.Contains(result.ForLLM, "not found") && !strings.Contains(result.ForUser, "not found") { t.Errorf("Expected 'file not found' message, got ForLLM: %s", result.ForLLM) } } // TestEditTool_EditFile_OldTextNotFound verifies error when old_text doesn't exist + func TestEditTool_EditFile_OldTextNotFound(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "test.txt") + os.WriteFile(testFile, []byte("Hello World"), 0o644) tool := NewEditFileTool(tmpDir, true) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "old_text": "Goodbye", + "new_text": "Hello", } result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when old_text not found") } // Should mention old_text not found + if !strings.Contains(result.ForLLM, "not found") && !strings.Contains(result.ForUser, "not found") { t.Errorf("Expected 'not found' message, got ForLLM: %s", result.ForLLM) } } // TestEditTool_EditFile_MultipleMatches verifies error when old_text appears multiple times + func TestEditTool_EditFile_MultipleMatches(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "test.txt") + os.WriteFile(testFile, []byte("test test test"), 0o644) tool := NewEditFileTool(tmpDir, true) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "old_text": "test", + "new_text": "done", } result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when old_text appears multiple times") } // Should mention multiple occurrences + if !strings.Contains(result.ForLLM, "times") && !strings.Contains(result.ForUser, "times") { t.Errorf("Expected 'multiple times' message, got ForLLM: %s", result.ForLLM) } } // TestEditTool_EditFile_OutsideAllowedDir verifies error when path is outside allowed directory + func TestEditTool_EditFile_OutsideAllowedDir(t *testing.T) { tmpDir := t.TempDir() + otherDir := t.TempDir() + testFile := filepath.Join(otherDir, "test.txt") + os.WriteFile(testFile, []byte("content"), 0o644) tool := NewEditFileTool(tmpDir, true) // Restrict to tmpDir + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "old_text": "content", + "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result + assert.True(t, result.IsError, "Expected error when path is outside allowed directory") // Should mention outside allowed directory + // Note: ErrorResult only sets ForLLM by default, so ForUser might be empty. + // We check ForLLM as it's the primary error channel. + assert.True( + t, + strings.Contains(result.ForLLM, "outside") || strings.Contains(result.ForLLM, "access denied") || + strings.Contains(result.ForLLM, "escapes"), + "Expected 'outside allowed' or 'access denied' message, got ForLLM: %s", + result.ForLLM, ) } // TestEditTool_EditFile_MissingPath verifies error handling for missing path + func TestEditTool_EditFile_MissingPath(t *testing.T) { tool := NewEditFileTool("", false) + ctx := context.Background() + args := map[string]any{ "old_text": "old", + "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when path is missing") } } // TestEditTool_EditFile_MissingOldText verifies error handling for missing old_text + func TestEditTool_EditFile_MissingOldText(t *testing.T) { tool := NewEditFileTool("", false) + ctx := context.Background() + args := map[string]any{ - "path": "/tmp/test.txt", + "path": "/tmp/test.txt", + "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when old_text is missing") } } // TestEditTool_EditFile_MissingNewText verifies error handling for missing new_text + func TestEditTool_EditFile_MissingNewText(t *testing.T) { tool := NewEditFileTool("", false) + ctx := context.Background() + args := map[string]any{ - "path": "/tmp/test.txt", + "path": "/tmp/test.txt", + "old_text": "old", } result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when new_text is missing") } } // TestEditTool_AppendFile_Success verifies successful file appending + func TestEditTool_AppendFile_Success(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "test.txt") + os.WriteFile(testFile, []byte("Initial content"), 0o644) tool := NewAppendFileTool("", false) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "content": "\nAppended content", } result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // Should return SilentResult + if !result.Silent { t.Errorf("Expected Silent=true for AppendFile, got false") } // ForUser should be empty (silent result) + if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) } // Verify content was actually appended + content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read file: %v", err) } + contentStr := string(content) + if !strings.Contains(contentStr, "Initial content") { t.Errorf("Expected original content to remain, got: %s", contentStr) } + if !strings.Contains(contentStr, "Appended content") { t.Errorf("Expected appended content, got: %s", contentStr) } } // TestEditTool_AppendFile_MissingPath verifies error handling for missing path + func TestEditTool_AppendFile_MissingPath(t *testing.T) { tool := NewAppendFileTool("", false) + ctx := context.Background() + args := map[string]any{ "content": "test", } @@ -273,15 +361,19 @@ func TestEditTool_AppendFile_MissingPath(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when path is missing") } } // TestEditTool_AppendFile_MissingContent verifies error handling for missing content + func TestEditTool_AppendFile_MissingContent(t *testing.T) { tool := NewAppendFileTool("", false) + ctx := context.Background() + args := map[string]any{ "path": "/tmp/test.txt", } @@ -289,43 +381,67 @@ func TestEditTool_AppendFile_MissingContent(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when content is missing") } } // TestReplaceEditContent verifies the helper function replaceEditContent + func TestReplaceEditContent(t *testing.T) { tests := []struct { - name string - content []byte - oldText string - newText string - expected []byte + name string + + content []byte + + oldText string + + newText string + + expected []byte + expectError bool }{ { - name: "successful replacement", - content: []byte("hello world"), - oldText: "world", - newText: "universe", - expected: []byte("hello universe"), + name: "successful replacement", + + content: []byte("hello world"), + + oldText: "world", + + newText: "universe", + + expected: []byte("hello universe"), + expectError: false, }, + { - name: "old text not found", - content: []byte("hello world"), - oldText: "golang", - newText: "rust", - expected: nil, + name: "old text not found", + + content: []byte("hello world"), + + oldText: "golang", + + newText: "rust", + + expected: nil, + expectError: true, }, + { - name: "multiple matches found", - content: []byte("test text test"), - oldText: "test", - newText: "done", - expected: nil, + name: "multiple matches found", + + content: []byte("test text test"), + + oldText: "test", + + newText: "done", + + expected: nil, + expectError: true, }, } @@ -333,10 +449,12 @@ func TestReplaceEditContent(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result, err := replaceEditContent(tt.content, tt.oldText, tt.newText) + if tt.expectError { assert.Error(t, err) } else { assert.NoError(t, err) + assert.Equal(t, tt.expected, result) } }) @@ -344,94 +462,142 @@ func TestReplaceEditContent(t *testing.T) { } // TestAppendFileTool_AppendToNonExistent_Restricted verifies that AppendFileTool in restricted mode + // can append to a file that does not yet exist — it should silently create the file. + // This exercises the errors.Is(err, fs.ErrNotExist) path in appendFile + sandboxFs. + func TestAppendFileTool_AppendToNonExistent_Restricted(t *testing.T) { workspace := t.TempDir() + tool := NewAppendFileTool(workspace, true) + ctx := context.Background() args := map[string]any{ - "path": "brand_new_file.txt", + "path": "brand_new_file.txt", + "content": "first content", } result := tool.Execute(ctx, args) + assert.False( + t, + result.IsError, + "Expected success when appending to non-existent file in restricted mode, got: %s", + result.ForLLM, ) // Verify the file was created with correct content + data, err := os.ReadFile(filepath.Join(workspace, "brand_new_file.txt")) + assert.NoError(t, err) + assert.Equal(t, "first content", string(data)) } // TestAppendFileTool_Restricted_Success verifies that AppendFileTool in restricted mode + // correctly appends to an existing file within the sandbox. + func TestAppendFileTool_Restricted_Success(t *testing.T) { workspace := t.TempDir() + testFile := "existing.txt" + err := os.WriteFile(filepath.Join(workspace, testFile), []byte("initial"), 0o644) + assert.NoError(t, err) tool := NewAppendFileTool(workspace, true) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "content": " appended", } result := tool.Execute(ctx, args) + assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) + assert.True(t, result.Silent) data, err := os.ReadFile(filepath.Join(workspace, testFile)) + assert.NoError(t, err) + assert.Equal(t, "initial appended", string(data)) } // TestEditFileTool_Restricted_InPlaceEdit verifies that EditFileTool in restricted mode + // correctly edits a file using the sandboxFs path. + func TestEditFileTool_Restricted_InPlaceEdit(t *testing.T) { workspace := t.TempDir() + testFile := "edit_target.txt" + err := os.WriteFile(filepath.Join(workspace, testFile), []byte("Hello World"), 0o644) + assert.NoError(t, err) tool := NewEditFileTool(workspace, true) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "old_text": "World", + "new_text": "Go", } result := tool.Execute(ctx, args) + assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) + assert.True(t, result.Silent) data, err := os.ReadFile(filepath.Join(workspace, testFile)) + assert.NoError(t, err) + assert.Equal(t, "Hello Go", string(data)) } // TestEditFileTool_Restricted_FileNotFound verifies that editFile returns a proper + // error message when the target file does not exist. + func TestEditFileTool_Restricted_FileNotFound(t *testing.T) { workspace := t.TempDir() + tool := NewEditFileTool(workspace, true) + ctx := context.Background() + args := map[string]any{ - "path": "no_such_file.txt", + "path": "no_such_file.txt", + "old_text": "old", + "new_text": "new", } result := tool.Execute(ctx, args) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "not found") } diff --git a/pkg/tools/filesystem.go b/pkg/tools/filesystem.go index 072512b36..7b53c5e73 100644 --- a/pkg/tools/filesystem.go +++ b/pkg/tools/filesystem.go @@ -13,7 +13,9 @@ import ( ) // validatePath ensures the given path is within the workspace if restrict is true. + // Used by shell.go for working directory validation. + func validatePath(path, workspace string, restrict bool) (string, error) { if workspace == "" { return path, fmt.Errorf("workspace is not defined") @@ -25,6 +27,7 @@ func validatePath(path, workspace string, restrict bool) (string, error) { } var absPath string + if filepath.IsAbs(path) { absPath = filepath.Clean(path) } else { @@ -40,7 +43,9 @@ func validatePath(path, workspace string, restrict bool) (string, error) { } var resolved string + workspaceReal := absWorkspace + if resolved, err = filepath.EvalSymlinks(absWorkspace); err == nil { workspaceReal = resolved } @@ -51,6 +56,7 @@ func validatePath(path, workspace string, restrict bool) (string, error) { } } else if os.IsNotExist(err) { var parentResolved string + if parentResolved, err = resolveExistingAncestor(filepath.Dir(absPath)); err == nil { if !isWithinWorkspace(parentResolved, workspaceReal) { return "", fmt.Errorf("access denied: symlink resolves outside workspace") @@ -73,6 +79,7 @@ func resolveExistingAncestor(path string) (string, error) { } else if !os.IsNotExist(err) { return "", err } + if filepath.Dir(current) == current { return "", os.ErrNotExist } @@ -81,6 +88,7 @@ func resolveExistingAncestor(path string) (string, error) { func isWithinWorkspace(candidate, workspace string) bool { rel, err := filepath.Rel(filepath.Clean(workspace), filepath.Clean(candidate)) + return err == nil && filepath.IsLocal(rel) } @@ -90,11 +98,13 @@ type ReadFileTool struct { func NewReadFileTool(workspace string, restrict bool) *ReadFileTool { var fs fileSystem + if restrict { fs = &sandboxFs{workspace: workspace} } else { fs = &hostFs{} } + return &ReadFileTool{fs: fs} } @@ -109,18 +119,22 @@ func (t *ReadFileTool) Description() string { func (t *ReadFileTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "path": map[string]any{ - "type": "string", + "type": "string", + "description": "Path to the file to read", }, }, + "required": []string{"path"}, } } func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { path, ok := args["path"].(string) + if !ok { return ErrorResult("path is required") } @@ -129,6 +143,7 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe if err != nil { return ErrorResult(err.Error()) } + return NewToolResult(string(content)) } @@ -138,11 +153,13 @@ type WriteFileTool struct { func NewWriteFileTool(workspace string, restrict bool) *WriteFileTool { var fs fileSystem + if restrict { fs = &sandboxFs{workspace: workspace} } else { fs = &hostFs{} } + return &WriteFileTool{fs: fs} } @@ -157,27 +174,34 @@ func (t *WriteFileTool) Description() string { func (t *WriteFileTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "path": map[string]any{ - "type": "string", + "type": "string", + "description": "Path to the file to write", }, + "content": map[string]any{ - "type": "string", + "type": "string", + "description": "Content to write to the file", }, }, + "required": []string{"path", "content"}, } } func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { path, ok := args["path"].(string) + if !ok { return ErrorResult("path is required") } content, ok := args["content"].(string) + if !ok { return ErrorResult("content is required") } @@ -195,11 +219,13 @@ type ListDirTool struct { func NewListDirTool(workspace string, restrict bool) *ListDirTool { var fs fileSystem + if restrict { fs = &sandboxFs{workspace: workspace} } else { fs = &hostFs{} } + return &ListDirTool{fs: fs} } @@ -214,18 +240,22 @@ func (t *ListDirTool) Description() string { func (t *ListDirTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "path": map[string]any{ - "type": "string", + "type": "string", + "description": "Path to list", }, }, + "required": []string{"path"}, } } func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolResult { path, ok := args["path"].(string) + if !ok { path = "." } @@ -234,32 +264,42 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolRes if err != nil { return ErrorResult(err.Error()) } + return formatDirEntries(entries) } func formatDirEntries(entries []os.DirEntry) *ToolResult { var result strings.Builder + for _, entry := range entries { if entry.IsDir() { result.WriteString("DIR: ") } else { result.WriteString("FILE: ") } + result.WriteString(entry.Name()) + result.WriteByte('\n') } + return NewToolResult(result.String()) } // fileSystem abstracts reading, writing, and listing files, allowing both + // unrestricted (host filesystem) and sandbox (os.Root) implementations to share the same polymorphic interface. + type fileSystem interface { ReadFile(path string) ([]byte, error) + WriteFile(path string, data []byte) error + ReadDir(path string) ([]os.DirEntry, error) } // hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem. + type hostFs struct{} func (h *hostFs) ReadFile(path string) ([]byte, error) { @@ -268,11 +308,14 @@ func (h *hostFs) ReadFile(path string) ([]byte, error) { if os.IsNotExist(err) { return nil, fmt.Errorf("failed to read file: file not found: %w", err) } + if os.IsPermission(err) { return nil, fmt.Errorf("failed to read file: access denied: %w", err) } + return nil, fmt.Errorf("failed to read file: %w", err) } + return content, nil } @@ -281,16 +324,20 @@ func (h *hostFs) ReadDir(path string) ([]os.DirEntry, error) { if err != nil { return nil, fmt.Errorf("failed to read directory: %w", err) } + return entries, nil } func (h *hostFs) WriteFile(path string, data []byte) error { // Use unified atomic write utility with explicit sync for flash storage reliability. + // Using 0o600 (owner read/write only) for secure default permissions. + return fileutil.WriteFileAtomic(path, data, 0o600) } // sandboxFs is a sandboxed fileSystem that operates within a strictly defined workspace using os.Root. + type sandboxFs struct { workspace string } @@ -304,6 +351,7 @@ func (r *sandboxFs) execute(path string, fn func(root *os.Root, relPath string) if err != nil { return fmt.Errorf("failed to open workspace: %w", err) } + defer root.Close() relPath, err := getSafeRelPath(r.workspace, path) @@ -316,28 +364,37 @@ func (r *sandboxFs) execute(path string, fn func(root *os.Root, relPath string) func (r *sandboxFs) ReadFile(path string) ([]byte, error) { var content []byte + err := r.execute(path, func(root *os.Root, relPath string) error { fileContent, err := root.ReadFile(relPath) if err != nil { if os.IsNotExist(err) { return fmt.Errorf("failed to read file: file not found: %w", err) } + // os.Root returns "escapes from parent" for paths outside the root + if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") || + strings.Contains(err.Error(), "permission denied") { return fmt.Errorf("failed to read file: access denied: %w", err) } + return fmt.Errorf("failed to read file: %w", err) } + content = fileContent + return nil }) + return content, err } func (r *sandboxFs) WriteFile(path string, data []byte) error { return r.execute(path, func(root *os.Root, relPath string) error { dir := filepath.Dir(relPath) + if dir != "." && dir != "/" { if err := root.MkdirAll(dir, 0o755); err != nil { return fmt.Errorf("failed to create parent directories: %w", err) @@ -345,42 +402,55 @@ func (r *sandboxFs) WriteFile(path string, data []byte) error { } // Use atomic write pattern with explicit sync for flash storage reliability. + // Using 0o600 (owner read/write only) for secure default permissions. + tmpRelPath := fmt.Sprintf(".tmp-%d-%d", os.Getpid(), time.Now().UnixNano()) tmpFile, err := root.OpenFile(tmpRelPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) if err != nil { root.Remove(tmpRelPath) + return fmt.Errorf("failed to open temp file: %w", err) } if _, err := tmpFile.Write(data); err != nil { tmpFile.Close() + root.Remove(tmpRelPath) + return fmt.Errorf("failed to write temp file: %w", err) } // CRITICAL: Force sync to storage medium before rename. + // This ensures data is physically written to disk, not just cached. + if err := tmpFile.Sync(); err != nil { tmpFile.Close() + root.Remove(tmpRelPath) + return fmt.Errorf("failed to sync temp file: %w", err) } if err := tmpFile.Close(); err != nil { root.Remove(tmpRelPath) + return fmt.Errorf("failed to close temp file: %w", err) } if err := root.Rename(tmpRelPath, relPath); err != nil { root.Remove(tmpRelPath) + return fmt.Errorf("failed to rename temp file over target: %w", err) } // Sync directory to ensure rename is durable + if dirFile, err := root.Open("."); err == nil { _ = dirFile.Sync() + dirFile.Close() } @@ -390,26 +460,33 @@ func (r *sandboxFs) WriteFile(path string, data []byte) error { func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) { var entries []os.DirEntry + err := r.execute(path, func(root *os.Root, relPath string) error { dirEntries, err := fs.ReadDir(root.FS(), relPath) if err != nil { return err } + entries = dirEntries + return nil }) + return entries, err } // Helper to get a safe relative path for os.Root usage + func getSafeRelPath(workspace, path string) (string, error) { if workspace == "" { return "", fmt.Errorf("workspace is not defined") } rel := filepath.Clean(path) + if filepath.IsAbs(rel) { var err error + rel, err = filepath.Rel(workspace, rel) if err != nil { return "", fmt.Errorf("failed to calculate relative path: %w", err) diff --git a/pkg/tools/filesystem_test.go b/pkg/tools/filesystem_test.go index b3bc2affe..456f8fbd3 100644 --- a/pkg/tools/filesystem_test.go +++ b/pkg/tools/filesystem_test.go @@ -12,13 +12,18 @@ import ( ) // TestFilesystemTool_ReadFile_Success verifies successful file reading + func TestFilesystemTool_ReadFile_Success(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "test.txt") + os.WriteFile(testFile, []byte("test content"), 0o644) tool := NewReadFileTool("", false) + ctx := context.Background() + args := map[string]any{ "path": testFile, } @@ -26,26 +31,33 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForLLM should contain file content + if !strings.Contains(result.ForLLM, "test content") { t.Errorf("Expected ForLLM to contain 'test content', got: %s", result.ForLLM) } // ReadFile returns NewToolResult which only sets ForLLM, not ForUser + // This is the expected behavior - file content goes to LLM, not directly to user + if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for NewToolResult, got: %s", result.ForUser) } } // TestFilesystemTool_ReadFile_NotFound verifies error handling for missing file + func TestFilesystemTool_ReadFile_NotFound(t *testing.T) { tool := NewReadFileTool("", false) + ctx := context.Background() + args := map[string]any{ "path": "/nonexistent_file_12345.txt", } @@ -53,107 +65,135 @@ func TestFilesystemTool_ReadFile_NotFound(t *testing.T) { result := tool.Execute(ctx, args) // Failure should be marked as error + if !result.IsError { t.Errorf("Expected error for missing file, got IsError=false") } // Should contain error message + if !strings.Contains(result.ForLLM, "failed to read") && !strings.Contains(result.ForUser, "failed to read") { t.Errorf("Expected error message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser) } } // TestFilesystemTool_ReadFile_MissingPath verifies error handling for missing path + func TestFilesystemTool_ReadFile_MissingPath(t *testing.T) { tool := &ReadFileTool{} + ctx := context.Background() + args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when path is missing") } // Should mention required parameter + if !strings.Contains(result.ForLLM, "path is required") && !strings.Contains(result.ForUser, "path is required") { t.Errorf("Expected 'path is required' message, got ForLLM: %s", result.ForLLM) } } // TestFilesystemTool_WriteFile_Success verifies successful file writing + func TestFilesystemTool_WriteFile_Success(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "newfile.txt") tool := NewWriteFileTool("", false) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "content": "hello world", } result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // WriteFile returns SilentResult + if !result.Silent { t.Errorf("Expected Silent=true for WriteFile, got false") } // ForUser should be empty (silent result) + if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) } // Verify file was actually written + content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read written file: %v", err) } + if string(content) != "hello world" { t.Errorf("Expected file content 'hello world', got: %s", string(content)) } } // TestFilesystemTool_WriteFile_CreateDir verifies directory creation + func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "subdir", "newfile.txt") tool := NewWriteFileTool("", false) + ctx := context.Background() + args := map[string]any{ - "path": testFile, + "path": testFile, + "content": "test", } result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success with directory creation, got IsError=true: %s", result.ForLLM) } // Verify directory was created and file written + content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read written file: %v", err) } + if string(content) != "test" { t.Errorf("Expected file content 'test', got: %s", string(content)) } } // TestFilesystemTool_WriteFile_MissingPath verifies error handling for missing path + func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) { tool := NewWriteFileTool("", false) + ctx := context.Background() + args := map[string]any{ "content": "test", } @@ -161,15 +201,19 @@ func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when path is missing") } } // TestFilesystemTool_WriteFile_MissingContent verifies error handling for missing content + func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) { tool := NewWriteFileTool("", false) + ctx := context.Background() + args := map[string]any{ "path": "/tmp/test.txt", } @@ -177,26 +221,35 @@ func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when content is missing") } // Should mention required parameter + if !strings.Contains(result.ForLLM, "content is required") && + !strings.Contains(result.ForUser, "content is required") { t.Errorf("Expected 'content is required' message, got ForLLM: %s", result.ForLLM) } } // TestFilesystemTool_ListDir_Success verifies successful directory listing + func TestFilesystemTool_ListDir_Success(t *testing.T) { tmpDir := t.TempDir() + os.WriteFile(filepath.Join(tmpDir, "file1.txt"), []byte("content"), 0o644) + os.WriteFile(filepath.Join(tmpDir, "file2.txt"), []byte("content"), 0o644) + os.Mkdir(filepath.Join(tmpDir, "subdir"), 0o755) tool := NewListDirTool("", false) + ctx := context.Background() + args := map[string]any{ "path": tmpDir, } @@ -204,23 +257,29 @@ func TestFilesystemTool_ListDir_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // Should list files and directories + if !strings.Contains(result.ForLLM, "file1.txt") || !strings.Contains(result.ForLLM, "file2.txt") { t.Errorf("Expected files in listing, got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "subdir") { t.Errorf("Expected subdir in listing, got: %s", result.ForLLM) } } // TestFilesystemTool_ListDir_NotFound verifies error handling for non-existent directory + func TestFilesystemTool_ListDir_NotFound(t *testing.T) { tool := NewListDirTool("", false) + ctx := context.Background() + args := map[string]any{ "path": "/nonexistent_directory_12345", } @@ -228,49 +287,61 @@ func TestFilesystemTool_ListDir_NotFound(t *testing.T) { result := tool.Execute(ctx, args) // Failure should be marked as error + if !result.IsError { t.Errorf("Expected error for non-existent directory, got IsError=false") } // Should contain error message + if !strings.Contains(result.ForLLM, "failed to read") && !strings.Contains(result.ForUser, "failed to read") { t.Errorf("Expected error message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser) } } // TestFilesystemTool_ListDir_DefaultPath verifies default to current directory + func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) { tool := NewListDirTool("", false) + ctx := context.Background() + args := map[string]any{} result := tool.Execute(ctx, args) // Should use "." as default path + if result.IsError { t.Errorf("Expected success with default path '.', got IsError=true: %s", result.ForLLM) } } // Block paths that look inside workspace but point outside via symlink. + func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) { root := t.TempDir() + workspace := filepath.Join(root, "workspace") + if err := os.MkdirAll(workspace, 0o755); err != nil { t.Fatalf("failed to create workspace: %v", err) } secret := filepath.Join(root, "secret.txt") + if err := os.WriteFile(secret, []byte("top secret"), 0o644); err != nil { t.Fatalf("failed to write secret file: %v", err) } link := filepath.Join(workspace, "leak.txt") + if err := os.Symlink(secret, link); err != nil { t.Skipf("symlink not supported in this environment: %v", err) } tool := NewReadFileTool(workspace, true) + result := tool.Execute(context.Background(), map[string]any{ "path": link, }) @@ -278,10 +349,15 @@ func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) { if !result.IsError { t.Fatalf("expected symlink escape to be blocked") } + // os.Root might return different errors depending on platform/implementation + // but it definitely should error. + // Our wrapper returns "access denied or file not found" + if !strings.Contains(result.ForLLM, "access denied") && !strings.Contains(result.ForLLM, "file not found") && + !strings.Contains(result.ForLLM, "no such file") { t.Fatalf("expected symlink escape error, got: %s", result.ForLLM) } @@ -291,8 +367,11 @@ func TestFilesystemTool_EmptyWorkspace_AccessDenied(t *testing.T) { tool := NewReadFileTool("", true) // restrict=true but workspace="" // Try to read a sensitive file (simulated by a temp file outside workspace) + tmpDir := t.TempDir() + secretFile := filepath.Join(tmpDir, "shadow") + os.WriteFile(secretFile, []byte("secret data"), 0o600) result := tool.Execute(context.Background(), map[string]any{ @@ -300,201 +379,293 @@ func TestFilesystemTool_EmptyWorkspace_AccessDenied(t *testing.T) { }) // We EXPECT IsError=true (access blocked due to empty workspace) + assert.True(t, result.IsError, "Security Regression: Empty workspace allowed access! content: %s", result.ForLLM) // Verify it failed for the right reason + assert.Contains(t, result.ForLLM, "workspace is not defined", "Expected 'workspace is not defined' error") } // TestRootMkdirAll verifies that root.MkdirAll (used by sandboxFs.WriteFile) handles all cases: + // single dir, deeply nested dirs, already-existing dirs, and a file blocking a directory path. + func TestRootMkdirAll(t *testing.T) { workspace := t.TempDir() + root, err := os.OpenRoot(workspace) if err != nil { t.Fatalf("failed to open root: %v", err) } + defer root.Close() // Case 1: Single directory + err = root.MkdirAll("dir1", 0o755) + assert.NoError(t, err) + _, err = os.Stat(filepath.Join(workspace, "dir1")) + assert.NoError(t, err) // Case 2: Deeply nested directory + err = root.MkdirAll("a/b/c/d", 0o755) + assert.NoError(t, err) + _, err = os.Stat(filepath.Join(workspace, "a/b/c/d")) + assert.NoError(t, err) // Case 3: Already exists — must be idempotent + err = root.MkdirAll("a/b/c/d", 0o755) + assert.NoError(t, err) // Case 4: A regular file blocks directory creation — must error + err = os.WriteFile(filepath.Join(workspace, "file_exists"), []byte("data"), 0o644) + assert.NoError(t, err) + err = root.MkdirAll("file_exists", 0o755) + assert.Error(t, err, "expected error when a file exists at the directory path") } func TestFilesystemTool_WriteFile_Restricted_CreateDir(t *testing.T) { workspace := t.TempDir() + tool := NewWriteFileTool(workspace, true) + ctx := context.Background() testFile := "deep/nested/path/to/file.txt" + content := "deep content" + args := map[string]any{ - "path": testFile, + "path": testFile, + "content": content, } result := tool.Execute(ctx, args) + assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) // Verify file content + actualPath := filepath.Join(workspace, testFile) + data, err := os.ReadFile(actualPath) + assert.NoError(t, err) + assert.Equal(t, content, string(data)) } // TestHostFs_Read_PermissionDenied verifies that hostFs.ReadFile surfaces access denied errors. + func TestHostFs_Read_PermissionDenied(t *testing.T) { if os.Getuid() == 0 { t.Skip("skipping permission test: running as root") } + tmpDir := t.TempDir() + protected := filepath.Join(tmpDir, "protected.txt") + err := os.WriteFile(protected, []byte("secret"), 0o000) + assert.NoError(t, err) + defer os.Chmod(protected, 0o644) // ensure cleanup _, err = (&hostFs{}).ReadFile(protected) + assert.Error(t, err) + assert.Contains(t, err.Error(), "access denied") } // TestHostFs_Read_Directory verifies that hostFs.ReadFile returns an error when given a directory path. + func TestHostFs_Read_Directory(t *testing.T) { tmpDir := t.TempDir() _, err := (&hostFs{}).ReadFile(tmpDir) + assert.Error(t, err, "expected error when reading a directory as a file") } // TestSandboxFs_Read_Directory verifies that sandboxFs.ReadFile returns an error when given a directory. + func TestSandboxFs_Read_Directory(t *testing.T) { workspace := t.TempDir() + root, err := os.OpenRoot(workspace) + assert.NoError(t, err) + defer root.Close() // Create a subdirectory + err = root.Mkdir("subdir", 0o755) + assert.NoError(t, err) _, err = (&sandboxFs{workspace: workspace}).ReadFile("subdir") + assert.Error(t, err, "expected error when reading a directory as a file") } // TestHostFs_Write_ParentDirMissing verifies that hostFs.WriteFile creates parent dirs automatically. + func TestHostFs_Write_ParentDirMissing(t *testing.T) { tmpDir := t.TempDir() + target := filepath.Join(tmpDir, "a", "b", "c", "file.txt") err := (&hostFs{}).WriteFile(target, []byte("hello")) + assert.NoError(t, err) data, err := os.ReadFile(target) + assert.NoError(t, err) + assert.Equal(t, "hello", string(data)) } // TestSandboxFs_Write_ParentDirMissing verifies that sandboxFs.WriteFile creates + // nested parent directories automatically within the sandbox. + func TestSandboxFs_Write_ParentDirMissing(t *testing.T) { workspace := t.TempDir() relPath := "x/y/z/file.txt" + err := (&sandboxFs{workspace: workspace}).WriteFile(relPath, []byte("nested")) + assert.NoError(t, err) data, err := os.ReadFile(filepath.Join(workspace, relPath)) + assert.NoError(t, err) + assert.Equal(t, "nested", string(data)) } // TestHostFs_Write verifies the hostFs.WriteFile helper function + func TestHostFs_Write(t *testing.T) { tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "atomic_test.txt") + testData := []byte("atomic test content") err := (&hostFs{}).WriteFile(testFile, testData) + assert.NoError(t, err) content, err := os.ReadFile(testFile) + assert.NoError(t, err) + assert.Equal(t, testData, content) // Verify it overwrites correctly + newData := []byte("new atomic content") + err = (&hostFs{}).WriteFile(testFile, newData) + assert.NoError(t, err) content, err = os.ReadFile(testFile) + assert.NoError(t, err) + assert.Equal(t, newData, content) } // TestSandboxFs_Write verifies the sandboxFs.WriteFile helper function + func TestSandboxFs_Write(t *testing.T) { tmpDir := t.TempDir() relPath := "atomic_root_test.txt" + testData := []byte("atomic root test content") erw := &sandboxFs{workspace: tmpDir} + err := erw.WriteFile(relPath, testData) + assert.NoError(t, err) root, err := os.OpenRoot(tmpDir) + assert.NoError(t, err) + defer root.Close() f, err := root.Open(relPath) + assert.NoError(t, err) + defer f.Close() content, err := io.ReadAll(f) + assert.NoError(t, err) + assert.Equal(t, testData, content) // Verify it overwrites correctly + newData := []byte("new root atomic content") + err = erw.WriteFile(relPath, newData) + assert.NoError(t, err) f2, err := root.Open(relPath) + assert.NoError(t, err) + defer f2.Close() content, err = io.ReadAll(f2) + assert.NoError(t, err) + assert.Equal(t, newData, content) } // TestValidatePath_OutsideWorkspace_IncludesPath verifies that the access + // denied error includes the workspace path so the caller knows the boundary. + func TestValidatePath_OutsideWorkspace_IncludesPath(t *testing.T) { workspace := t.TempDir() + outsidePath := filepath.Join(t.TempDir(), "secret.txt") _, err := validatePath(outsidePath, workspace, true) + assert.Error(t, err) + assert.Contains(t, err.Error(), "access denied") + assert.Contains(t, err.Error(), workspace) } diff --git a/pkg/tools/gitpush.go b/pkg/tools/gitpush.go index 3d8bc2b9b..8e9055ed4 100644 --- a/pkg/tools/gitpush.go +++ b/pkg/tools/gitpush.go @@ -12,35 +12,49 @@ import ( ) // worktreeInfoKey is the context key for passing WorktreeInfo to tools. + type worktreeInfoKey struct{} // WithWorktreeInfo returns a context carrying the active WorktreeInfo. + func WithWorktreeInfo(ctx context.Context, wt *git.WorktreeInfo) context.Context { return context.WithValue(ctx, worktreeInfoKey{}, wt) } // WorktreeInfoFromCtx extracts the WorktreeInfo from context, or nil. + func WorktreeInfoFromCtx(ctx context.Context) *git.WorktreeInfo { if v, ok := ctx.Value(worktreeInfoKey{}).(*git.WorktreeInfo); ok { return v } + return nil } // protectedBranches are branch names that can never be pushed to. + var protectedBranches = regexp.MustCompile(`^(main|master|develop|release/.*)$`) // GitPushTool implements safe git push restricted to worktree branches. + // + // Safety invariants: + // - Only works inside a worktree (WorktreeInfo must be in context) + // - Pushes only the worktree's branch — no arbitrary branch targets + // - Protected branches (main, master, develop, release/*) are blocked + // - Force push is never allowed + // - Auto-commits uncommitted changes before pushing + type GitPushTool struct{} // NewGitPushTool creates a GitPushTool. + func NewGitPushTool() *GitPushTool { return &GitPushTool{} } @@ -49,94 +63,138 @@ func (t *GitPushTool) Name() string { return "git_push" } func (t *GitPushTool) Description() string { return "Push the current worktree branch to origin. Only works inside a git worktree. " + + "Auto-commits uncommitted changes before pushing. " + + "Protected branches (main, master, develop) cannot be pushed to. Force push is not allowed." } func (t *GitPushTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "commit_message": map[string]any{ - "type": "string", + "type": "string", + "description": "Commit message for uncommitted changes. If omitted, uncommitted changes are auto-committed with a default message.", }, }, + "required": []string{}, } } func (t *GitPushTool) Execute(ctx context.Context, args map[string]any) *ToolResult { wt := WorktreeInfoFromCtx(ctx) + if wt == nil { return ErrorResult( + "git_push requires an active worktree.\n" + + "This tool can only be used during worktree-based sessions " + + "(e.g., heartbeat tasks or plan executing phase).\n" + + "The worktree provides the branch name and isolation boundary — " + + "without it, git_push cannot determine which branch to push.") } branch := wt.Branch + if branch == "" { return ErrorResult( + "worktree has no branch name.\n" + + "The WorktreeInfo was set but Branch is empty. " + + "This is an internal error — the worktree may not have been created correctly.") } // Block protected branches + if protectedBranches.MatchString(branch) { return ErrorResult(fmt.Sprintf( + "cannot push to protected branch %q.\n"+ + "Protected branches (main, master, develop, release/*) are blocked to prevent "+ + "accidental overwrites. Work should be done on feature branches created by worktrees.", + branch)) } // Auto-commit uncommitted changes + if git.HasUncommittedChanges(wt.Path) { commitMsg := "auto: save before push" + if msg, ok := args["commit_message"].(string); ok && msg != "" { commitMsg = msg } + if err := git.AutoCommit(wt.Path, commitMsg); err != nil { return ErrorResult(fmt.Sprintf( + "auto-commit failed before push: %v\n"+ + "git_push auto-commits uncommitted changes before pushing. "+ + "The commit failed, so no push was attempted. "+ + "Check if the worktree at %q is in a valid state (e.g., no merge conflicts).", + err, wt.Path)) } } // Check there are commits to push + ahead := git.CommitsAhead(wt.RepoRoot, wt.BaseBranch, branch) + if ahead == 0 { return NewToolResult(fmt.Sprintf( + "Nothing to push: branch %q has no commits ahead of %s.\n"+ + "The branch is identical to the base. Make changes and commit before pushing.", + branch, wt.BaseBranch)) } // Push with -u (set upstream tracking) + pushCtx, cancel := context.WithTimeout(ctx, 2*time.Minute) + defer cancel() cmd := exec.CommandContext(pushCtx, "git", "push", "-u", "origin", branch) + cmd.Dir = wt.Path + out, err := cmd.CombinedOutput() + output := strings.TrimSpace(string(out)) if err != nil { return ErrorResult(fmt.Sprintf( + "git push failed for branch %q: %s\n%s\n"+ + "Possible causes: network error, authentication failure, or remote rejected the push. "+ + "If the remote branch has diverged, resolve the divergence in the worktree first — "+ + "force push is not available.", + branch, err, output)) } return NewToolResult(fmt.Sprintf("Pushed branch %q to origin (%d commit(s) ahead of %s)\n%s", + branch, ahead, wt.BaseBranch, output)) } diff --git a/pkg/tools/gitpush_test.go b/pkg/tools/gitpush_test.go index d16c446d4..a1aaa28bb 100644 --- a/pkg/tools/gitpush_test.go +++ b/pkg/tools/gitpush_test.go @@ -9,22 +9,29 @@ import ( ) // TestGitPushTool_NoWorktree verifies that git_push fails without worktree context. + func TestGitPushTool_NoWorktree(t *testing.T) { tool := NewGitPushTool() result := tool.Execute(context.Background(), map[string]any{}) + if !result.IsError { t.Fatal("expected error when no worktree in context") } + if result.ForLLM == "" { t.Fatal("error message should not be empty") } + // Verify helpful guidance is included + assertContains(t, result.ForLLM, "worktree") + assertContains(t, result.ForLLM, "heartbeat") } // TestGitPushTool_ProtectedBranch verifies that protected branches are blocked. + func TestGitPushTool_ProtectedBranch(t *testing.T) { tool := NewGitPushTool() @@ -33,40 +40,56 @@ func TestGitPushTool_ProtectedBranch(t *testing.T) { for _, branch := range protectedNames { t.Run(branch, func(t *testing.T) { ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ - Branch: branch, + Branch: branch, + BaseBranch: "main", - Path: t.TempDir(), - RepoRoot: t.TempDir(), + + Path: t.TempDir(), + + RepoRoot: t.TempDir(), }) + result := tool.Execute(ctx, map[string]any{}) + if !result.IsError { t.Fatalf("expected error for protected branch %q", branch) } + assertContains(t, result.ForLLM, "protected") + assertContains(t, result.ForLLM, branch) }) } } // TestGitPushTool_EmptyBranch verifies that empty branch name is rejected. + func TestGitPushTool_EmptyBranch(t *testing.T) { tool := NewGitPushTool() ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ - Branch: "", + Branch: "", + BaseBranch: "main", - Path: t.TempDir(), - RepoRoot: t.TempDir(), + + Path: t.TempDir(), + + RepoRoot: t.TempDir(), }) + result := tool.Execute(ctx, map[string]any{}) + if !result.IsError { t.Fatal("expected error for empty branch") } + assertContains(t, result.ForLLM, "no branch name") } // TestGitPushTool_AllowedBranch verifies that non-protected branches pass the branch check. + // (Push itself will fail because there's no real git repo, but it should get past validation.) + func TestGitPushTool_AllowedBranch(t *testing.T) { tool := NewGitPushTool() @@ -75,13 +98,19 @@ func TestGitPushTool_AllowedBranch(t *testing.T) { for _, branch := range allowedNames { t.Run(branch, func(t *testing.T) { ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ - Branch: branch, + Branch: branch, + BaseBranch: "main", - Path: t.TempDir(), - RepoRoot: t.TempDir(), + + Path: t.TempDir(), + + RepoRoot: t.TempDir(), }) + result := tool.Execute(ctx, map[string]any{}) + // Should NOT fail with "protected branch" error + if result.IsError && strings.Contains(result.ForLLM, "protected") { t.Fatalf("branch %q should not be blocked as protected", branch) } @@ -90,25 +119,36 @@ func TestGitPushTool_AllowedBranch(t *testing.T) { } // TestProtectedBranchesRegex tests the regex directly. + func TestProtectedBranchesRegex(t *testing.T) { tests := []struct { - branch string + branch string + protected bool }{ {"main", true}, + {"master", true}, + {"develop", true}, + {"release/v1.0", true}, + {"release/2026-03", true}, + {"plan/add-feature", false}, + {"feature/main", false}, // "main" not at start - {"main-backup", false}, // "main" followed by suffix + + {"main-backup", false}, // "main" followed by suffix + {"hotfix/urgent", false}, } for _, tt := range tests { t.Run(tt.branch, func(t *testing.T) { got := protectedBranches.MatchString(tt.branch) + if got != tt.protected { t.Errorf("branch %q: got protected=%v, want %v", tt.branch, got, tt.protected) } @@ -117,48 +157,64 @@ func TestProtectedBranchesRegex(t *testing.T) { } // TestWorktreeInfoContext verifies context round-trip. + func TestWorktreeInfoContext(t *testing.T) { wt := &git.WorktreeInfo{ - Branch: "plan/test", + Branch: "plan/test", + BaseBranch: "main", - Path: "/tmp/wt", - RepoRoot: "/tmp/repo", + + Path: "/tmp/wt", + + RepoRoot: "/tmp/repo", } ctx := WithWorktreeInfo(context.Background(), wt) + got := WorktreeInfoFromCtx(ctx) + if got == nil { t.Fatal("expected non-nil WorktreeInfo from context") } + if got.Branch != wt.Branch { t.Errorf("Branch: got %q, want %q", got.Branch, wt.Branch) } + if got.BaseBranch != wt.BaseBranch { t.Errorf("BaseBranch: got %q, want %q", got.BaseBranch, wt.BaseBranch) } // Nil case + got2 := WorktreeInfoFromCtx(context.Background()) + if got2 != nil { t.Errorf("expected nil WorktreeInfo from bare context, got %+v", got2) } } // TestGitPushTool_Interface verifies the tool satisfies the Tool interface. + func TestGitPushTool_Interface(t *testing.T) { var _ Tool = (*GitPushTool)(nil) tool := NewGitPushTool() + if tool.Name() != "git_push" { t.Errorf("Name: got %q, want %q", tool.Name(), "git_push") } + if tool.Description() == "" { t.Error("Description should not be empty") } + params := tool.Parameters() + if params == nil { t.Fatal("Parameters should not be nil") } + if params["type"] != "object" { t.Errorf("Parameters type: got %v, want object", params["type"]) } @@ -166,6 +222,7 @@ func TestGitPushTool_Interface(t *testing.T) { func assertContains(t *testing.T, s, substr string) { t.Helper() + if !strings.Contains(s, substr) { t.Errorf("expected %q to contain %q", s, substr) } diff --git a/pkg/tools/i2c.go b/pkg/tools/i2c.go index 779b1d5a7..d3d7ebe10 100644 --- a/pkg/tools/i2c.go +++ b/pkg/tools/i2c.go @@ -10,6 +10,7 @@ import ( ) // I2CTool provides I2C bus interaction for reading sensors and controlling peripherals. + type I2CTool struct{} func NewI2CTool() *I2CTool { @@ -27,38 +28,55 @@ func (t *I2CTool) Description() string { func (t *I2CTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "action": map[string]any{ - "type": "string", - "enum": []string{"detect", "scan", "read", "write"}, + "type": "string", + + "enum": []string{"detect", "scan", "read", "write"}, + "description": "Action to perform: detect (list available I2C buses), scan (find devices on a bus), read (read bytes from a device), write (send bytes to a device)", }, + "bus": map[string]any{ - "type": "string", + "type": "string", + "description": "I2C bus number (e.g. \"1\" for /dev/i2c-1). Required for scan/read/write.", }, + "address": map[string]any{ - "type": "integer", + "type": "integer", + "description": "7-bit I2C device address (0x03-0x77). Required for read/write.", }, + "register": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Register address to read from or write to. If set, sends register byte before read/write.", }, + "data": map[string]any{ - "type": "array", - "items": map[string]any{"type": "integer"}, + "type": "array", + + "items": map[string]any{"type": "integer"}, + "description": "Bytes to write (0-255 each). Required for write action.", }, + "length": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Number of bytes to read (1-256). Default: 1. Used with read action.", }, + "confirm": map[string]any{ - "type": "boolean", + "type": "boolean", + "description": "Must be true for write operations. Safety guard to prevent accidental writes.", }, }, + "required": []string{"action"}, } } @@ -69,25 +87,36 @@ func (t *I2CTool) Execute(ctx context.Context, args map[string]any) *ToolResult } action, ok := args["action"].(string) + if !ok { return ErrorResult("action is required") } switch action { case "detect": + return t.detect() + case "scan": + return t.scan(args) + case "read": + return t.readDevice(args) + case "write": + return t.writeDevice(args) + default: + return ErrorResult(fmt.Sprintf("unknown action: %s (valid: detect, scan, read, write)", action)) } } // detect lists available I2C buses by globbing /dev/i2c-* + func (t *I2CTool) detect() *ToolResult { matches, err := filepath.Glob("/dev/i2c-*") if err != nil { @@ -102,11 +131,14 @@ func (t *I2CTool) detect() *ToolResult { type busInfo struct { Path string `json:"path"` - Bus string `json:"bus"` + + Bus string `json:"bus"` } buses := make([]busInfo, 0, len(matches)) + re := regexp.MustCompile(`/dev/i2c-(\d+)`) + for _, m := range matches { if sub := re.FindStringSubmatch(m); sub != nil { buses = append(buses, busInfo{Path: m, Bus: sub[1]}) @@ -114,44 +146,62 @@ func (t *I2CTool) detect() *ToolResult { } result, _ := json.MarshalIndent(buses, "", " ") + return SilentResult(fmt.Sprintf("Found %d I2C bus(es):\n%s", len(buses), string(result))) } // Helper functions for I2C operations (used by platform-specific implementations) // isValidBusID checks that a bus identifier is a simple number (prevents path injection) + // + //nolint:unused // Used by i2c_linux.go + func isValidBusID(id string) bool { matched, _ := regexp.MatchString(`^\d+$`, id) + return matched } // parseI2CAddress extracts and validates an I2C address from args + // + //nolint:unused // Used by i2c_linux.go + func parseI2CAddress(args map[string]any) (int, *ToolResult) { addrFloat, ok := args["address"].(float64) + if !ok { return 0, ErrorResult("address is required (e.g. 0x38 for AHT20)") } + addr := int(addrFloat) + if addr < 0x03 || addr > 0x77 { return 0, ErrorResult("address must be in valid 7-bit range (0x03-0x77)") } + return addr, nil } // parseI2CBus extracts and validates an I2C bus from args + // + //nolint:unused // Used by i2c_linux.go + func parseI2CBus(args map[string]any) (string, *ToolResult) { bus, ok := args["bus"].(string) + if !ok || bus == "" { return "", ErrorResult("bus is required (e.g. \"1\" for /dev/i2c-1)") } + if !isValidBusID(bus) { return "", ErrorResult("invalid bus identifier: must be a number (e.g. \"1\")") } + return bus, nil } diff --git a/pkg/tools/i2c_linux.go b/pkg/tools/i2c_linux.go index 4eaaf8f09..8b710f37a 100644 --- a/pkg/tools/i2c_linux.go +++ b/pkg/tools/i2c_linux.go @@ -8,279 +8,465 @@ import ( ) // I2C ioctl constants from Linux kernel headers (<linux/i2c-dev.h>, <linux/i2c.h>) + const ( i2cSlave = 0x0703 // Set slave address (fails if in use by driver) + i2cFuncs = 0x0705 // Query adapter functionality bitmask + i2cSmbus = 0x0720 // Perform SMBus transaction // I2C_FUNC capability bits - i2cFuncSmbusQuick = 0x00010000 + + i2cFuncSmbusQuick = 0x00010000 + i2cFuncSmbusReadByte = 0x00020000 // SMBus transaction types - i2cSmbusRead = 0 + + i2cSmbusRead = 0 + i2cSmbusWrite = 1 // SMBus protocol sizes + i2cSmbusQuick = 0 - i2cSmbusByte = 1 + + i2cSmbusByte = 1 ) // i2cSmbusData matches the kernel union i2c_smbus_data (34 bytes max). + // For quick and byte transactions only the first byte is used (if at all). + type i2cSmbusData [34]byte // i2cSmbusArgs matches the kernel struct i2c_smbus_ioctl_data. + type i2cSmbusArgs struct { readWrite uint8 - command uint8 - size uint32 - data *i2cSmbusData + + command uint8 + + size uint32 + + data *i2cSmbusData } // smbusProbe performs a single SMBus probe at the given address. + // Uses SMBus Quick Write (safest) or falls back to SMBus Read Byte for + // EEPROM address ranges where quick write can corrupt AT24RF08 chips. + // This matches i2cdetect's MODE_AUTO behavior. + func smbusProbe(fd int, addr int, hasQuick bool) bool { + // EEPROM ranges: use read byte (quick write can corrupt AT24RF08) + useReadByte := (addr >= 0x30 && addr <= 0x37) || (addr >= 0x50 && addr <= 0x5F) if !useReadByte && hasQuick { + // SMBus Quick Write: [START] [ADDR|W] [ACK/NACK] [STOP] + // Safest probe — no data transferred + args := i2cSmbusArgs{ + readWrite: i2cSmbusWrite, - command: 0, - size: i2cSmbusQuick, - data: nil, + + command: 0, + + size: i2cSmbusQuick, + + data: nil, } + _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSmbus, uintptr(unsafe.Pointer(&args))) + return errno == 0 + } // SMBus Read Byte: [START] [ADDR|R] [ACK/NACK] [DATA] [STOP] + var data i2cSmbusData + args := i2cSmbusArgs{ + readWrite: i2cSmbusRead, - command: 0, - size: i2cSmbusByte, - data: &data, + + command: 0, + + size: i2cSmbusByte, + + data: &data, } + _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSmbus, uintptr(unsafe.Pointer(&args))) + return errno == 0 + } // scan probes valid 7-bit addresses on a bus for connected devices. + // Uses the same hybrid probe strategy as i2cdetect's MODE_AUTO: + // SMBus Quick Write for most addresses, SMBus Read Byte for EEPROM ranges. + func (t *I2CTool) scan(args map[string]any) *ToolResult { + bus, errResult := parseI2CBus(args) + if errResult != nil { + return errResult + } devPath := fmt.Sprintf("/dev/i2c-%s", bus) + fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to open %s: %v (check permissions and i2c-dev module)", devPath, err)) + } + defer syscall.Close(fd) // Query adapter capabilities to determine available probe methods. + // I2C_FUNCS writes an unsigned long, which is word-sized on Linux. + var funcs uintptr + _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cFuncs, uintptr(unsafe.Pointer(&funcs))) + if errno != 0 { + return ErrorResult(fmt.Sprintf("failed to query I2C adapter capabilities on %s: %v", devPath, errno)) + } hasQuick := funcs&i2cFuncSmbusQuick != 0 + hasReadByte := funcs&i2cFuncSmbusReadByte != 0 if !hasQuick && !hasReadByte { + return ErrorResult( + fmt.Sprintf("I2C adapter %s supports neither SMBus Quick nor Read Byte — cannot probe safely", devPath), ) + } type deviceEntry struct { Address string `json:"address"` - Status string `json:"status,omitempty"` + + Status string `json:"status,omitempty"` } var found []deviceEntry + // Scan 0x08-0x77, skipping I2C reserved addresses 0x00-0x07 + for addr := 0x08; addr <= 0x77; addr++ { + // Set slave address — EBUSY means a kernel driver owns this address + _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr)) + if errno != 0 { + if errno == syscall.EBUSY { + found = append(found, deviceEntry{ + Address: fmt.Sprintf("0x%02x", addr), - Status: "busy (in use by kernel driver)", + + Status: "busy (in use by kernel driver)", }) + } + continue + } if smbusProbe(fd, addr, hasQuick) { + found = append(found, deviceEntry{ + Address: fmt.Sprintf("0x%02x", addr), }) + } + } if len(found) == 0 { + return SilentResult(fmt.Sprintf("No devices found on %s. Check wiring and pull-up resistors.", devPath)) + } result, _ := json.MarshalIndent(map[string]any{ - "bus": devPath, + + "bus": devPath, + "devices": found, - "count": len(found), + + "count": len(found), }, "", " ") + return SilentResult(fmt.Sprintf("Scan of %s:\n%s", devPath, string(result))) + } // readDevice reads bytes from an I2C device, optionally at a specific register + func (t *I2CTool) readDevice(args map[string]any) *ToolResult { + bus, errResult := parseI2CBus(args) + if errResult != nil { + return errResult + } addr, errResult := parseI2CAddress(args) + if errResult != nil { + return errResult + } length := 1 + if l, ok := args["length"].(float64); ok { + length = int(l) + } + if length < 1 || length > 256 { + return ErrorResult("length must be between 1 and 256") + } devPath := fmt.Sprintf("/dev/i2c-%s", bus) + fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to open %s: %v", devPath, err)) + } + defer syscall.Close(fd) // Set slave address + _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr)) + if errno != 0 { + return ErrorResult(fmt.Sprintf("failed to set I2C address 0x%02x: %v", addr, errno)) + } // If register is specified, write it first + if regFloat, ok := args["register"].(float64); ok { + reg := int(regFloat) + if reg < 0 || reg > 255 { + return ErrorResult("register must be between 0x00 and 0xFF") + } + _, err = syscall.Write(fd, []byte{byte(reg)}) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to write register 0x%02x: %v", reg, err)) + } + } // Read data + buf := make([]byte, length) + n, err := syscall.Read(fd, buf) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to read from device 0x%02x: %v", addr, err)) + } // Format as hex bytes + hexBytes := make([]string, n) + intBytes := make([]int, n) + for i := 0; i < n; i++ { + hexBytes[i] = fmt.Sprintf("0x%02x", buf[i]) + intBytes[i] = int(buf[i]) + } result, _ := json.MarshalIndent(map[string]any{ - "bus": devPath, + + "bus": devPath, + "address": fmt.Sprintf("0x%02x", addr), - "bytes": intBytes, - "hex": hexBytes, - "length": n, + + "bytes": intBytes, + + "hex": hexBytes, + + "length": n, }, "", " ") + return SilentResult(string(result)) + } // writeDevice writes bytes to an I2C device, optionally at a specific register + func (t *I2CTool) writeDevice(args map[string]any) *ToolResult { + confirm, _ := args["confirm"].(bool) + if !confirm { + return ErrorResult( + "write operations require confirm: true. Please confirm with the user before writing to I2C devices, as incorrect writes can misconfigure hardware.", ) + } bus, errResult := parseI2CBus(args) + if errResult != nil { + return errResult + } addr, errResult := parseI2CAddress(args) + if errResult != nil { + return errResult + } dataRaw, ok := args["data"].([]any) + if !ok || len(dataRaw) == 0 { + return ErrorResult("data is required for write (array of byte values 0-255)") + } + if len(dataRaw) > 256 { + return ErrorResult("data too long: maximum 256 bytes per I2C transaction") + } data := make([]byte, 0, len(dataRaw)+1) // If register is specified, prepend it to the data + if regFloat, ok := args["register"].(float64); ok { + reg := int(regFloat) + if reg < 0 || reg > 255 { + return ErrorResult("register must be between 0x00 and 0xFF") + } + data = append(data, byte(reg)) + } for i, v := range dataRaw { + f, ok := v.(float64) + if !ok { + return ErrorResult(fmt.Sprintf("data[%d] is not a valid byte value", i)) + } + b := int(f) + if b < 0 || b > 255 { + return ErrorResult(fmt.Sprintf("data[%d] = %d is out of byte range (0-255)", i, b)) + } + data = append(data, byte(b)) + } devPath := fmt.Sprintf("/dev/i2c-%s", bus) + fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to open %s: %v", devPath, err)) + } + defer syscall.Close(fd) // Set slave address + _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr)) + if errno != 0 { + return ErrorResult(fmt.Sprintf("failed to set I2C address 0x%02x: %v", addr, errno)) + } // Write data + n, err := syscall.Write(fd, data) + if err != nil { + return ErrorResult(fmt.Sprintf("failed to write to device 0x%02x: %v", addr, err)) + } return SilentResult(fmt.Sprintf("Wrote %d byte(s) to device 0x%02x on %s", n, addr, devPath)) + } diff --git a/pkg/tools/logs.go b/pkg/tools/logs.go index 00fea3395..33d63a28b 100644 --- a/pkg/tools/logs.go +++ b/pkg/tools/logs.go @@ -9,7 +9,9 @@ import ( ) // LogsTool provides on-demand access to application logs from the in-memory ring buffer. + // Designed for token-efficient log analysis: defaults to WARN level to exclude noise. + type LogsTool struct{} func NewLogsTool() *LogsTool { @@ -20,29 +22,40 @@ func (t *LogsTool) Name() string { return "logs" } func (t *LogsTool) Description() string { return "Retrieve recent application logs from the in-memory ring buffer. " + + "Use level filter to minimize token usage (default: WARN). " + + "Call this when the user asks about errors, issues, or system health." } func (t *LogsTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "level": map[string]any{ - "type": "string", + "type": "string", + "description": "Minimum log level: DEBUG, INFO, WARN, ERROR. Default: WARN", - "enum": []string{"DEBUG", "INFO", "WARN", "ERROR"}, + + "enum": []string{"DEBUG", "INFO", "WARN", "ERROR"}, }, + "component": map[string]any{ - "type": "string", + "type": "string", + "description": "Filter by component name (e.g. telegram, discord, slack, agent)", }, + "limit": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Maximum number of log entries to return. Default: 50", }, + "query": map[string]any{ - "type": "string", + "type": "string", + "description": "Filter by substring match in log message", }, }, @@ -51,38 +64,50 @@ func (t *LogsTool) Parameters() map[string]any { func (t *LogsTool) Execute(_ context.Context, args map[string]any) *ToolResult { // Parse level (default: WARN) + level := logger.WARN + if lvlStr, ok := args["level"].(string); ok && lvlStr != "" { level = logger.ParseLevel(lvlStr) } // Parse component + component, _ := args["component"].(string) // Parse limit (default: 50, max: 300) + limit := 50 + if l, ok := args["limit"].(float64); ok && l > 0 { limit = int(l) } + if limit > 300 { limit = 300 } // Parse query + query, _ := args["query"].(string) // Fetch from ring buffer (already sanitized by RecentLogs) + entries := logger.RecentLogs(level, component, limit) // Apply query filter if specified + if query != "" { filtered := make([]logger.LogEntry, 0, len(entries)) + queryLower := strings.ToLower(query) + for _, e := range entries { if strings.Contains(strings.ToLower(e.Message), queryLower) { filtered = append(filtered, e) } } + entries = filtered } diff --git a/pkg/tools/logs_test.go b/pkg/tools/logs_test.go index 0b268a2e0..d7ec4c024 100644 --- a/pkg/tools/logs_test.go +++ b/pkg/tools/logs_test.go @@ -11,22 +11,31 @@ import ( func setupTestLogs(t *testing.T) { t.Helper() + prev := logger.GetLevel() + t.Cleanup(func() { logger.SetLevel(prev) }) + logger.SetLevel(logger.DEBUG) logger.DebugC("agent", "debug message") + logger.InfoC("telegram", "message received") + logger.WarnC("telegram", "webhook retry") + logger.ErrorC("discord", "connection timeout") + logger.WarnCF("wecom", "signature failed", map[string]any{ "token": "secret-value", + "nonce": "safe-value", }) } func TestLogsTool_DefaultLevel(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{}) @@ -36,6 +45,7 @@ func TestLogsTool_DefaultLevel(t *testing.T) { } var entries []logger.LogEntry + if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { t.Fatalf("failed to parse result: %v", err) } @@ -49,6 +59,7 @@ func TestLogsTool_DefaultLevel(t *testing.T) { func TestLogsTool_LevelFilter(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{ @@ -60,6 +71,7 @@ func TestLogsTool_LevelFilter(t *testing.T) { } var entries []logger.LogEntry + if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { t.Fatalf("failed to parse result: %v", err) } @@ -73,10 +85,12 @@ func TestLogsTool_LevelFilter(t *testing.T) { func TestLogsTool_ComponentFilter(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{ - "level": "DEBUG", + "level": "DEBUG", + "component": "telegram", }) @@ -85,6 +99,7 @@ func TestLogsTool_ComponentFilter(t *testing.T) { } var entries []logger.LogEntry + if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { t.Fatalf("failed to parse result: %v", err) } @@ -98,10 +113,12 @@ func TestLogsTool_ComponentFilter(t *testing.T) { func TestLogsTool_QueryFilter(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{ "level": "DEBUG", + "query": "timeout", }) @@ -110,6 +127,7 @@ func TestLogsTool_QueryFilter(t *testing.T) { } var entries []logger.LogEntry + if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { t.Fatalf("failed to parse result: %v", err) } @@ -117,6 +135,7 @@ func TestLogsTool_QueryFilter(t *testing.T) { if len(entries) == 0 { t.Fatal("expected at least one entry matching 'timeout'") } + for _, e := range entries { if !strings.Contains(strings.ToLower(e.Message), "timeout") { t.Errorf("entry should contain 'timeout': %s", e.Message) @@ -126,14 +145,17 @@ func TestLogsTool_QueryFilter(t *testing.T) { func TestLogsTool_QueryCaseInsensitive(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{ "level": "DEBUG", + "query": "TIMEOUT", }) var entries []logger.LogEntry + if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { t.Fatalf("failed to parse result: %v", err) } @@ -145,10 +167,12 @@ func TestLogsTool_QueryCaseInsensitive(t *testing.T) { func TestLogsTool_Limit(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{ "level": "DEBUG", + "limit": float64(2), }) @@ -157,6 +181,7 @@ func TestLogsTool_Limit(t *testing.T) { } var entries []logger.LogEntry + if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { t.Fatalf("failed to parse result: %v", err) } @@ -170,12 +195,15 @@ func TestLogsTool_LimitMax(t *testing.T) { tool := NewLogsTool() // limit > 300 should be capped + result := tool.Execute(context.Background(), map[string]any{ "level": "DEBUG", + "limit": float64(999), }) // Should not error, just cap silently + if result.IsError { t.Fatalf("unexpected error: %s", result.ForLLM) } @@ -183,10 +211,12 @@ func TestLogsTool_LimitMax(t *testing.T) { func TestLogsTool_FieldsSanitized(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{ - "level": "WARN", + "level": "WARN", + "component": "wecom", }) @@ -195,22 +225,27 @@ func TestLogsTool_FieldsSanitized(t *testing.T) { } var entries []logger.LogEntry + if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { t.Fatalf("failed to parse result: %v", err) } found := false + for _, e := range entries { if e.Fields != nil && e.Fields["token"] != nil { found = true + if e.Fields["token"] != "***" { t.Errorf("token field should be sanitized, got %v", e.Fields["token"]) } + if e.Fields["nonce"] != "safe-value" { t.Errorf("nonce field should be preserved, got %v", e.Fields["nonce"]) } } } + if !found { t.Error("expected to find wecom entry with token field") } @@ -218,19 +253,23 @@ func TestLogsTool_FieldsSanitized(t *testing.T) { func TestLogsTool_NoResults(t *testing.T) { prev := logger.GetLevel() + defer logger.SetLevel(prev) + logger.SetLevel(logger.DEBUG) tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{ - "level": "ERROR", + "level": "ERROR", + "component": "nonexistent-component-xyz", }) if result.IsError { t.Fatalf("should not be an error result: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "No log entries found") { t.Errorf("expected 'No log entries found' message, got: %s", result.ForLLM) } @@ -238,9 +277,11 @@ func TestLogsTool_NoResults(t *testing.T) { func TestLogsTool_Silent(t *testing.T) { setupTestLogs(t) + tool := NewLogsTool() result := tool.Execute(context.Background(), map[string]any{}) + if !result.Silent { t.Error("logs tool result should be Silent") } @@ -252,10 +293,13 @@ func TestLogsTool_ToolInterface(t *testing.T) { if tool.Name() != "logs" { t.Errorf("expected name 'logs', got %q", tool.Name()) } + if tool.Description() == "" { t.Error("description should not be empty") } + params := tool.Parameters() + if params == nil { t.Error("parameters should not be nil") } diff --git a/pkg/tools/message.go b/pkg/tools/message.go index 15ef4ff73..7efebd5ad 100644 --- a/pkg/tools/message.go +++ b/pkg/tools/message.go @@ -8,10 +8,13 @@ import ( type SendCallback func(channel, chatID, content string) error type MessageTool struct { - sendCallback SendCallback + sendCallback SendCallback + defaultChannel string - defaultChatID string - sentInRound bool // Tracks whether a message was sent in the current processing round + + defaultChatID string + + sentInRound bool // Tracks whether a message was sent in the current processing round } func NewMessageTool() *MessageTool { @@ -29,31 +32,41 @@ func (t *MessageTool) Description() string { func (t *MessageTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "content": map[string]any{ - "type": "string", + "type": "string", + "description": "The message content to send", }, + "channel": map[string]any{ - "type": "string", + "type": "string", + "description": "Optional: target channel (telegram, whatsapp, etc.)", }, + "chat_id": map[string]any{ - "type": "string", + "type": "string", + "description": "Optional: target chat/user ID", }, }, + "required": []string{"content"}, } } func (t *MessageTool) SetContext(channel, chatID string) { t.defaultChannel = channel + t.defaultChatID = chatID + t.sentInRound = false // Reset send tracking for new processing round } // HasSentInRound returns true if the message tool sent a message during the current round. + func (t *MessageTool) HasSentInRound() bool { return t.sentInRound } @@ -64,16 +77,19 @@ func (t *MessageTool) SetSendCallback(callback SendCallback) { func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolResult { content, ok := args["content"].(string) + if !ok { return &ToolResult{ForLLM: "content is required", IsError: true} } channel, _ := args["channel"].(string) + chatID, _ := args["chat_id"].(string) if channel == "" { channel = t.defaultChannel } + if chatID == "" { chatID = t.defaultChatID } @@ -88,16 +104,21 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolRes if err := t.sendCallback(channel, chatID, content); err != nil { return &ToolResult{ - ForLLM: fmt.Sprintf("sending message: %v", err), + ForLLM: fmt.Sprintf("sending message: %v", err), + IsError: true, - Err: err, + + Err: err, } } t.sentInRound = true + // Silent: user already received the message directly + return &ToolResult{ ForLLM: fmt.Sprintf("Message sent to %s:%s", channel, chatID), + Silent: true, } } diff --git a/pkg/tools/message_test.go b/pkg/tools/message_test.go index 717c1117b..f49beab76 100644 --- a/pkg/tools/message_test.go +++ b/pkg/tools/message_test.go @@ -8,17 +8,23 @@ import ( func TestMessageTool_Execute_Success(t *testing.T) { tool := NewMessageTool() + tool.SetContext("test-channel", "test-chat-id") var sentChannel, sentChatID, sentContent string + tool.SetSendCallback(func(channel, chatID, content string) error { sentChannel = channel + sentChatID = chatID + sentContent = content + return nil }) ctx := context.Background() + args := map[string]any{ "content": "Hello, world!", } @@ -26,33 +32,41 @@ func TestMessageTool_Execute_Success(t *testing.T) { result := tool.Execute(ctx, args) // Verify message was sent with correct parameters + if sentChannel != "test-channel" { t.Errorf("Expected channel 'test-channel', got '%s'", sentChannel) } + if sentChatID != "test-chat-id" { t.Errorf("Expected chatID 'test-chat-id', got '%s'", sentChatID) } + if sentContent != "Hello, world!" { t.Errorf("Expected content 'Hello, world!', got '%s'", sentContent) } // Verify ToolResult meets US-011 criteria: + // - Send success returns SilentResult (Silent=true) + if !result.Silent { t.Error("Expected Silent=true for successful send") } // - ForLLM contains send status description + if result.ForLLM != "Message sent to test-channel:test-chat-id" { t.Errorf("Expected ForLLM 'Message sent to test-channel:test-chat-id', got '%s'", result.ForLLM) } // - ForUser is empty (user already received message directly) + if result.ForUser != "" { t.Errorf("Expected ForUser to be empty, got '%s'", result.ForUser) } // - IsError should be false + if result.IsError { t.Error("Expected IsError=false for successful send") } @@ -60,28 +74,37 @@ func TestMessageTool_Execute_Success(t *testing.T) { func TestMessageTool_Execute_WithCustomChannel(t *testing.T) { tool := NewMessageTool() + tool.SetContext("default-channel", "default-chat-id") var sentChannel, sentChatID string + tool.SetSendCallback(func(channel, chatID, content string) error { sentChannel = channel + sentChatID = chatID + return nil }) ctx := context.Background() + args := map[string]any{ "content": "Test message", + "channel": "custom-channel", + "chat_id": "custom-chat-id", } result := tool.Execute(ctx, args) // Verify custom channel/chatID were used instead of defaults + if sentChannel != "custom-channel" { t.Errorf("Expected channel 'custom-channel', got '%s'", sentChannel) } + if sentChatID != "custom-chat-id" { t.Errorf("Expected chatID 'custom-chat-id', got '%s'", sentChatID) } @@ -89,6 +112,7 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) { if !result.Silent { t.Error("Expected Silent=true") } + if result.ForLLM != "Message sent to custom-channel:custom-chat-id" { t.Errorf("Expected ForLLM 'Message sent to custom-channel:custom-chat-id', got '%s'", result.ForLLM) } @@ -96,14 +120,17 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) { func TestMessageTool_Execute_SendFailure(t *testing.T) { tool := NewMessageTool() + tool.SetContext("test-channel", "test-chat-id") sendErr := errors.New("network error") + tool.SetSendCallback(func(channel, chatID, content string) error { return sendErr }) ctx := context.Background() + args := map[string]any{ "content": "Test message", } @@ -111,21 +138,27 @@ func TestMessageTool_Execute_SendFailure(t *testing.T) { result := tool.Execute(ctx, args) // Verify ToolResult for send failure: + // - Send failure returns ErrorResult (IsError=true) + if !result.IsError { t.Error("Expected IsError=true for failed send") } // - ForLLM contains error description + expectedErrMsg := "sending message: network error" + if result.ForLLM != expectedErrMsg { t.Errorf("Expected ForLLM '%s', got '%s'", expectedErrMsg, result.ForLLM) } // - Err field should contain original error + if result.Err == nil { t.Error("Expected Err to be set") } + if result.Err != sendErr { t.Errorf("Expected Err to be sendErr, got %v", result.Err) } @@ -133,17 +166,21 @@ func TestMessageTool_Execute_SendFailure(t *testing.T) { func TestMessageTool_Execute_MissingContent(t *testing.T) { tool := NewMessageTool() + tool.SetContext("test-channel", "test-chat-id") ctx := context.Background() + args := map[string]any{} // content missing result := tool.Execute(ctx, args) // Verify error result for missing content + if !result.IsError { t.Error("Expected IsError=true for missing content") } + if result.ForLLM != "content is required" { t.Errorf("Expected ForLLM 'content is required', got '%s'", result.ForLLM) } @@ -151,6 +188,7 @@ func TestMessageTool_Execute_MissingContent(t *testing.T) { func TestMessageTool_Execute_NoTargetChannel(t *testing.T) { tool := NewMessageTool() + // No SetContext called, so defaultChannel and defaultChatID are empty tool.SetSendCallback(func(channel, chatID, content string) error { @@ -158,6 +196,7 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) { }) ctx := context.Background() + args := map[string]any{ "content": "Test message", } @@ -165,9 +204,11 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) { result := tool.Execute(ctx, args) // Verify error when no target channel specified + if !result.IsError { t.Error("Expected IsError=true when no target channel") } + if result.ForLLM != "No target channel/chat specified" { t.Errorf("Expected ForLLM 'No target channel/chat specified', got '%s'", result.ForLLM) } @@ -175,10 +216,13 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) { func TestMessageTool_Execute_NotConfigured(t *testing.T) { tool := NewMessageTool() + tool.SetContext("test-channel", "test-chat-id") + // No SetSendCallback called ctx := context.Background() + args := map[string]any{ "content": "Test message", } @@ -186,9 +230,11 @@ func TestMessageTool_Execute_NotConfigured(t *testing.T) { result := tool.Execute(ctx, args) // Verify error when send callback not configured + if !result.IsError { t.Error("Expected IsError=true when send callback not configured") } + if result.ForLLM != "Message sending not configured" { t.Errorf("Expected ForLLM 'Message sending not configured', got '%s'", result.ForLLM) } @@ -196,6 +242,7 @@ func TestMessageTool_Execute_NotConfigured(t *testing.T) { func TestMessageTool_Name(t *testing.T) { tool := NewMessageTool() + if tool.Name() != "message" { t.Errorf("Expected name 'message', got '%s'", tool.Name()) } @@ -203,7 +250,9 @@ func TestMessageTool_Name(t *testing.T) { func TestMessageTool_Description(t *testing.T) { tool := NewMessageTool() + desc := tool.Description() + if desc == "" { t.Error("Description should not be empty") } @@ -211,48 +260,63 @@ func TestMessageTool_Description(t *testing.T) { func TestMessageTool_Parameters(t *testing.T) { tool := NewMessageTool() + params := tool.Parameters() // Verify parameters structure + typ, ok := params["type"].(string) + if !ok || typ != "object" { t.Error("Expected type 'object'") } props, ok := params["properties"].(map[string]any) + if !ok { t.Fatal("Expected properties to be a map") } // Check required properties + required, ok := params["required"].([]string) + if !ok || len(required) != 1 || required[0] != "content" { t.Error("Expected 'content' to be required") } // Check content property + contentProp, ok := props["content"].(map[string]any) + if !ok { t.Error("Expected 'content' property") } + if contentProp["type"] != "string" { t.Error("Expected content type to be 'string'") } // Check channel property (optional) + channelProp, ok := props["channel"].(map[string]any) + if !ok { t.Error("Expected 'channel' property") } + if channelProp["type"] != "string" { t.Error("Expected channel type to be 'string'") } // Check chat_id property (optional) + chatIDProp, ok := props["chat_id"].(map[string]any) + if !ok { t.Error("Expected 'chat_id' property") } + if chatIDProp["type"] != "string" { t.Error("Expected chat_id type to be 'string'") } diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index 13fbdd5bc..57e5f3818 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -14,9 +14,12 @@ import ( ) // NormalizeToolName keeps only lowercase ASCII letters. + // "read_file" → "readfile", "ReadFile" → "readfile", "read-file" → "readfile". + func NormalizeToolName(s string) string { var b strings.Builder + for _, r := range s { if r >= 'A' && r <= 'Z' { b.WriteRune(r + 32) @@ -24,12 +27,14 @@ func NormalizeToolName(s string) string { b.WriteRune(r) } } + return b.String() } type ToolRegistry struct { tools map[string]Tool - mu sync.RWMutex + + mu sync.RWMutex } func NewToolRegistry() *ToolRegistry { @@ -40,24 +45,33 @@ func NewToolRegistry() *ToolRegistry { func (r *ToolRegistry) Register(tool Tool) { r.mu.Lock() + defer r.mu.Unlock() + r.tools[tool.Name()] = tool } func (r *ToolRegistry) Get(name string) (Tool, bool) { r.mu.RLock() + defer r.mu.RUnlock() + // Exact match first + if tool, ok := r.tools[name]; ok { return tool, true } + // Fuzzy fallback: normalize and compare (handles "readfile" → "read_file" etc.) + norm := NormalizeToolName(name) + for _, tool := range r.tools { if NormalizeToolName(tool.Name()) == norm { return tool, true } } + return nil, false } @@ -66,70 +80,99 @@ func (r *ToolRegistry) Execute(ctx context.Context, name string, args map[string } // ExecuteWithContext executes a tool with channel/chatID context and optional async callback. + // If the tool implements AsyncTool and a non-nil callback is provided, + // the callback will be set on the tool before execution. + func (r *ToolRegistry) ExecuteWithContext( ctx context.Context, + name string, + args map[string]any, + channel, chatID string, + asyncCallback AsyncCallback, ) *ToolResult { logger.InfoCF("tool", "Tool execution started", + map[string]any{ "tool": name, + "args": args, }) tool, ok := r.Get(name) + if !ok { available := strings.Join(r.List(), ", ") + logger.ErrorCF("tool", "Tool not found", + map[string]any{ "tool": name, }) + return ErrorResult(fmt.Sprintf( + "tool %q not found. Available tools: %s", name, available, )).WithError(fmt.Errorf("tool not found")) } // If tool implements ContextualTool, set context + if contextualTool, ok := tool.(ContextualTool); ok && channel != "" && chatID != "" { contextualTool.SetContext(channel, chatID) } // If tool implements AsyncTool and callback is provided, set callback + if asyncTool, ok := tool.(AsyncTool); ok && asyncCallback != nil { asyncTool.SetCallback(asyncCallback) + logger.DebugCF("tool", "Async callback injected", + map[string]any{ "tool": name, }) } start := time.Now() + result := tool.Execute(ctx, args) + duration := time.Since(start) // Log based on result type + if result.IsError { logger.ErrorCF("tool", "Tool execution failed", + map[string]any{ - "tool": name, + "tool": name, + "duration": duration.Milliseconds(), - "error": result.ForLLM, + + "error": result.ForLLM, }) } else if result.Async { logger.InfoCF("tool", "Tool started (async)", + map[string]any{ - "tool": name, + "tool": name, + "duration": duration.Milliseconds(), }) } else { logger.InfoCF("tool", "Tool execution completed", + map[string]any{ - "tool": name, - "duration_ms": duration.Milliseconds(), + "tool": name, + + "duration_ms": duration.Milliseconds(), + "result_length": len(result.ForLLM), }) } @@ -138,53 +181,75 @@ func (r *ToolRegistry) ExecuteWithContext( } // sortedToolNames returns tool names in sorted order for deterministic iteration. + // This is critical for KV cache stability: non-deterministic map iteration would + // produce different system prompts and tool definitions on each call, invalidating + // the LLM's prefix cache even when no tools have changed. + func (r *ToolRegistry) sortedToolNames() []string { names := make([]string, 0, len(r.tools)) + for name := range r.tools { names = append(names, name) } + sort.Strings(names) + return names } func (r *ToolRegistry) GetDefinitions() []map[string]any { r.mu.RLock() + defer r.mu.RUnlock() sorted := r.sortedToolNames() + definitions := make([]map[string]any, 0, len(sorted)) + for _, name := range sorted { definitions = append(definitions, ToolToSchema(r.tools[name])) } + return definitions } // ToProviderDefs converts tool definitions to provider-compatible format. + // This is the format expected by LLM provider APIs. + func (r *ToolRegistry) ToProviderDefs() []providers.ToolDefinition { r.mu.RLock() + defer r.mu.RUnlock() sorted := r.sortedToolNames() + definitions := make([]providers.ToolDefinition, 0, len(sorted)) + for _, name := range sorted { tool := r.tools[name] + schema := ToolToSchema(tool) // Safely extract nested values with type checks + fn, ok := schema["function"].(map[string]any) + if !ok { continue } name, _ := fn["name"].(string) + desc, _ := fn["description"].(string) + params, _ := fn["parameters"].(map[string]any) paramsRaw := json.RawMessage(`{}`) + if len(params) > 0 { if payload, err := json.Marshal(params); err == nil { paramsRaw = json.RawMessage(payload) @@ -193,38 +258,51 @@ func (r *ToolRegistry) ToProviderDefs() []providers.ToolDefinition { definitions = append(definitions, providers.ToolDefinition{ Type: "function", + Function: providers.ToolFunctionDefinition{ - Name: name, + Name: name, + Description: desc, - Parameters: paramsRaw, + + Parameters: paramsRaw, }, }) } + return definitions } // List returns a list of all registered tool names. + func (r *ToolRegistry) List() []string { r.mu.RLock() + defer r.mu.RUnlock() return r.sortedToolNames() } // Count returns the number of registered tools. + func (r *ToolRegistry) Count() int { r.mu.RLock() + defer r.mu.RUnlock() + return len(r.tools) } // GetRuntimeStatus aggregates runtime status from all tools that implement StatusProvider. + // Returns empty string if no tool has status to report. + func (r *ToolRegistry) GetRuntimeStatus() string { r.mu.RLock() + defer r.mu.RUnlock() var parts []string + for _, tool := range r.tools { if sp, ok := tool.(StatusProvider); ok { if s := sp.RuntimeStatus(); s != "" { @@ -232,40 +310,53 @@ func (r *ToolRegistry) GetRuntimeStatus() string { } } } + if len(parts) == 0 { return "" } + return strings.Join(parts, "\n\n") } // buildParamHint extracts parameter names from a JSON schema and returns + // a hint string like "(task, label?, preset?)". Required params are bare, + // optional params have a trailing "?". + func buildParamHint(schema map[string]any) string { props, _ := schema["properties"].(map[string]any) + if len(props) == 0 { return "" } reqSlice, _ := schema["required"].([]string) + reqSet := make(map[string]bool, len(reqSlice)) + for _, r := range reqSlice { reqSet[r] = true } names := make([]string, 0, len(props)) + for name := range props { names = append(names, name) } + sort.Strings(names) parts := make([]string, 0, len(names)) + // Required params first, then optional + for _, name := range names { if reqSet[name] { parts = append(parts, name) } } + for _, name := range names { if !reqSet[name] { parts = append(parts, name+"?") @@ -276,17 +367,25 @@ func buildParamHint(schema map[string]any) string { } // GetSummaries returns human-readable summaries of all registered tools. + // Returns a slice of "- `name`(params) - description" strings. + func (r *ToolRegistry) GetSummaries() []string { r.mu.RLock() + defer r.mu.RUnlock() sorted := r.sortedToolNames() + summaries := make([]string, 0, len(sorted)) + for _, name := range sorted { tool := r.tools[name] + hint := buildParamHint(tool.Parameters()) + summaries = append(summaries, fmt.Sprintf("- `%s`%s - %s", tool.Name(), hint, tool.Description())) } + return summaries } diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go index 05c2676b3..cff843c3e 100644 --- a/pkg/tools/registry_test.go +++ b/pkg/tools/registry_test.go @@ -12,32 +12,42 @@ import ( // --- mock types --- type mockRegistryTool struct { - name string - desc string + name string + + desc string + params map[string]any + result *ToolResult } -func (m *mockRegistryTool) Name() string { return m.name } -func (m *mockRegistryTool) Description() string { return m.desc } +func (m *mockRegistryTool) Name() string { return m.name } + +func (m *mockRegistryTool) Description() string { return m.desc } + func (m *mockRegistryTool) Parameters() map[string]any { return m.params } + func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]any) *ToolResult { return m.result } type mockCtxTool struct { mockRegistryTool + channel string - chatID string + + chatID string } func (m *mockCtxTool) SetContext(channel, chatID string) { m.channel = channel + m.chatID = chatID } type mockAsyncRegistryTool struct { mockRegistryTool + cb AsyncCallback } @@ -49,9 +59,12 @@ func (m *mockAsyncRegistryTool) SetCallback(cb AsyncCallback) { func newMockTool(name, desc string) *mockRegistryTool { return &mockRegistryTool{ - name: name, - desc: desc, + name: name, + + desc: desc, + params: map[string]any{"type": "object"}, + result: SilentResult("ok"), } } @@ -63,15 +76,23 @@ func TestNormalizeToolName(t *testing.T) { input, want string }{ {"read_file", "readfile"}, + {"readfile", "readfile"}, + {"ReadFile", "readfile"}, + {"read-file", "readfile"}, + {"edit_file", "editfile"}, + {"web_search", "websearch"}, + {"EXEC", "exec"}, } + for _, tt := range tests { got := NormalizeToolName(tt.input) + if got != tt.want { t.Errorf("NormalizeToolName(%q) = %q, want %q", tt.input, got, tt.want) } @@ -80,9 +101,11 @@ func TestNormalizeToolName(t *testing.T) { func TestNewToolRegistry(t *testing.T) { r := NewToolRegistry() + if r.Count() != 0 { t.Errorf("expected empty registry, got count %d", r.Count()) } + if len(r.List()) != 0 { t.Errorf("expected empty list, got %v", r.List()) } @@ -90,13 +113,17 @@ func TestNewToolRegistry(t *testing.T) { func TestToolRegistry_RegisterAndGet(t *testing.T) { r := NewToolRegistry() + tool := newMockTool("echo", "echoes input") + r.Register(tool) got, ok := r.Get("echo") + if !ok { t.Fatal("expected to find registered tool") } + if got.Name() != "echo" { t.Errorf("expected name 'echo', got %q", got.Name()) } @@ -104,7 +131,9 @@ func TestToolRegistry_RegisterAndGet(t *testing.T) { func TestToolRegistry_Get_NotFound(t *testing.T) { r := NewToolRegistry() + _, ok := r.Get("nonexistent") + if ok { t.Error("expected ok=false for unregistered tool") } @@ -112,28 +141,42 @@ func TestToolRegistry_Get_NotFound(t *testing.T) { func TestToolRegistry_Get_FuzzyMatch(t *testing.T) { r := NewToolRegistry() + r.Register(newMockTool("read_file", "reads a file")) + r.Register(newMockTool("edit_file", "edits a file")) + r.Register(newMockTool("web_search", "searches the web")) tests := []struct { - query string + query string + wantName string }{ {"readfile", "read_file"}, + {"ReadFile", "read_file"}, + {"read-file", "read_file"}, + {"editfile", "edit_file"}, + {"EditFile", "edit_file"}, + {"websearch", "web_search"}, + {"WebSearch", "web_search"}, } + for _, tt := range tests { tool, ok := r.Get(tt.query) + if !ok { t.Errorf("Get(%q) not found, want %q", tt.query, tt.wantName) + continue } + if tool.Name() != tt.wantName { t.Errorf("Get(%q).Name() = %q, want %q", tt.query, tool.Name(), tt.wantName) } @@ -142,13 +185,17 @@ func TestToolRegistry_Get_FuzzyMatch(t *testing.T) { func TestToolRegistry_RegisterOverwrite(t *testing.T) { r := NewToolRegistry() + r.Register(newMockTool("dup", "first")) + r.Register(newMockTool("dup", "second")) if r.Count() != 1 { t.Errorf("expected count 1 after overwrite, got %d", r.Count()) } + tool, _ := r.Get("dup") + if tool.Description() != "second" { t.Errorf("expected overwritten description 'second', got %q", tool.Description()) } @@ -156,17 +203,23 @@ func TestToolRegistry_RegisterOverwrite(t *testing.T) { func TestToolRegistry_Execute_Success(t *testing.T) { r := NewToolRegistry() + r.Register(&mockRegistryTool{ - name: "greet", - desc: "says hello", + name: "greet", + + desc: "says hello", + params: map[string]any{}, + result: SilentResult("hello"), }) result := r.Execute(context.Background(), "greet", nil) + if result.IsError { t.Errorf("expected success, got error: %s", result.ForLLM) } + if result.ForLLM != "hello" { t.Errorf("expected ForLLM 'hello', got %q", result.ForLLM) } @@ -174,13 +227,17 @@ func TestToolRegistry_Execute_Success(t *testing.T) { func TestToolRegistry_Execute_NotFound(t *testing.T) { r := NewToolRegistry() + result := r.Execute(context.Background(), "missing", nil) + if !result.IsError { t.Error("expected error for missing tool") } + if !strings.Contains(result.ForLLM, "not found") { t.Errorf("expected 'not found' in error, got %q", result.ForLLM) } + if result.Err == nil { t.Error("expected Err to be set via WithError") } @@ -188,9 +245,11 @@ func TestToolRegistry_Execute_NotFound(t *testing.T) { func TestToolRegistry_ExecuteWithContext_ContextualTool(t *testing.T) { r := NewToolRegistry() + ct := &mockCtxTool{ mockRegistryTool: *newMockTool("ctx_tool", "needs context"), } + r.Register(ct) r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "telegram", "chat-42", nil) @@ -198,6 +257,7 @@ func TestToolRegistry_ExecuteWithContext_ContextualTool(t *testing.T) { if ct.channel != "telegram" { t.Errorf("expected channel 'telegram', got %q", ct.channel) } + if ct.chatID != "chat-42" { t.Errorf("expected chatID 'chat-42', got %q", ct.chatID) } @@ -205,9 +265,11 @@ func TestToolRegistry_ExecuteWithContext_ContextualTool(t *testing.T) { func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) { r := NewToolRegistry() + ct := &mockCtxTool{ mockRegistryTool: *newMockTool("ctx_tool", "needs context"), } + r.Register(ct) r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "", "", nil) @@ -219,24 +281,31 @@ func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) { func TestToolRegistry_ExecuteWithContext_AsyncCallback(t *testing.T) { r := NewToolRegistry() + at := &mockAsyncRegistryTool{ mockRegistryTool: *newMockTool("async_tool", "async work"), } + at.result = AsyncResult("started") + r.Register(at) called := false + cb := func(_ context.Context, _ *ToolResult) { called = true } result := r.ExecuteWithContext(context.Background(), "async_tool", nil, "", "", cb) + if at.cb == nil { t.Error("expected SetCallback to have been called") } + if !result.Async { t.Error("expected async result") } at.cb(context.Background(), SilentResult("done")) + if !called { t.Error("expected callback to be invoked") } @@ -244,22 +313,29 @@ func TestToolRegistry_ExecuteWithContext_AsyncCallback(t *testing.T) { func TestToolRegistry_GetDefinitions(t *testing.T) { r := NewToolRegistry() + r.Register(newMockTool("alpha", "tool A")) defs := r.GetDefinitions() + if len(defs) != 1 { t.Fatalf("expected 1 definition, got %d", len(defs)) } + if defs[0]["type"] != "function" { t.Errorf("expected type 'function', got %v", defs[0]["type"]) } + fn, ok := defs[0]["function"].(map[string]any) + if !ok { t.Fatal("expected 'function' key to be a map") } + if fn["name"] != "alpha" { t.Errorf("expected name 'alpha', got %v", fn["name"]) } + if fn["description"] != "tool A" { t.Errorf("expected description 'tool A', got %v", fn["description"]) } @@ -267,34 +343,47 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { func TestToolRegistry_ToProviderDefs(t *testing.T) { r := NewToolRegistry() + params := map[string]any{"type": "object", "properties": map[string]any{}} + r.Register(&mockRegistryTool{ - name: "beta", - desc: "tool B", + name: "beta", + + desc: "tool B", + params: params, + result: SilentResult("ok"), }) defs := r.ToProviderDefs() + if len(defs) != 1 { t.Fatalf("expected 1 provider def, got %d", len(defs)) } want := providers.ToolDefinition{ Type: "function", + Function: providers.ToolFunctionDefinition{ - Name: "beta", + Name: "beta", + Description: "tool B", - Parameters: providers.MustMarshalParameters(params), + + Parameters: providers.MustMarshalParameters(params), }, } + got := defs[0] + if got.Type != want.Type { t.Errorf("Type: want %q, got %q", want.Type, got.Type) } + if got.Function.Name != want.Function.Name { t.Errorf("Name: want %q, got %q", want.Function.Name, got.Function.Name) } + if got.Function.Description != want.Function.Description { t.Errorf("Description: want %q, got %q", want.Function.Description, got.Function.Description) } @@ -302,18 +391,23 @@ func TestToolRegistry_ToProviderDefs(t *testing.T) { func TestToolRegistry_List(t *testing.T) { r := NewToolRegistry() + r.Register(newMockTool("x", "")) + r.Register(newMockTool("y", "")) names := r.List() + if len(names) != 2 { t.Fatalf("expected 2 names, got %d", len(names)) } nameSet := map[string]bool{} + for _, n := range names { nameSet[n] = true } + if !nameSet["x"] || !nameSet["y"] { t.Errorf("expected names {x, y}, got %v", names) } @@ -321,17 +415,21 @@ func TestToolRegistry_List(t *testing.T) { func TestToolRegistry_Count(t *testing.T) { r := NewToolRegistry() + if r.Count() != 0 { t.Errorf("expected 0, got %d", r.Count()) } r.Register(newMockTool("a", "")) + r.Register(newMockTool("b", "")) + if r.Count() != 2 { t.Errorf("expected 2, got %d", r.Count()) } r.Register(newMockTool("a", "replaced")) + if r.Count() != 2 { t.Errorf("expected 2 after overwrite, got %d", r.Count()) } @@ -339,62 +437,91 @@ func TestToolRegistry_Count(t *testing.T) { func TestBuildParamHint(t *testing.T) { tests := []struct { - name string + name string + schema map[string]any - want string + + want string }{ { name: "required and optional", + schema: map[string]any{ "type": "object", + "properties": map[string]any{ - "task": map[string]any{"type": "string"}, + "task": map[string]any{"type": "string"}, + "label": map[string]any{"type": "string"}, }, + "required": []string{"task"}, }, + want: "(task, label?)", }, + { name: "all required", + schema: map[string]any{ "type": "object", + "properties": map[string]any{ "command": map[string]any{"type": "string"}, }, + "required": []string{"command"}, }, + want: "(command)", }, + { name: "no properties", + schema: map[string]any{ "type": "object", }, + want: "", }, + { - name: "empty schema", + name: "empty schema", + schema: map[string]any{}, - want: "", + + want: "", }, + { - name: "nil schema", + name: "nil schema", + schema: nil, - want: "", + + want: "", }, + { name: "multiple optional sorted", + schema: map[string]any{ "type": "object", + "properties": map[string]any{ - "task": map[string]any{"type": "string"}, - "preset": map[string]any{"type": "string"}, - "label": map[string]any{"type": "string"}, + "task": map[string]any{"type": "string"}, + + "preset": map[string]any{"type": "string"}, + + "label": map[string]any{"type": "string"}, + "agent_id": map[string]any{"type": "string"}, }, + "required": []string{"task"}, }, + want: "(task, agent_id?, label?, preset?)", }, } @@ -402,6 +529,7 @@ func TestBuildParamHint(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := buildParamHint(tt.schema) + if got != tt.want { t.Errorf("buildParamHint() = %q, want %q", got, tt.want) } @@ -411,15 +539,19 @@ func TestBuildParamHint(t *testing.T) { func TestToolRegistry_GetSummaries(t *testing.T) { r := NewToolRegistry() + r.Register(newMockTool("read_file", "Reads a file")) summaries := r.GetSummaries() + if len(summaries) != 1 { t.Fatalf("expected 1 summary, got %d", len(summaries)) } + if !strings.Contains(summaries[0], "`read_file`") { t.Errorf("expected backtick-quoted name in summary, got %q", summaries[0]) } + if !strings.Contains(summaries[0], "Reads a file") { t.Errorf("expected description in summary, got %q", summaries[0]) } @@ -427,25 +559,35 @@ func TestToolRegistry_GetSummaries(t *testing.T) { func TestToolRegistry_GetSummaries_WithParamHint(t *testing.T) { r := NewToolRegistry() + r.Register(&mockRegistryTool{ name: "spawn", + desc: "Spawn a subagent", + params: map[string]any{ "type": "object", + "properties": map[string]any{ - "task": map[string]any{"type": "string"}, + "task": map[string]any{"type": "string"}, + "preset": map[string]any{"type": "string"}, }, + "required": []string{"task"}, }, + result: SilentResult("ok"), }) summaries := r.GetSummaries() + if len(summaries) != 1 { t.Fatalf("expected 1 summary, got %d", len(summaries)) } + // Should contain param hint + if !strings.Contains(summaries[0], "(task, preset?)") { t.Errorf("expected param hint in summary, got %q", summaries[0]) } @@ -453,21 +595,27 @@ func TestToolRegistry_GetSummaries_WithParamHint(t *testing.T) { func TestToolToSchema(t *testing.T) { tool := newMockTool("demo", "demo tool") + schema := ToolToSchema(tool) if schema["type"] != "function" { t.Errorf("expected type 'function', got %v", schema["type"]) } + fn, ok := schema["function"].(map[string]any) + if !ok { t.Fatal("expected 'function' to be a map") } + if fn["name"] != "demo" { t.Errorf("expected name 'demo', got %v", fn["name"]) } + if fn["description"] != "demo tool" { t.Errorf("expected description 'demo tool', got %v", fn["description"]) } + if fn["parameters"] == nil { t.Error("expected parameters to be set") } @@ -475,17 +623,25 @@ func TestToolToSchema(t *testing.T) { func TestToolRegistry_ConcurrentAccess(t *testing.T) { r := NewToolRegistry() + var wg sync.WaitGroup for i := range 50 { wg.Add(1) + go func(n int) { defer wg.Done() + name := string(rune('A' + n%26)) + r.Register(newMockTool(name, "concurrent")) + r.Get(name) + r.Count() + r.List() + r.GetDefinitions() }(i) } diff --git a/pkg/tools/result_test.go b/pkg/tools/result_test.go index a234e33f3..ac7d1bfda 100644 --- a/pkg/tools/result_test.go +++ b/pkg/tools/result_test.go @@ -12,12 +12,15 @@ func TestNewToolResult(t *testing.T) { if result.ForLLM != "test content" { t.Errorf("Expected ForLLM 'test content', got '%s'", result.ForLLM) } + if result.Silent { t.Error("Expected Silent to be false") } + if result.IsError { t.Error("Expected IsError to be false") } + if result.Async { t.Error("Expected Async to be false") } @@ -29,12 +32,15 @@ func TestSilentResult(t *testing.T) { if result.ForLLM != "silent operation" { t.Errorf("Expected ForLLM 'silent operation', got '%s'", result.ForLLM) } + if !result.Silent { t.Error("Expected Silent to be true") } + if result.IsError { t.Error("Expected IsError to be false") } + if result.Async { t.Error("Expected Async to be false") } @@ -46,12 +52,15 @@ func TestAsyncResult(t *testing.T) { if result.ForLLM != "async task started" { t.Errorf("Expected ForLLM 'async task started', got '%s'", result.ForLLM) } + if result.Silent { t.Error("Expected Silent to be false") } + if result.IsError { t.Error("Expected IsError to be false") } + if !result.Async { t.Error("Expected Async to be true") } @@ -63,12 +72,15 @@ func TestErrorResult(t *testing.T) { if result.ForLLM != "operation failed" { t.Errorf("Expected ForLLM 'operation failed', got '%s'", result.ForLLM) } + if result.Silent { t.Error("Expected Silent to be false") } + if !result.IsError { t.Error("Expected IsError to be true") } + if result.Async { t.Error("Expected Async to be false") } @@ -76,20 +88,25 @@ func TestErrorResult(t *testing.T) { func TestUserResult(t *testing.T) { content := "user visible message" + result := UserResult(content) if result.ForLLM != content { t.Errorf("Expected ForLLM '%s', got '%s'", content, result.ForLLM) } + if result.ForUser != content { t.Errorf("Expected ForUser '%s', got '%s'", content, result.ForUser) } + if result.Silent { t.Error("Expected Silent to be false") } + if result.IsError { t.Error("Expected IsError to be false") } + if result.Async { t.Error("Expected Async to be false") } @@ -97,27 +114,37 @@ func TestUserResult(t *testing.T) { func TestToolResultJSONSerialization(t *testing.T) { tests := []struct { - name string + name string + result *ToolResult }{ { - name: "basic result", + name: "basic result", + result: NewToolResult("basic content"), }, + { - name: "silent result", + name: "silent result", + result: SilentResult("silent content"), }, + { - name: "async result", + name: "async result", + result: AsyncResult("async content"), }, + { - name: "error result", + name: "error result", + result: ErrorResult("error content"), }, + { - name: "user result", + name: "user result", + result: UserResult("user content"), }, } @@ -125,30 +152,38 @@ func TestToolResultJSONSerialization(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Marshal to JSON + data, err := json.Marshal(tt.result) if err != nil { t.Fatalf("Failed to marshal: %v", err) } // Unmarshal back + var decoded ToolResult + if err := json.Unmarshal(data, &decoded); err != nil { t.Fatalf("Failed to unmarshal: %v", err) } // Verify fields match (Err should be excluded) + if decoded.ForLLM != tt.result.ForLLM { t.Errorf("ForLLM mismatch: got '%s', want '%s'", decoded.ForLLM, tt.result.ForLLM) } + if decoded.ForUser != tt.result.ForUser { t.Errorf("ForUser mismatch: got '%s', want '%s'", decoded.ForUser, tt.result.ForUser) } + if decoded.Silent != tt.result.Silent { t.Errorf("Silent mismatch: got %v, want %v", decoded.Silent, tt.result.Silent) } + if decoded.IsError != tt.result.IsError { t.Errorf("IsError mismatch: got %v, want %v", decoded.IsError, tt.result.IsError) } + if decoded.Async != tt.result.Async { t.Errorf("Async mismatch: got %v, want %v", decoded.Async, tt.result.Async) } @@ -158,22 +193,27 @@ func TestToolResultJSONSerialization(t *testing.T) { func TestToolResultWithErrors(t *testing.T) { err := errors.New("underlying error") + result := ErrorResult("error message").WithError(err) if result.Err == nil { t.Error("Expected Err to be set") } + if result.Err.Error() != "underlying error" { t.Errorf("Expected Err message 'underlying error', got '%s'", result.Err.Error()) } // Verify Err is not serialized + data, marshalErr := json.Marshal(result) + if marshalErr != nil { t.Fatalf("Failed to marshal: %v", marshalErr) } var decoded ToolResult + if unmarshalErr := json.Unmarshal(data, &decoded); unmarshalErr != nil { t.Fatalf("Failed to unmarshal: %v", unmarshalErr) } @@ -192,37 +232,47 @@ func TestToolResultJSONStructure(t *testing.T) { } // Verify JSON structure + var parsed map[string]any + if err := json.Unmarshal(data, &parsed); err != nil { t.Fatalf("Failed to parse JSON: %v", err) } // Check expected keys exist + if _, ok := parsed["for_llm"]; !ok { t.Error("Expected 'for_llm' key in JSON") } + if _, ok := parsed["for_user"]; !ok { t.Error("Expected 'for_user' key in JSON") } + if _, ok := parsed["silent"]; !ok { t.Error("Expected 'silent' key in JSON") } + if _, ok := parsed["is_error"]; !ok { t.Error("Expected 'is_error' key in JSON") } + if _, ok := parsed["async"]; !ok { t.Error("Expected 'async' key in JSON") } // Check that 'err' is NOT present (it should have json:"-" tag) + if _, ok := parsed["err"]; ok { t.Error("Expected 'err' key to be excluded from JSON") } // Verify values + if parsed["for_llm"] != "test content" { t.Errorf("Expected for_llm 'test content', got %v", parsed["for_llm"]) } + if parsed["silent"] != false { t.Errorf("Expected silent false, got %v", parsed["silent"]) } diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index 98eb4c388..45672c2cb 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -22,16 +22,22 @@ import ( ) const ( - bgMaxLifetime = 45 * time.Minute - bgRingBufSize = 32 * 1024 // 32KB - bgInitCapture = 3 * time.Second + bgMaxLifetime = 45 * time.Minute + + bgRingBufSize = 32 * 1024 // 32KB + + bgInitCapture = 3 * time.Second + bgMaxProcesses = 10 ) // ringBuffer is a thread-safe circular buffer that retains the most recent bytes. + type ringBuffer struct { - mu sync.Mutex - buf []byte + mu sync.Mutex + + buf []byte + size int } @@ -40,147 +46,239 @@ func newRingBuffer(size int) *ringBuffer { } // Write appends data to the ring buffer, dropping oldest bytes if capacity is exceeded. + func (rb *ringBuffer) Write(p []byte) (int, error) { rb.mu.Lock() + defer rb.mu.Unlock() + rb.buf = append(rb.buf, p...) + if len(rb.buf) > rb.size { rb.buf = rb.buf[len(rb.buf)-rb.size:] } + return len(p), nil } // String returns the current buffer contents. + func (rb *ringBuffer) String() string { rb.mu.Lock() + defer rb.mu.Unlock() + return string(rb.buf) } // Lines returns the last n lines from the buffer. + func (rb *ringBuffer) Lines(n int) []string { rb.mu.Lock() + defer rb.mu.Unlock() + if len(rb.buf) == 0 { return nil } + all := strings.Split(string(rb.buf), "\n") + // Remove trailing empty element from final newline + if len(all) > 0 && all[len(all)-1] == "" { all = all[:len(all)-1] } + if n <= 0 || n >= len(all) { return all } + return all[len(all)-n:] } // Match checks if any line in the buffer matches the given regex pattern. + // Returns the first matching line, or empty string if no match. + func (rb *ringBuffer) Match(pattern *regexp.Regexp) string { rb.mu.Lock() + defer rb.mu.Unlock() + for _, line := range strings.Split(string(rb.buf), "\n") { if pattern.MatchString(line) { return line } } + return "" } // Len returns the current number of bytes in the buffer. + func (rb *ringBuffer) Len() int { rb.mu.Lock() + defer rb.mu.Unlock() + return len(rb.buf) } // bgProcess represents a background process managed by ExecTool. + type bgProcess struct { - id string - command string - cmd *exec.Cmd - pid int + id string + + command string + + cmd *exec.Cmd + + pid int + startedAt time.Time - output *ringBuffer - done chan struct{} // closed when process exits - exitErr error - cancel context.CancelFunc // cancels the monitor goroutine + + output *ringBuffer + + done chan struct{} // closed when process exits + + exitErr error + + cancel context.CancelFunc // cancels the monitor goroutine } // isRunning returns true if the process has not yet exited. + func (bp *bgProcess) isRunning() bool { select { case <-bp.done: + return false + default: + return true } } type ExecTool struct { - workingDir string - timeout time.Duration - denyPatterns []*regexp.Regexp - allowRules [][]string // pre-split command prefix allowlist + workingDir string + + timeout time.Duration + + denyPatterns []*regexp.Regexp + + allowRules [][]string // pre-split command prefix allowlist + restrictToWorkspace bool - localNetOnly bool // restrict curl/wget to localhost + RFC 1918 + + localNetOnly bool // restrict curl/wget to localhost + RFC 1918 // Background process management - bgMu sync.Mutex + + bgMu sync.Mutex + bgProcesses map[string]*bgProcess - bgNextID int - bgShutdown context.CancelFunc // cancels all bg monitor goroutines - bgCtx context.Context + + bgNextID int + + bgShutdown context.CancelFunc // cancels all bg monitor goroutines + + bgCtx context.Context } var defaultDenyPatterns = []*regexp.Regexp{ regexp.MustCompile(`\brm\s+-[rf]{1,2}\b`), + regexp.MustCompile(`\bdel\s+/[fq]\b`), + regexp.MustCompile(`\brmdir\s+/s\b`), + // Match disk wiping commands (must be followed by space/args) + regexp.MustCompile( + `\b(format|mkfs|diskpart)\b\s`, ), + regexp.MustCompile(`\bdd\s+if=`), + regexp.MustCompile(`>\s*/dev/sd[a-z]\b`), // Block writes to disk devices (but allow /dev/null) + regexp.MustCompile(`\b(shutdown|reboot|poweroff)\b`), + regexp.MustCompile(`:\(\)\s*\{.*\};\s*:`), + regexp.MustCompile(`\$\([^)]+\)`), + regexp.MustCompile(`\$\{[^}]+\}`), + regexp.MustCompile("`[^`]+`"), + regexp.MustCompile(`\|\s*sh\b`), + regexp.MustCompile(`\|\s*bash\b`), + regexp.MustCompile(`;\s*rm\s+-[rf]`), + regexp.MustCompile(`&&\s*rm\s+-[rf]`), + regexp.MustCompile(`\|\|\s*rm\s+-[rf]`), + regexp.MustCompile(`>\s*/dev/null\s*>&?\s*\d?`), + regexp.MustCompile(`<<\s*EOF`), + regexp.MustCompile(`\$\(\s*cat\s+`), + regexp.MustCompile(`\$\(\s*curl\s+`), + regexp.MustCompile(`\$\(\s*wget\s+`), + regexp.MustCompile(`\$\(\s*which\s+`), + regexp.MustCompile(`\bsudo\b`), + regexp.MustCompile(`\bchmod\s+[0-7]{3,4}\b`), + regexp.MustCompile(`\bchown\b`), + regexp.MustCompile(`\bpkill\b`), + regexp.MustCompile(`\bkillall\b`), + regexp.MustCompile(`\bkill\s+-[9]\b`), + regexp.MustCompile(`\bcurl\b.*\|\s*(sh|bash)`), + regexp.MustCompile(`\bwget\b.*\|\s*(sh|bash)`), + regexp.MustCompile(`\bnpm\s+install\s+-g\b`), + regexp.MustCompile(`\bpip\s+install\s+--user\b`), + regexp.MustCompile(`\bapt\s+(install|remove|purge)\b`), + regexp.MustCompile(`\byum\s+(install|remove)\b`), + regexp.MustCompile(`\bdnf\s+(install|remove)\b`), + regexp.MustCompile(`\bdocker\s+run\b`), + regexp.MustCompile(`\bdocker\s+exec\b`), + regexp.MustCompile(`\bgit\s+push\b`), + regexp.MustCompile(`\bgit\s+force\b`), + regexp.MustCompile(`\bgit\s+checkout\b`), + regexp.MustCompile(`\bgit\s+switch\b`), + regexp.MustCompile(`\bssh\b.*@`), + regexp.MustCompile(`\beval\b`), + regexp.MustCompile(`\bsource\s+.*\.sh\b`), } @@ -193,21 +291,27 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf if config != nil { execConfig := config.Tools.Exec + enableDenyPatterns := execConfig.EnableDenyPatterns + if enableDenyPatterns { denyPatterns = append(denyPatterns, defaultDenyPatterns...) + if len(execConfig.CustomDenyPatterns) > 0 { fmt.Printf("Using custom deny patterns: %v\n", execConfig.CustomDenyPatterns) + for _, pattern := range execConfig.CustomDenyPatterns { re, err := regexp.Compile(pattern) if err != nil { return nil, fmt.Errorf("invalid custom deny pattern %q: %w", pattern, err) } + denyPatterns = append(denyPatterns, re) } } } else { // If deny patterns are disabled, we won't add any patterns, allowing all commands. + fmt.Println("Warning: deny patterns are disabled. All commands will be allowed.") } } else { @@ -217,14 +321,21 @@ func NewExecToolWithConfig(workingDir string, restrict bool, config *config.Conf bgCtx, bgCancel := context.WithCancel(context.Background()) return &ExecTool{ - workingDir: workingDir, - timeout: 5 * time.Minute, - denyPatterns: denyPatterns, - allowRules: nil, + workingDir: workingDir, + + timeout: 5 * time.Minute, + + denyPatterns: denyPatterns, + + allowRules: nil, + restrictToWorkspace: restrict, - bgProcesses: make(map[string]*bgProcess), - bgCtx: bgCtx, - bgShutdown: bgCancel, + + bgProcesses: make(map[string]*bgProcess), + + bgCtx: bgCtx, + + bgShutdown: bgCancel, }, nil } @@ -239,58 +350,77 @@ func (t *ExecTool) Description() string { func (t *ExecTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "command": map[string]any{ - "type": "string", + "type": "string", + "description": "The shell command to execute", }, + "working_dir": map[string]any{ - "type": "string", + "type": "string", + "description": "Optional working directory for the command", }, + "background": map[string]any{ - "type": "boolean", + "type": "boolean", + "description": "Run the command in the background. Returns immediately with a process ID.", }, + "bg_action": map[string]any{ - "type": "string", - "enum": []string{"output", "kill"}, + "type": "string", + + "enum": []string{"output", "kill"}, + "description": "Action on a background process: 'output' to get latest output, 'kill' to stop it.", }, + "bg_id": map[string]any{ - "type": "string", + "type": "string", + "description": "Background process ID (e.g. 'bg-1'). Required with bg_action.", }, }, + "required": []string{}, } } func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult { // Handle bg_action first (output/kill) + if bgAction, ok := args["bg_action"].(string); ok && bgAction != "" { bgID, _ := args["bg_id"].(string) + return t.handleBgAction(bgAction, bgID) } // Check for background execution + bg, _ := args["background"].(bool) command, ok := args["command"].(string) + if !ok || command == "" { return ErrorResult("command is required") } cwd := t.workingDir + if override := WorkspaceOverrideFromCtx(ctx); override != "" { cwd = override } + if wd, ok := args["working_dir"].(string); ok && wd != "" { if t.restrictToWorkspace && cwd != "" { resolvedWD, err := validatePath(wd, cwd, true) if err != nil { return ErrorResult("Command blocked by safety guard (" + err.Error() + ")") } + cwd = resolvedWD } else { cwd = wd @@ -299,6 +429,7 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult if cwd == "" { wd, err := os.Getwd() + if err == nil { cwd = wd } @@ -316,23 +447,30 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]any) *ToolResult } // executeSync runs a command synchronously (existing behavior). + func (t *ExecTool) executeSync(ctx context.Context, command, cwd string) *ToolResult { // timeout == 0 means no timeout + var cmdCtx context.Context + var cancel context.CancelFunc + if t.timeout > 0 { cmdCtx, cancel = context.WithTimeout(ctx, t.timeout) } else { cmdCtx, cancel = context.WithCancel(ctx) } + defer cancel() var cmd *exec.Cmd + if runtime.GOOS == "windows" { cmd = exec.CommandContext(cmdCtx, "powershell", "-NoProfile", "-NonInteractive", "-Command", command) } else { cmd = exec.CommandContext(cmdCtx, "sh", "-c", command) } + if cwd != "" { cmd.Dir = cwd } @@ -340,7 +478,9 @@ func (t *ExecTool) executeSync(ctx context.Context, command, cwd string) *ToolRe prepareCommandForTermination(cmd) var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr if err := cmd.Start(); err != nil { @@ -348,43 +488,59 @@ func (t *ExecTool) executeSync(ctx context.Context, command, cwd string) *ToolRe } done := make(chan error, 1) + go func() { done <- cmd.Wait() }() var err error + select { case err = <-done: + case <-cmdCtx.Done(): + _ = terminateProcessTree(cmd) + select { case err = <-done: + case <-time.After(2 * time.Second): + if cmd.Process != nil { _ = cmd.Process.Kill() } + err = <-done } } var ob strings.Builder + ob.WriteString(stdout.String()) + if stderr.Len() > 0 { ob.WriteString("\nSTDERR:\n") + ob.WriteString(stderr.String()) } if err != nil { if errors.Is(cmdCtx.Err(), context.DeadlineExceeded) { msg := fmt.Sprintf("Command timed out after %v", t.timeout) + return &ToolResult{ - ForLLM: msg, + ForLLM: msg, + ForUser: msg, + IsError: true, } } + fmt.Fprintf(&ob, "\nExit code: %v", err) } + output := ob.String() if output == "" { @@ -392,53 +548,68 @@ func (t *ExecTool) executeSync(ctx context.Context, command, cwd string) *ToolRe } maxLen := 10000 + if len(output) > maxLen { output = output[:maxLen] + fmt.Sprintf("\n... (truncated, %d more chars)", len(output)-maxLen) } if err != nil { return &ToolResult{ - ForLLM: output, + ForLLM: output, + ForUser: output, + IsError: true, } } return &ToolResult{ - ForLLM: output, + ForLLM: output, + ForUser: output, + IsError: false, } } // executeBg starts a background process and returns immediately. + func (t *ExecTool) executeBg(command, cwd string) *ToolResult { t.bgMu.Lock() // Check max processes limit + running := 0 + for _, bp := range t.bgProcesses { if bp.isRunning() { running++ } } + if running >= bgMaxProcesses { t.bgMu.Unlock() + return ErrorResult( + fmt.Sprintf("maximum background processes reached (%d). Kill an existing one first.", bgMaxProcesses), ) } t.bgNextID++ + id := fmt.Sprintf("bg-%d", t.bgNextID) + t.bgMu.Unlock() var cmd *exec.Cmd + if runtime.GOOS == "windows" { cmd = exec.Command("powershell", "-NoProfile", "-NonInteractive", "-Command", command) } else { cmd = exec.Command("sh", "-c", command) } + if cwd != "" { cmd.Dir = cwd } @@ -448,10 +619,12 @@ func (t *ExecTool) executeBg(command, cwd string) *ToolResult { output := newRingBuffer(bgRingBufSize) // Use pipes to capture output + stdoutPipe, err := cmd.StdoutPipe() if err != nil { return ErrorResult(fmt.Sprintf("failed to create stdout pipe: %v", err)) } + stderrPipe, err := cmd.StderrPipe() if err != nil { return ErrorResult(fmt.Sprintf("failed to create stderr pipe: %v", err)) @@ -464,101 +637,151 @@ func (t *ExecTool) executeBg(command, cwd string) *ToolResult { monitorCtx, monitorCancel := context.WithCancel(t.bgCtx) bp := &bgProcess{ - id: id, - command: command, - cmd: cmd, - pid: cmd.Process.Pid, + id: id, + + command: command, + + cmd: cmd, + + pid: cmd.Process.Pid, + startedAt: time.Now(), - output: output, - done: make(chan struct{}), - cancel: monitorCancel, + + output: output, + + done: make(chan struct{}), + + cancel: monitorCancel, } t.bgMu.Lock() + t.bgProcesses[id] = bp + t.bgMu.Unlock() // io.Copy goroutines: pipe stdout/stderr into ring buffer + go io.Copy(output, stdoutPipe) + go io.Copy(output, stderrPipe) // cmd.Wait goroutine + waitDone := make(chan error, 1) + go func() { waitDone <- cmd.Wait() }() // Monitor goroutine: handles lifetime timer, process exit, and shutdown + go func() { lifetime := time.NewTimer(getBgMaxLifetime()) + defer lifetime.Stop() select { case err := <-waitDone: + // Process exited naturally + bp.exitErr = err + close(bp.done) + case <-lifetime.C: + // Max lifetime exceeded — kill + _ = terminateProcessTree(cmd) + select { case err := <-waitDone: + bp.exitErr = err + case <-time.After(2 * time.Second): + if cmd.Process != nil { _ = cmd.Process.Kill() } + bp.exitErr = <-waitDone } + close(bp.done) + case <-monitorCtx.Done(): + // Shutdown or explicit kill via cancel + _ = terminateProcessTree(cmd) + select { case err := <-waitDone: + bp.exitErr = err + case <-time.After(2 * time.Second): + if cmd.Process != nil { _ = cmd.Process.Kill() } + bp.exitErr = <-waitDone } + select { case <-bp.done: + default: + close(bp.done) } } }() // Capture initial output (wait up to bgInitCapture) + time.Sleep(bgInitCapture) + initialOutput := output.String() var sb strings.Builder + fmt.Fprintf(&sb, "Background process started.\n") + fmt.Fprintf(&sb, " id: %s\n", id) + fmt.Fprintf(&sb, " pid: %d\n", bp.pid) + fmt.Fprintf(&sb, " cmd: %s\n", command) + fmt.Fprintf(&sb, " max lifetime: %s\n", getBgMaxLifetime()) + if initialOutput != "" { fmt.Fprintf(&sb, "\nInitial output:\n%s", initialOutput) } return &ToolResult{ - ForLLM: sb.String(), + ForLLM: sb.String(), + ForUser: fmt.Sprintf("Background process %s (pid=%d) started: %s", id, bp.pid, command), } } // handleBgAction handles bg_action=output and bg_action=kill. + func (t *ExecTool) handleBgAction(action, bgID string) *ToolResult { if bgID == "" { return ErrorResult("bg_id is required for bg_action") } t.bgMu.Lock() + bp, ok := t.bgProcesses[bgID] + t.bgMu.Unlock() if !ok { @@ -567,23 +790,31 @@ func (t *ExecTool) handleBgAction(action, bgID string) *ToolResult { switch action { case "output": + return t.bgOutput(bp) + case "kill": + return t.bgKill(bp) + default: + return ErrorResult(fmt.Sprintf("unknown bg_action %q (use 'output' or 'kill')", action)) } } func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { var sb strings.Builder + fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command) if bp.isRunning() { uptime := time.Since(bp.startedAt).Truncate(time.Second) + fmt.Fprintf(&sb, "Status: running (uptime: %s, max: %s)\n", uptime, getBgMaxLifetime()) } else { ran := time.Since(bp.startedAt).Truncate(time.Second) + if bp.exitErr != nil { fmt.Fprintf(&sb, "Status: exited with error (ran: %s): %v\n", ran, bp.exitErr) } else { @@ -592,6 +823,7 @@ func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { } output := bp.output.String() + if output == "" { fmt.Fprintf(&sb, "\n(no output)") } else { @@ -599,7 +831,8 @@ func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { } return &ToolResult{ - ForLLM: sb.String(), + ForLLM: sb.String(), + ForUser: sb.String(), } } @@ -607,38 +840,52 @@ func (t *ExecTool) bgOutput(bp *bgProcess) *ToolResult { func (t *ExecTool) bgKill(bp *bgProcess) *ToolResult { if bp.isRunning() { bp.cancel() // triggers monitor goroutine cleanup + // Wait for process to actually exit + select { case <-bp.done: + case <-time.After(5 * time.Second): } } t.bgMu.Lock() + delete(t.bgProcesses, bp.id) + t.bgMu.Unlock() msg := fmt.Sprintf("Background process %s (pid=%d) terminated: %s", bp.id, bp.pid, bp.command) + return &ToolResult{ - ForLLM: msg, + ForLLM: msg, + ForUser: msg, } } // BgProcesses returns a snapshot of background processes for use by bg_monitor. + func (t *ExecTool) BgProcesses() map[string]*bgProcess { t.bgMu.Lock() + defer t.bgMu.Unlock() + snapshot := make(map[string]*bgProcess, len(t.bgProcesses)) + for k, v := range t.bgProcesses { snapshot[k] = v } + return snapshot } // RuntimeStatus implements StatusProvider for system prompt injection. + func (t *ExecTool) RuntimeStatus() string { t.bgMu.Lock() + defer t.bgMu.Unlock() if len(t.bgProcesses) == 0 { @@ -646,54 +893,75 @@ func (t *ExecTool) RuntimeStatus() string { } // Sort by ID for stable output + ids := make([]string, 0, len(t.bgProcesses)) + for id := range t.bgProcesses { ids = append(ids, id) } + sort.Strings(ids) var sb strings.Builder + sb.WriteString("## Background Processes\n\n") + for _, id := range ids { bp := t.bgProcesses[id] + if bp.isRunning() { uptime := time.Since(bp.startedAt).Truncate(time.Second) + fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n", + id, bp.pid, uptime, getBgMaxLifetime(), bp.command) } else { ran := time.Since(bp.startedAt).Truncate(time.Second) + if bp.exitErr != nil { fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n", + id, bp.pid, ran, bp.command) } else { fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n", + id, bp.pid, ran, bp.command) } } } + sb.WriteString("\nUse exec with bg_action=\"output\" / \"kill\" and bg_id to manage.\n") + sb.WriteString("Use bg_monitor for list/watch/tail operations.") return sb.String() } // Shutdown terminates all background processes. Call on application exit. + func (t *ExecTool) Shutdown() { t.bgShutdown() // cancel all monitor goroutines t.bgMu.Lock() + procs := make([]*bgProcess, 0, len(t.bgProcesses)) + for _, bp := range t.bgProcesses { procs = append(procs, bp) } + t.bgMu.Unlock() // Wait for all processes to exit + for _, bp := range procs { select { case <-bp.done: + case <-time.After(5 * time.Second): + // Force kill if still running + if bp.cmd.Process != nil { _ = bp.cmd.Process.Kill() } @@ -703,6 +971,7 @@ func (t *ExecTool) Shutdown() { func (t *ExecTool) guardCommand(command, cwd string) string { cmd := strings.TrimSpace(command) + lower := strings.ToLower(cmd) for _, pattern := range t.denyPatterns { @@ -714,20 +983,27 @@ func (t *ExecTool) guardCommand(command, cwd string) string { if len(t.allowRules) > 0 { if !matchAllowRules(lower, t.allowRules) { var b strings.Builder + b.WriteString("Command blocked: not in allowlist [") + for i, rule := range t.allowRules { if i > 0 { b.WriteByte(',') } + b.WriteString(strings.Join(rule, " ")) } + b.WriteByte(']') + return b.String() } } // Restrict curl/wget to localhost and RFC 1918 private addresses. + // External HTTP access is available via the web_fetch tool. + if t.localNetOnly && isCurlOrWget(cmd) { if errMsg := checkCurlLocalNet(cmd); errMsg != "" { return errMsg @@ -745,16 +1021,27 @@ func (t *ExecTool) guardCommand(command, cwd string) string { } // Token-based absolute path detection. + // Uses strings.Fields so relative paths (e.g., "tests/cold/file.py") + // are not falsely flagged. + // Flags like -I/usr/local/include are naturally skipped because + // filepath.IsAbs returns false for tokens starting with "-". + // + // Agent CLI tools (claude, codex, gemini) accept slash commands + // (e.g., "/review") that look like absolute paths but are not. + // For these tools we check whether the token is an existing path + // before blocking. + agentCLI := isAgentCLICommand(cmd) + for _, token := range strings.Fields(cmd) { token = strings.Trim(token, "\"'") @@ -763,6 +1050,7 @@ func (t *ExecTool) guardCommand(command, cwd string) string { } p := filepath.Clean(token) + rel, err := filepath.Rel(cwdPath, p) if err != nil { continue @@ -770,22 +1058,31 @@ func (t *ExecTool) guardCommand(command, cwd string) string { if strings.HasPrefix(rel, "..") { // Path is outside workspace — allow if it's an executable binary + if isExecutable(p) { continue } + // Allow /dev/* paths (e.g. /dev/null, /dev/urandom). + // Device files are not regular filesystem paths and pose + // no workspace-escape risk. + if strings.HasPrefix(p, "/dev/") { continue } + // Agent CLI slash commands: skip non-existent paths + // (e.g., "/review" is a command, not a file). + if agentCLI { if _, statErr := os.Stat(p); os.IsNotExist(statErr) { continue } } + return fmt.Sprintf("Command blocked: path outside working dir %s", p) } } @@ -795,43 +1092,59 @@ func (t *ExecTool) guardCommand(command, cwd string) string { } // agentCLINames lists agent CLI tools that use slash commands + // (e.g., "/review", "/help") which look like absolute paths. + var agentCLINames = []string{"claude", "codex", "gemini"} // isAgentCLICommand returns true if the command invokes an agent CLI tool. + func isAgentCLICommand(cmd string) bool { fields := strings.Fields(cmd) + if len(fields) == 0 { return false } + base := filepath.Base(fields[0]) + for _, name := range agentCLINames { if base == name { return true } } + return false } // isExecutable checks if a path points to an executable file. + // On Unix, checks the execute permission bits. + // On Windows, checks for known executable extensions. + func isExecutable(path string) bool { info, err := os.Stat(path) if err != nil { return false } + if info.IsDir() { return false } + if runtime.GOOS == "windows" { ext := strings.ToLower(filepath.Ext(path)) + switch ext { case ".exe", ".cmd", ".bat", ".ps1", ".com": + return true } + return false } + return info.Mode()&0o111 != 0 } @@ -844,12 +1157,17 @@ func (t *ExecTool) SetRestrictToWorkspace(restrict bool) { } // SetAllowRules sets the command prefix allowlist. + // Each rule is a space-separated command prefix (e.g. "go test", "pnpm run lint"). + // A command is allowed if its first N words match any rule's N words exactly. + func (t *ExecTool) SetAllowRules(rules []string) { t.allowRules = make([][]string, 0, len(rules)) + for _, r := range rules { words := strings.Fields(strings.ToLower(r)) + if len(words) > 0 { t.allowRules = append(t.allowRules, words) } @@ -857,23 +1175,30 @@ func (t *ExecTool) SetAllowRules(rules []string) { } // matchAllowRules checks if cmd matches any prefix in the allowlist. + func matchAllowRules(cmd string, rules [][]string) bool { cmdWords := strings.Fields(cmd) + for _, ruleWords := range rules { if len(cmdWords) < len(ruleWords) { continue } + match := true + for i, rw := range ruleWords { if cmdWords[i] != rw { match = false + break } } + if match { return true } } + return false } @@ -882,60 +1207,84 @@ func (t *ExecTool) SetLocalNetOnly(v bool) { } // isCurlOrWget reports whether command is a curl or wget invocation. + func isCurlOrWget(command string) bool { fields := strings.Fields(command) + if len(fields) == 0 { return false } + base := filepath.Base(fields[0]) + return base == "curl" || base == "wget" } // checkCurlLocalNet validates that all http/https URLs in a curl/wget command + // target localhost or RFC 1918 private addresses. + // Returns an error message string, or empty string if the command is allowed. + func checkCurlLocalNet(command string) string { for _, token := range strings.Fields(command) { token = strings.Trim(token, "\"'") + if !strings.HasPrefix(token, "http://") && !strings.HasPrefix(token, "https://") { continue } + u, err := url.Parse(token) if err != nil { continue } + host := u.Hostname() + if !isLocalHost(host) { return fmt.Sprintf( + "Command blocked by safety guard "+ + "(curl/wget is restricted to localhost and private network; %q is a public address)", + host, ) } } + return "" } // isLocalHost reports whether host is localhost or a loopback/RFC 1918 private IP. + // DNS resolution is intentionally avoided to prevent DNS rebinding attacks. + func isLocalHost(host string) bool { if strings.EqualFold(host, "localhost") { return true } + ip := net.ParseIP(host) + if ip == nil { return false } + return ip.IsLoopback() || ip.IsPrivate() } // SetBgMaxLifetimeForTest overrides bgMaxLifetime for testing purposes. + // This is exposed only for tests; the returned function restores the original value. + var bgMaxLifetimeOverride time.Duration func SetBgMaxLifetimeForTest(d time.Duration) func() { old := bgMaxLifetimeOverride + bgMaxLifetimeOverride = d + return func() { bgMaxLifetimeOverride = old } } @@ -943,5 +1292,6 @@ func getBgMaxLifetime() time.Duration { if bgMaxLifetimeOverride > 0 { return bgMaxLifetimeOverride } + return bgMaxLifetime } diff --git a/pkg/tools/shell_process_unix.go b/pkg/tools/shell_process_unix.go index fa96d75da..d513f87ad 100644 --- a/pkg/tools/shell_process_unix.go +++ b/pkg/tools/shell_process_unix.go @@ -11,75 +11,129 @@ import ( ) func prepareCommandForTermination(cmd *exec.Cmd) { + if cmd == nil { + return + } + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + } func terminateProcessTree(cmd *exec.Cmd) error { + if cmd == nil || cmd.Process == nil { + return nil + } pid := cmd.Process.Pid + if pid <= 0 { + return nil + } // Kill the entire process group spawned by the shell command. + _ = syscall.Kill(-pid, syscall.SIGKILL) + // Some shells/background jobs may still leave descendants around + // briefly; aggressively walk /proc and kill child processes too. + killDescendants(pid) + // Fallback kill on the shell process itself. + _ = cmd.Process.Kill() + return nil + } func killDescendants(ppid int) { + if ppid <= 0 { + return + } entries, err := os.ReadDir("/proc") + if err != nil { + return + } for _, e := range entries { + if !e.IsDir() { + continue + } + childPID, err := strconv.Atoi(e.Name()) + if err != nil || childPID <= 0 || childPID == ppid { + continue + } statPath := "/proc/" + e.Name() + "/stat" + data, err := os.ReadFile(statPath) + if err != nil { + continue + } // /proc/<pid>/stat: pid (comm) state ppid ... + raw := string(data) + end := strings.LastIndex(raw, ")") + if end == -1 || end+2 >= len(raw) { + continue + } + fields := strings.Fields(raw[end+2:]) + if len(fields) < 2 { + continue + } + parent, err := strconv.Atoi(fields[1]) + if err != nil || parent != ppid { + continue + } // Recurse first, then kill child process/group. + killDescendants(childPID) + _ = syscall.Kill(-childPID, syscall.SIGKILL) + _ = syscall.Kill(childPID, syscall.SIGKILL) + } + } diff --git a/pkg/tools/shell_process_windows.go b/pkg/tools/shell_process_windows.go index fe23b5c96..fbd28b0fa 100644 --- a/pkg/tools/shell_process_windows.go +++ b/pkg/tools/shell_process_windows.go @@ -17,11 +17,14 @@ func terminateProcessTree(cmd *exec.Cmd) error { } pid := cmd.Process.Pid + if pid <= 0 { return nil } _ = exec.Command("taskkill", "/T", "/F", "/PID", strconv.Itoa(pid)).Run() + _ = cmd.Process.Kill() + return nil } diff --git a/pkg/tools/shell_test.go b/pkg/tools/shell_test.go index 2203beffd..c4f4530be 100644 --- a/pkg/tools/shell_test.go +++ b/pkg/tools/shell_test.go @@ -12,6 +12,7 @@ import ( ) // TestShellTool_Success verifies successful command execution + func TestShellTool_Success(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -19,6 +20,7 @@ func TestShellTool_Success(t *testing.T) { } ctx := context.Background() + args := map[string]any{ "command": "echo 'hello world'", } @@ -26,22 +28,26 @@ func TestShellTool_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain command output + if !strings.Contains(result.ForUser, "hello world") { t.Errorf("Expected ForUser to contain 'hello world', got: %s", result.ForUser) } // ForLLM should contain full output + if !strings.Contains(result.ForLLM, "hello world") { t.Errorf("Expected ForLLM to contain 'hello world', got: %s", result.ForLLM) } } // TestShellTool_Failure verifies failed command execution + func TestShellTool_Failure(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -49,6 +55,7 @@ func TestShellTool_Failure(t *testing.T) { } ctx := context.Background() + args := map[string]any{ "command": "ls /nonexistent_directory_12345", } @@ -56,22 +63,26 @@ func TestShellTool_Failure(t *testing.T) { result := tool.Execute(ctx, args) // Failure should be marked as error + if !result.IsError { t.Errorf("Expected error for failed command, got IsError=false") } // ForUser should contain error information + if result.ForUser == "" { t.Errorf("Expected ForUser to contain error info, got empty string") } // ForLLM should contain exit code or error + if !strings.Contains(result.ForLLM, "Exit code") && result.ForUser == "" { t.Errorf("Expected ForLLM to contain exit code or error, got: %s", result.ForLLM) } } // TestShellTool_Timeout verifies command timeout handling + func TestShellTool_Timeout(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -81,6 +92,7 @@ func TestShellTool_Timeout(t *testing.T) { tool.SetTimeout(100 * time.Millisecond) ctx := context.Background() + args := map[string]any{ "command": "sleep 10", } @@ -88,21 +100,27 @@ func TestShellTool_Timeout(t *testing.T) { result := tool.Execute(ctx, args) // Timeout should be marked as error + if !result.IsError { t.Errorf("Expected error for timeout, got IsError=false") } // Should mention timeout + if !strings.Contains(result.ForLLM, "timed out") && !strings.Contains(result.ForUser, "timed out") { t.Errorf("Expected timeout message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser) } } // TestShellTool_WorkingDir verifies custom working directory + func TestShellTool_WorkingDir(t *testing.T) { // Create temp directory + tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "test.txt") + os.WriteFile(testFile, []byte("test content"), 0o644) tool, err := NewExecTool("", false) @@ -111,8 +129,10 @@ func TestShellTool_WorkingDir(t *testing.T) { } ctx := context.Background() + args := map[string]any{ - "command": "cat test.txt", + "command": "cat test.txt", + "working_dir": tmpDir, } @@ -128,6 +148,7 @@ func TestShellTool_WorkingDir(t *testing.T) { } // TestShellTool_DangerousCommand verifies safety guard blocks dangerous commands + func TestShellTool_DangerousCommand(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -135,6 +156,7 @@ func TestShellTool_DangerousCommand(t *testing.T) { } ctx := context.Background() + args := map[string]any{ "command": "rm -rf /", } @@ -142,6 +164,7 @@ func TestShellTool_DangerousCommand(t *testing.T) { result := tool.Execute(ctx, args) // Dangerous command should be blocked + if !result.IsError { t.Errorf("Expected dangerous command to be blocked (IsError=true)") } @@ -152,6 +175,7 @@ func TestShellTool_DangerousCommand(t *testing.T) { } // TestShellTool_MissingCommand verifies error handling for missing command + func TestShellTool_MissingCommand(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -159,17 +183,20 @@ func TestShellTool_MissingCommand(t *testing.T) { } ctx := context.Background() + args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when command is missing") } } // TestShellTool_StderrCapture verifies stderr is captured and included + func TestShellTool_StderrCapture(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -177,6 +204,7 @@ func TestShellTool_StderrCapture(t *testing.T) { } ctx := context.Background() + args := map[string]any{ "command": "sh -c 'echo stdout; echo stderr >&2'", } @@ -184,15 +212,18 @@ func TestShellTool_StderrCapture(t *testing.T) { result := tool.Execute(ctx, args) // Both stdout and stderr should be in output + if !strings.Contains(result.ForLLM, "stdout") { t.Errorf("Expected stdout in output, got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "stderr") { t.Errorf("Expected stderr in output, got: %s", result.ForLLM) } } // TestShellTool_OutputTruncation verifies long output is truncated + func TestShellTool_OutputTruncation(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -200,7 +231,9 @@ func TestShellTool_OutputTruncation(t *testing.T) { } ctx := context.Background() + // Generate long output (>10000 chars) + args := map[string]any{ "command": "python3 -c \"print('x' * 20000)\" || echo " + strings.Repeat("x", 20000), } @@ -208,19 +241,25 @@ func TestShellTool_OutputTruncation(t *testing.T) { result := tool.Execute(ctx, args) // Should have truncation message or be truncated + if len(result.ForLLM) > 15000 { t.Errorf("Expected output to be truncated, got length: %d", len(result.ForLLM)) } } // TestShellTool_WorkingDir_OutsideWorkspace verifies that working_dir cannot escape the workspace directly + func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) { root := t.TempDir() + workspace := filepath.Join(root, "workspace") + outsideDir := filepath.Join(root, "outside") + if err := os.MkdirAll(workspace, 0o755); err != nil { t.Fatalf("failed to create workspace: %v", err) } + if err := os.MkdirAll(outsideDir, 0o755); err != nil { t.Fatalf("failed to create outside dir: %v", err) } @@ -231,34 +270,45 @@ func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": "pwd", + "command": "pwd", + "working_dir": outsideDir, }) if !result.IsError { t.Fatalf("expected working_dir outside workspace to be blocked, got output: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "blocked") { t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM) } } // TestShellTool_WorkingDir_SymlinkEscape verifies that a symlink inside the workspace + // pointing outside cannot be used as working_dir to escape the sandbox. + func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) { root := t.TempDir() + workspace := filepath.Join(root, "workspace") + secretDir := filepath.Join(root, "secret") + if err := os.MkdirAll(workspace, 0o755); err != nil { t.Fatalf("failed to create workspace: %v", err) } + if err := os.MkdirAll(secretDir, 0o755); err != nil { t.Fatalf("failed to create secret dir: %v", err) } + os.WriteFile(filepath.Join(secretDir, "secret.txt"), []byte("top secret"), 0o644) // symlink lives inside the workspace but resolves to secretDir outside it + link := filepath.Join(workspace, "escape") + if err := os.Symlink(secretDir, link); err != nil { t.Skipf("symlinks not supported in this environment: %v", err) } @@ -269,21 +319,25 @@ func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": "cat secret.txt", + "command": "cat secret.txt", + "working_dir": link, }) if !result.IsError { t.Fatalf("expected symlink working_dir escape to be blocked, got output: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "blocked") { t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM) } } // TestShellTool_RestrictToWorkspace verifies workspace restriction + func TestShellTool_RestrictToWorkspace(t *testing.T) { tmpDir := t.TempDir() + tool, err := NewExecTool(tmpDir, false) if err != nil { t.Errorf("unable to configure exec tool: %s", err) @@ -292,6 +346,7 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) { tool.SetRestrictToWorkspace(true) ctx := context.Background() + args := map[string]any{ "command": "cat ../../etc/passwd", } @@ -299,14 +354,18 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) { result := tool.Execute(ctx, args) // Path traversal should be blocked + if !result.IsError { t.Errorf("Expected path traversal to be blocked with restrictToWorkspace=true") } if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") { t.Errorf( + "Expected 'blocked' message for path traversal, got ForLLM: %s, ForUser: %s", + result.ForLLM, + result.ForUser, ) } @@ -315,23 +374,33 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) { // --- guardCommand unit tests --- // TestGuardCommand_RelativePathWithSlashes verifies that relative paths + // containing slashes (e.g., tests/cold/test.py, projects/terra-py-form) + // are NOT falsely blocked. This was a regression caused by the old regex + // matching "/cold/test.py" from "tests/cold/test.py" as an absolute path. + func TestGuardCommand_RelativePathWithSlashes(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) cmds := []string{ "pytest tests/cold/test_solver.py -v --tb=short", + "cd projects/terra-py-form && pytest", + "uv run pytest tests/cold/test_solver.py -v --tb=short", + "cat src/terra_py_form/cold/parser.py", + "python src/main.py --config config/dev.json", } for _, cmd := range cmds { result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("Relative path should not be blocked: %q → %s", cmd, result) } @@ -339,19 +408,25 @@ func TestGuardCommand_RelativePathWithSlashes(t *testing.T) { } // TestGuardCommand_VenvBinary verifies that .venv/bin/... paths are allowed + // (they are relative paths, not absolute). + func TestGuardCommand_VenvBinary(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) cmds := []string{ ".venv/bin/python -m pytest", + ".venv/bin/pytest tests/ -v", + ".venv/bin/pip install -e .", } for _, cmd := range cmds { result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("Venv relative path should not be blocked: %q → %s", cmd, result) } @@ -359,111 +434,147 @@ func TestGuardCommand_VenvBinary(t *testing.T) { } // TestGuardCommand_ExecutableBinaryAllowed verifies that absolute paths + // to executable files outside the workspace are allowed (system binaries). + func TestGuardCommand_ExecutableBinaryAllowed(t *testing.T) { if runtime.GOOS == "windows" { t.Skip("Unix executable permission test not applicable on Windows") } workspace := t.TempDir() + externalDir := t.TempDir() // Create a fake executable outside the workspace + execPath := filepath.Join(externalDir, "mybin") + os.WriteFile(execPath, []byte("#!/bin/sh\necho ok"), 0o755) tool, _ := NewExecTool(workspace, true) cmd := execPath + " --help" + result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("Executable binary outside workspace should be allowed: %q → %s", cmd, result) } } // TestGuardCommand_ExecutableBinaryAllowed_Windows verifies that .exe files + // outside the workspace are allowed on Windows. + func TestGuardCommand_ExecutableBinaryAllowed_Windows(t *testing.T) { if runtime.GOOS != "windows" { t.Skip("Windows-specific test") } workspace := t.TempDir() + externalDir := t.TempDir() // Create a fake .exe outside the workspace + execPath := filepath.Join(externalDir, "tool.exe") + os.WriteFile(execPath, []byte("MZ"), 0o644) tool, _ := NewExecTool(workspace, true) cmd := execPath + " --version" + result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("Windows .exe outside workspace should be allowed: %q → %s", cmd, result) } } // TestGuardCommand_NonExecutableOutsideBlocked verifies that non-executable + // files outside the workspace are blocked (e.g., reading /etc/shadow). + func TestGuardCommand_NonExecutableOutsideBlocked(t *testing.T) { if runtime.GOOS == "windows" { t.Skip("Unix permission test not applicable on Windows") } workspace := t.TempDir() + externalDir := t.TempDir() // Create a regular (non-executable) file outside workspace + dataFile := filepath.Join(externalDir, "secret.txt") + os.WriteFile(dataFile, []byte("secret data"), 0o644) tool, _ := NewExecTool(workspace, true) cmd := "cat " + dataFile + result := tool.guardCommand(cmd, workspace) + if result == "" { t.Errorf("Non-executable file outside workspace should be blocked: %q", cmd) } + if !strings.Contains(result, "path outside working dir") { t.Errorf("Expected 'path outside working dir' message, got: %s", result) } } // TestGuardCommand_NonExistentAbsolutePathBlocked verifies that absolute + // paths that don't exist are blocked (could be file creation outside workspace). + func TestGuardCommand_NonExistentAbsolutePathBlocked(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) // Use platform-appropriate absolute path + var cmd string + if runtime.GOOS == "windows" { cmd = "echo hello > C:\\nonexistent_picoclaw_test_output" } else { cmd = "echo hello > /tmp/nonexistent_picoclaw_test_output" } + result := tool.guardCommand(cmd, workspace) + if result == "" { t.Errorf("Non-existent absolute path outside workspace should be blocked: %q", cmd) } } // TestGuardCommand_FlagEmbeddedPathSkipped verifies that paths embedded in + // flags (e.g., -I/usr/local/include) are NOT extracted as absolute paths + // because the token starts with "-", not "/". + func TestGuardCommand_FlagEmbeddedPathSkipped(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) cmds := []string{ "gcc -I/usr/local/include -L/usr/lib main.c", + "g++ -std=c++17 -I/opt/include file.cpp", + "python --prefix=/usr/local script.py", } for _, cmd := range cmds { result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("Flag-embedded path should not be blocked: %q → %s", cmd, result) } @@ -471,38 +582,51 @@ func TestGuardCommand_FlagEmbeddedPathSkipped(t *testing.T) { } // TestGuardCommand_AbsolutePathInsideWorkspace verifies that absolute paths + // within the workspace are always allowed. + func TestGuardCommand_AbsolutePathInsideWorkspace(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) innerDir := filepath.Join(workspace, "projects", "myapp") + os.MkdirAll(innerDir, 0o755) cmd := "ls " + innerDir + result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("Absolute path inside workspace should be allowed: %q → %s", cmd, result) } } // TestGuardCommand_PathTraversal verifies that various path traversal + // patterns are blocked. + func TestGuardCommand_PathTraversal(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) cmds := []string{ "cat ../../etc/passwd", + "cat ../../../etc/shadow", + "ls projects/../../../../etc", } for _, cmd := range cmds { result := tool.guardCommand(cmd, workspace) + if result == "" { t.Errorf("Path traversal should be blocked: %q", cmd) } + if !strings.Contains(result, "path traversal") { t.Errorf("Expected 'path traversal' message, got: %s", result) } @@ -510,16 +634,22 @@ func TestGuardCommand_PathTraversal(t *testing.T) { } // TestGuardCommand_CdWithAbsoluteWorkspacePath verifies that cd to an + // absolute path within the workspace followed by other commands is allowed. + func TestGuardCommand_CdWithAbsoluteWorkspacePath(t *testing.T) { workspace := t.TempDir() + innerDir := filepath.Join(workspace, "projects", "foo") + os.MkdirAll(innerDir, 0o755) tool, _ := NewExecTool(workspace, true) cmd := "cd " + innerDir + " && ls -la" + result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("cd to workspace subdir should be allowed: %q → %s", cmd, result) } @@ -527,26 +657,36 @@ func TestGuardCommand_CdWithAbsoluteWorkspacePath(t *testing.T) { func TestGuardCommand_AgentCLISlashCommand(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) // Agent CLI slash commands (e.g., "/review") are not file paths. + // They should be allowed because they don't exist on disk. + cmds := []string{ `codex exec --yolo "/review skip-git-repo-check"`, + `claude "/review"`, + `gemini "/help"`, } + for _, cmd := range cmds { result := tool.guardCommand(cmd, workspace) + if result != "" { t.Errorf("Agent CLI slash command should not be blocked: %q → %s", cmd, result) } } // Non-agent commands with absolute paths should still be blocked. + if runtime.GOOS != "windows" { blocked := `cat /etc/hosts` + result := tool.guardCommand(blocked, workspace) + if result == "" { t.Errorf("Non-agent command with absolute path should be blocked: %q", blocked) } @@ -554,65 +694,88 @@ func TestGuardCommand_AgentCLISlashCommand(t *testing.T) { } // TestGuardCommand_DenyPattern_IncludesPattern verifies that deny-match + // error messages include the matched pattern string. + func TestGuardCommand_DenyPattern_IncludesPattern(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) + // Also add a custom deny pattern for precise matching. + tool.denyPatterns = append(tool.denyPatterns, regexp.MustCompile(`\bdangerous_cmd\b`)) result := tool.guardCommand("dangerous_cmd --force", workspace) + if result == "" { t.Fatal("expected deny pattern to block the command") } + if !strings.Contains(result, "deny pattern") { t.Errorf("expected 'deny pattern' in message, got: %s", result) } + if !strings.Contains(result, `\bdangerous_cmd\b`) { t.Errorf("expected pattern string in message, got: %s", result) } } // TestGuardCommand_Allowlist_ShowsRules verifies that allowlist violation + // messages include all configured rules. + func TestGuardCommand_Allowlist_ShowsRules(t *testing.T) { workspace := t.TempDir() + tool, _ := NewExecTool(workspace, true) + tool.SetAllowRules([]string{"go test", "git"}) result := tool.guardCommand("curl http://example.com", workspace) + if result == "" { t.Fatal("expected allowlist to block the command") } + if !strings.Contains(result, "not in allowlist") { t.Errorf("expected 'not in allowlist' in message, got: %s", result) } + if !strings.Contains(result, "go test") || !strings.Contains(result, "git") { t.Errorf("expected allowlist rules in message, got: %s", result) } } // TestGuardCommand_PathOutside_IncludesPath verifies that workspace-escape + // messages include the offending path token. + func TestGuardCommand_PathOutside_IncludesPath(t *testing.T) { if runtime.GOOS == "windows" { t.Skip("Unix absolute path test not applicable on Windows") } workspace := t.TempDir() + externalDir := t.TempDir() + dataFile := filepath.Join(externalDir, "secret.txt") + os.WriteFile(dataFile, []byte("secret"), 0o644) tool, _ := NewExecTool(workspace, true) result := tool.guardCommand("cat "+dataFile, workspace) + if result == "" { t.Fatal("expected path outside workspace to be blocked") } + if !strings.Contains(result, "path outside working dir") { t.Errorf("expected 'path outside working dir' in message, got: %s", result) } + if !strings.Contains(result, dataFile) { t.Errorf("expected offending path %q in message, got: %s", dataFile, result) } @@ -622,9 +785,11 @@ func TestGuardCommand_PathOutside_IncludesPath(t *testing.T) { func TestExecTool_Bg_StartAndOutput(t *testing.T) { tool, _ := NewExecTool("", false) + defer tool.Shutdown() var cmd string + if runtime.GOOS == "windows" { cmd = "Write-Output 'hello from bg'; Start-Sleep -Seconds 30" } else { @@ -632,30 +797,39 @@ func TestExecTool_Bg_StartAndOutput(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if result.IsError { t.Fatalf("failed to start bg process: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "bg-1") { t.Errorf("expected bg-1 in result, got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "Background process started") { t.Errorf("expected start message, got: %s", result.ForLLM) } // Get output + outputResult := tool.Execute(context.Background(), map[string]any{ "bg_action": "output", - "bg_id": "bg-1", + + "bg_id": "bg-1", }) + if outputResult.IsError { t.Fatalf("failed to get output: %s", outputResult.ForLLM) } + if !strings.Contains(outputResult.ForLLM, "hello from bg") { t.Errorf("expected 'hello from bg' in output, got: %s", outputResult.ForLLM) } + if !strings.Contains(outputResult.ForLLM, "running") { t.Errorf("expected 'running' status, got: %s", outputResult.ForLLM) } @@ -663,9 +837,11 @@ func TestExecTool_Bg_StartAndOutput(t *testing.T) { func TestExecTool_Bg_Kill(t *testing.T) { tool, _ := NewExecTool("", false) + defer tool.Shutdown() var cmd string + if runtime.GOOS == "windows" { cmd = "Start-Sleep -Seconds 60" } else { @@ -673,27 +849,35 @@ func TestExecTool_Bg_Kill(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if result.IsError { t.Fatalf("failed to start bg process: %s", result.ForLLM) } // Kill it + killResult := tool.Execute(context.Background(), map[string]any{ "bg_action": "kill", - "bg_id": "bg-1", + + "bg_id": "bg-1", }) + if killResult.IsError { t.Fatalf("failed to kill: %s", killResult.ForLLM) } + if !strings.Contains(killResult.ForLLM, "terminated") { t.Errorf("expected 'terminated' message, got: %s", killResult.ForLLM) } // Process should no longer be in the map + procs := tool.BgProcesses() + if _, ok := procs["bg-1"]; ok { t.Errorf("expected bg-1 to be removed after kill") } @@ -701,9 +885,11 @@ func TestExecTool_Bg_Kill(t *testing.T) { func TestExecTool_Bg_ExitedProcess(t *testing.T) { tool, _ := NewExecTool("", false) + defer tool.Shutdown() var cmd string + if runtime.GOOS == "windows" { cmd = "Write-Output 'quick exit'" } else { @@ -711,27 +897,35 @@ func TestExecTool_Bg_ExitedProcess(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if result.IsError { t.Fatalf("failed to start bg process: %s", result.ForLLM) } // Wait for process to exit (initial capture is 3s, so after that it should be done) + time.Sleep(4 * time.Second) // Get output — should show exited + outputResult := tool.Execute(context.Background(), map[string]any{ "bg_action": "output", - "bg_id": "bg-1", + + "bg_id": "bg-1", }) + if outputResult.IsError { t.Fatalf("failed to get output: %s", outputResult.ForLLM) } + if !strings.Contains(outputResult.ForLLM, "exited") { t.Errorf("expected 'exited' in output, got: %s", outputResult.ForLLM) } + if !strings.Contains(outputResult.ForLLM, "quick exit") { t.Errorf("expected 'quick exit' in output, got: %s", outputResult.ForLLM) } @@ -739,25 +933,33 @@ func TestExecTool_Bg_ExitedProcess(t *testing.T) { func TestExecTool_Bg_InvalidID(t *testing.T) { tool, _ := NewExecTool("", false) + defer tool.Shutdown() // Output for non-existent ID + result := tool.Execute(context.Background(), map[string]any{ "bg_action": "output", - "bg_id": "bg-999", + + "bg_id": "bg-999", }) + if !result.IsError { t.Fatalf("expected error for invalid bg_id") } + if !strings.Contains(result.ForLLM, "not found") { t.Errorf("expected 'not found' message, got: %s", result.ForLLM) } // Kill for non-existent ID + result = tool.Execute(context.Background(), map[string]any{ "bg_action": "kill", - "bg_id": "bg-999", + + "bg_id": "bg-999", }) + if !result.IsError { t.Fatalf("expected error for invalid bg_id") } @@ -765,9 +967,11 @@ func TestExecTool_Bg_InvalidID(t *testing.T) { func TestExecTool_Bg_InitialOutputCapture(t *testing.T) { tool, _ := NewExecTool("", false) + defer tool.Shutdown() var cmd string + if runtime.GOOS == "windows" { cmd = "Write-Output 'initial line 1'; Write-Output 'initial line 2'; Start-Sleep -Seconds 30" } else { @@ -775,15 +979,19 @@ func TestExecTool_Bg_InitialOutputCapture(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if result.IsError { t.Fatalf("failed to start bg process: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "initial line 1") { t.Errorf("expected 'initial line 1' in initial output, got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "initial line 2") { t.Errorf("expected 'initial line 2' in initial output, got: %s", result.ForLLM) } @@ -791,14 +999,17 @@ func TestExecTool_Bg_InitialOutputCapture(t *testing.T) { func TestExecTool_Bg_RuntimeStatus(t *testing.T) { tool, _ := NewExecTool("", false) + defer tool.Shutdown() // No bg processes — should return empty + if s := tool.RuntimeStatus(); s != "" { t.Errorf("expected empty runtime status with no bg processes, got: %s", s) } var cmd string + if runtime.GOOS == "windows" { cmd = "Start-Sleep -Seconds 30" } else { @@ -806,17 +1017,21 @@ func TestExecTool_Bg_RuntimeStatus(t *testing.T) { } tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) status := tool.RuntimeStatus() + if !strings.Contains(status, "Background Processes") { t.Errorf("expected 'Background Processes' section, got: %s", status) } + if !strings.Contains(status, "bg-1") { t.Errorf("expected 'bg-1' in status, got: %s", status) } + if !strings.Contains(status, "running") { t.Errorf("expected 'running' in status, got: %s", status) } @@ -826,6 +1041,7 @@ func TestExecTool_Bg_Shutdown(t *testing.T) { tool, _ := NewExecTool("", false) var cmd string + if runtime.GOOS == "windows" { cmd = "Start-Sleep -Seconds 60" } else { @@ -833,16 +1049,21 @@ func TestExecTool_Bg_Shutdown(t *testing.T) { } tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) // Both should be running + procs := tool.BgProcesses() + for _, bp := range procs { if !bp.isRunning() { t.Errorf("expected process to be running before shutdown") @@ -850,10 +1071,13 @@ func TestExecTool_Bg_Shutdown(t *testing.T) { } // Shutdown + tool.Shutdown() // All should be done + procs = tool.BgProcesses() + for _, bp := range procs { if bp.isRunning() { t.Errorf("expected process to be stopped after shutdown") @@ -864,8 +1088,11 @@ func TestExecTool_Bg_Shutdown(t *testing.T) { func TestRingBuffer(t *testing.T) { t.Run("Write and String", func(t *testing.T) { rb := newRingBuffer(100) + rb.Write([]byte("hello ")) + rb.Write([]byte("world")) + if got := rb.String(); got != "hello world" { t.Errorf("expected 'hello world', got %q", got) } @@ -873,11 +1100,15 @@ func TestRingBuffer(t *testing.T) { t.Run("Lines", func(t *testing.T) { rb := newRingBuffer(100) + rb.Write([]byte("line1\nline2\nline3\nline4\nline5\n")) + lines := rb.Lines(3) + if len(lines) != 3 { t.Fatalf("expected 3 lines, got %d", len(lines)) } + if lines[0] != "line3" || lines[1] != "line4" || lines[2] != "line5" { t.Errorf("unexpected lines: %v", lines) } @@ -885,20 +1116,27 @@ func TestRingBuffer(t *testing.T) { t.Run("Match", func(t *testing.T) { rb := newRingBuffer(100) + rb.Write([]byte("starting...\nServer ready on port 3000\nwaiting...\n")) re := regexp.MustCompile(`ready.*port`) + match := rb.Match(re) + if match == "" { t.Fatal("expected match but got empty string") } + if !strings.Contains(match, "ready") { t.Errorf("expected match to contain 'ready', got: %s", match) } // Non-matching pattern + re2 := regexp.MustCompile(`never_match`) + match2 := rb.Match(re2) + if match2 != "" { t.Errorf("expected no match, got: %s", match2) } @@ -906,12 +1144,17 @@ func TestRingBuffer(t *testing.T) { t.Run("Overflow", func(t *testing.T) { rb := newRingBuffer(10) // small buffer + rb.Write([]byte("1234567890ABCDEF")) + got := rb.String() + if len(got) != 10 { t.Errorf("expected buffer to be 10 bytes, got %d", len(got)) } + // Should keep the last 10 bytes + if got != "7890ABCDEF" { t.Errorf("expected '7890ABCDEF', got %q", got) } @@ -919,10 +1162,13 @@ func TestRingBuffer(t *testing.T) { t.Run("Len", func(t *testing.T) { rb := newRingBuffer(100) + if rb.Len() != 0 { t.Errorf("expected 0 length initially") } + rb.Write([]byte("hello")) + if rb.Len() != 5 { t.Errorf("expected 5, got %d", rb.Len()) } @@ -930,7 +1176,9 @@ func TestRingBuffer(t *testing.T) { t.Run("Empty Lines", func(t *testing.T) { rb := newRingBuffer(100) + lines := rb.Lines(5) + if lines != nil { t.Errorf("expected nil for empty buffer, got: %v", lines) } @@ -939,10 +1187,13 @@ func TestRingBuffer(t *testing.T) { func TestExecTool_Bg_RingBufferOverflow(t *testing.T) { tool, _ := NewExecTool("", false) + defer tool.Shutdown() // Generate output larger than 32KB ring buffer + var cmd string + if runtime.GOOS == "windows" { cmd = "1..2000 | ForEach-Object { Write-Output ('x' * 50) }; Start-Sleep -Seconds 30" } else { @@ -950,68 +1201,103 @@ func TestExecTool_Bg_RingBufferOverflow(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, + "command": cmd, + "background": true, }) + if result.IsError { t.Fatalf("failed to start bg process: %s", result.ForLLM) } // Wait for output to accumulate + time.Sleep(5 * time.Second) // Get output — ring buffer should have truncated old data + outputResult := tool.Execute(context.Background(), map[string]any{ "bg_action": "output", - "bg_id": "bg-1", + + "bg_id": "bg-1", }) + if outputResult.IsError { t.Fatalf("failed to get output: %s", outputResult.ForLLM) } // The output should contain data but be bounded by the ring buffer size + procs := tool.BgProcesses() + bp := procs["bg-1"] + if bp == nil { t.Fatal("bg-1 not found") } + bufLen := bp.output.Len() + if bufLen > bgRingBufSize { t.Errorf("ring buffer exceeded max size: %d > %d", bufLen, bgRingBufSize) } } // TestIsLocalHost verifies localhost and RFC 1918 detection using net package. + func TestIsLocalHost(t *testing.T) { tests := []struct { host string + want bool }{ // Loopback / localhost + {"localhost", true}, + {"LOCALHOST", true}, + {"127.0.0.1", true}, + {"127.0.0.2", true}, + {"::1", true}, + // RFC 1918 private ranges + {"10.0.0.1", true}, + {"10.255.255.255", true}, + {"172.16.0.1", true}, + {"172.31.255.255", true}, + {"192.168.0.1", true}, + {"192.168.1.100", true}, + // Public addresses + {"8.8.8.8", false}, + {"1.1.1.1", false}, + {"example.com", false}, + {"api.github.com", false}, + // Edge: non-private but routable private-looking address + {"172.15.255.255", false}, // just below 172.16/12 - {"172.32.0.0", false}, // just above 172.31/12 + + {"172.32.0.0", false}, // just above 172.31/12 + } for _, tt := range tests { got := isLocalHost(tt.host) + if got != tt.want { t.Errorf("isLocalHost(%q) = %v, want %v", tt.host, got, tt.want) } @@ -1019,56 +1305,82 @@ func TestIsLocalHost(t *testing.T) { } // TestCheckCurlLocalNet verifies URL-level enforcement for curl/wget commands. + func TestCheckCurlLocalNet(t *testing.T) { tests := []struct { - cmd string + cmd string + wantErr bool }{ // Allowed: localhost and private IPs + {"curl http://localhost:3000/health", false}, + {"curl -v http://127.0.0.1:8080/api/status", false}, + {"wget http://192.168.1.10/file.bin", false}, + {"curl -X POST http://10.0.0.5:9000/webhook", false}, + // Blocked: public addresses + {"curl http://example.com", true}, + {"wget https://releases.github.com/v1.tar.gz", true}, + {"curl http://8.8.8.8/data", true}, + // Allowed: no http URL (e.g. --help, --version — no network access) + {"curl --help", false}, + {"curl --version", false}, + {"wget --help", false}, } for _, tt := range tests { errMsg := checkCurlLocalNet(tt.cmd) + gotErr := errMsg != "" + if gotErr != tt.wantErr { t.Errorf("checkCurlLocalNet(%q): gotErr=%v wantErr=%v (msg: %q)", + tt.cmd, gotErr, tt.wantErr, errMsg) } } } // TestExecTool_LocalNetOnly verifies curl/wget blocking via SetLocalNetOnly. + func TestExecTool_LocalNetOnly(t *testing.T) { tool, _ := NewExecTool("", false) + tool.SetLocalNetOnly(true) tests := []struct { - cmd string + cmd string + wantErr bool }{ {"curl http://localhost:3000", false}, + {"curl http://example.com", true}, + {"echo hello", false}, // non-curl not affected + } ctx := context.Background() + for _, tt := range tests { result := tool.Execute(ctx, map[string]any{"command": tt.cmd}) + if tt.wantErr && !result.IsError { t.Errorf("cmd %q: expected blocked, but succeeded", tt.cmd) } + if !tt.wantErr && result.IsError && strings.Contains(result.ForLLM, "safety guard") { t.Errorf("cmd %q: expected allowed, but safety guard blocked: %s", tt.cmd, result.ForLLM) } diff --git a/pkg/tools/shell_timeout_unix_test.go b/pkg/tools/shell_timeout_unix_test.go index d0d4a5b9b..5dd861442 100644 --- a/pkg/tools/shell_timeout_unix_test.go +++ b/pkg/tools/shell_timeout_unix_test.go @@ -14,73 +14,122 @@ import ( ) func processRunning(pid int) bool { + if pid <= 0 { + return false + } + // kill(0) can return success for zombie processes too, so inspect /proc + // state and treat zombies as not-running for timeout cleanup assertions. + err := syscall.Kill(pid, 0) + if err != nil && err != syscall.EPERM { + return false + } data, readErr := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat") + if readErr != nil { + return false + } + raw := string(data) + end := strings.LastIndex(raw, ")") + if end == -1 || end+2 >= len(raw) { + return true // best effort fallback + } + fields := strings.Fields(raw[end+2:]) + if len(fields) == 0 { + return true // best effort fallback + } state := fields[0] + return state != "Z" + } func TestShellTool_TimeoutKillsChildProcess(t *testing.T) { + tool, err := NewExecTool(t.TempDir(), false) + if err != nil { + t.Errorf("unable to configure exec tool: %s", err) + } tool.SetTimeout(500 * time.Millisecond) args := map[string]any{ + // Spawn a child process that would outlive the shell unless process-group kill is used. + "command": "sleep 60 & echo $! > child.pid; wait", } result := tool.Execute(context.Background(), args) + if !result.IsError { + t.Fatalf("expected timeout error, got success: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "timed out") { + t.Fatalf("expected timeout message, got: %s", result.ForLLM) + } childPIDPath := filepath.Join(tool.workingDir, "child.pid") + data, err := os.ReadFile(childPIDPath) + if err != nil { + t.Fatalf("failed to read child pid file: %v", err) + } childPID, err := strconv.Atoi(strings.TrimSpace(string(data))) + if err != nil { + t.Fatalf("failed to parse child pid: %v", err) + } deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if !processRunning(childPID) { + return + } + time.Sleep(50 * time.Millisecond) + } t.Fatalf("child process %d is still running after timeout", childPID) + } diff --git a/pkg/tools/skills_install.go b/pkg/tools/skills_install.go index 71bfe730b..89c5dcfe0 100644 --- a/pkg/tools/skills_install.go +++ b/pkg/tools/skills_install.go @@ -16,22 +16,32 @@ import ( ) // InstallSkillTool allows the LLM agent to install skills from registries. + // It shares the same RegistryManager that FindSkillsTool uses, + // so all registries configured in config are available for installation. + type InstallSkillTool struct { registryMgr *skills.RegistryManager - workspace string - mu sync.Mutex + + workspace string + + mu sync.Mutex } // NewInstallSkillTool creates a new InstallSkillTool. + // registryMgr is the shared registry manager (same instance as FindSkillsTool). + // workspace is the root workspace directory; skills install to {workspace}/skills/{slug}/. + func NewInstallSkillTool(registryMgr *skills.RegistryManager, workspace string) *InstallSkillTool { return &InstallSkillTool{ registryMgr: registryMgr, - workspace: workspace, - mu: sync.Mutex{}, + + workspace: workspace, + + mu: sync.Mutex{}, } } @@ -46,151 +56,210 @@ func (t *InstallSkillTool) Description() string { func (t *InstallSkillTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "slug": map[string]any{ - "type": "string", + "type": "string", + "description": "The unique slug of the skill to install (e.g., 'github', 'docker-compose')", }, + "version": map[string]any{ - "type": "string", + "type": "string", + "description": "Specific version to install (optional, defaults to latest)", }, + "registry": map[string]any{ - "type": "string", + "type": "string", + "description": "Registry to install from (required, e.g., 'clawhub')", }, + "force": map[string]any{ - "type": "boolean", + "type": "boolean", + "description": "Force reinstall if skill already exists (default false)", }, }, + "required": []string{"slug", "registry"}, } } func (t *InstallSkillTool) Execute(ctx context.Context, args map[string]any) *ToolResult { // Install lock to prevent concurrent directory operations. + // Ideally this should be done at a `slug` level, currently, its at a `workspace` level. + t.mu.Lock() + defer t.mu.Unlock() // Validate slug + slug, _ := args["slug"].(string) + if err := utils.ValidateSkillIdentifier(slug); err != nil { return ErrorResult(fmt.Sprintf("invalid slug %q: error: %s", slug, err.Error())) } // Validate registry + registryName, _ := args["registry"].(string) + if err := utils.ValidateSkillIdentifier(registryName); err != nil { return ErrorResult(fmt.Sprintf("invalid registry %q: error: %s", registryName, err.Error())) } version, _ := args["version"].(string) + force, _ := args["force"].(bool) // Check if already installed. + skillsDir := filepath.Join(t.workspace, "skills") + targetDir := filepath.Join(skillsDir, slug) if !force { if _, err := os.Stat(targetDir); err == nil { return ErrorResult( + fmt.Sprintf("skill %q already installed at %s. Use force=true to reinstall.", slug, targetDir), ) } } else { // Force: remove existing if present. + os.RemoveAll(targetDir) } // Resolve which registry to use. + registry := t.registryMgr.GetRegistry(registryName) + if registry == nil { return ErrorResult(fmt.Sprintf("registry %q not found", registryName)) } // Ensure skills directory exists. + if err := os.MkdirAll(skillsDir, 0o755); err != nil { return ErrorResult(fmt.Sprintf("failed to create skills directory: %v", err)) } // Download and install (handles metadata, version resolution, extraction). + result, err := registry.DownloadAndInstall(ctx, slug, version, targetDir) if err != nil { // Clean up partial install. + rmErr := os.RemoveAll(targetDir) + if rmErr != nil { logger.ErrorCF("tool", "Failed to remove partial install", + map[string]any{ - "tool": "install_skill", + "tool": "install_skill", + "target_dir": targetDir, - "error": rmErr.Error(), + + "error": rmErr.Error(), }) } + return ErrorResult(fmt.Sprintf("failed to install %q: %v", slug, err)) } // Moderation: block malware. + if result.IsMalwareBlocked { rmErr := os.RemoveAll(targetDir) + if rmErr != nil { logger.ErrorCF("tool", "Failed to remove partial install", + map[string]any{ - "tool": "install_skill", + "tool": "install_skill", + "target_dir": targetDir, - "error": rmErr.Error(), + + "error": rmErr.Error(), }) } + return ErrorResult(fmt.Sprintf("skill %q is flagged as malicious and cannot be installed", slug)) } // Write origin metadata. + if err := writeOriginMeta(targetDir, registry.Name(), slug, result.Version); err != nil { logger.ErrorCF("tool", "Failed to write origin metadata", + map[string]any{ - "tool": "install_skill", - "error": err.Error(), - "target": targetDir, + "tool": "install_skill", + + "error": err.Error(), + + "target": targetDir, + "registry": registry.Name(), - "slug": slug, - "version": result.Version, + + "slug": slug, + + "version": result.Version, }) + _ = err } // Build result with moderation warning if suspicious. + var output string + if result.IsSuspicious { output = fmt.Sprintf("⚠️ Warning: skill %q is flagged as suspicious (may contain risky patterns).\n\n", slug) } + output += fmt.Sprintf("Successfully installed skill %q v%s from %s registry.\nLocation: %s\n", + slug, result.Version, registry.Name(), targetDir) if result.Summary != "" { output += fmt.Sprintf("Description: %s\n", result.Summary) } + output += "\nThe skill is now available and can be loaded in the current session." return SilentResult(output) } // originMeta tracks which registry a skill was installed from. + type originMeta struct { - Version int `json:"version"` - Registry string `json:"registry"` - Slug string `json:"slug"` + Version int `json:"version"` + + Registry string `json:"registry"` + + Slug string `json:"slug"` + InstalledVersion string `json:"installed_version"` - InstalledAt int64 `json:"installed_at"` + + InstalledAt int64 `json:"installed_at"` } func writeOriginMeta(targetDir, registryName, slug, version string) error { meta := originMeta{ - Version: 1, - Registry: registryName, - Slug: slug, + Version: 1, + + Registry: registryName, + + Slug: slug, + InstalledVersion: version, - InstalledAt: time.Now().UnixMilli(), + + InstalledAt: time.Now().UnixMilli(), } data, err := json.MarshalIndent(meta, "", " ") @@ -199,5 +268,6 @@ func writeOriginMeta(targetDir, registryName, slug, version string) error { } // Use unified atomic write utility with explicit sync for flash storage reliability. + return fileutil.WriteFileAtomic(filepath.Join(targetDir, ".skill-origin.json"), data, 0o600) } diff --git a/pkg/tools/skills_install_test.go b/pkg/tools/skills_install_test.go index 676fcecc0..882e446c6 100644 --- a/pkg/tools/skills_install_test.go +++ b/pkg/tools/skills_install_test.go @@ -14,22 +14,29 @@ import ( func TestInstallSkillToolName(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) + assert.Equal(t, "install_skill", tool.Name()) } func TestInstallSkillToolMissingSlug(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) + result := tool.Execute(context.Background(), map[string]any{}) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "identifier is required and must be a non-empty string") } func TestInstallSkillToolEmptySlug(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) + result := tool.Execute(context.Background(), map[string]any{ "slug": " ", }) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "identifier is required and must be a non-empty string") } @@ -38,7 +45,9 @@ func TestInstallSkillToolUnsafeSlug(t *testing.T) { cases := []string{ "../etc/passwd", + "path/traversal", + "path\\traversal", } @@ -46,59 +55,85 @@ func TestInstallSkillToolUnsafeSlug(t *testing.T) { result := tool.Execute(context.Background(), map[string]any{ "slug": slug, }) + assert.True(t, result.IsError, "slug %q should be rejected", slug) + assert.Contains(t, result.ForLLM, "invalid slug") } } func TestInstallSkillToolAlreadyExists(t *testing.T) { workspace := t.TempDir() + skillDir := filepath.Join(workspace, "skills", "existing-skill") + require.NoError(t, os.MkdirAll(skillDir, 0o755)) tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace) + result := tool.Execute(context.Background(), map[string]any{ - "slug": "existing-skill", + "slug": "existing-skill", + "registry": "clawhub", }) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "already installed") } func TestInstallSkillToolRegistryNotFound(t *testing.T) { workspace := t.TempDir() + tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace) + result := tool.Execute(context.Background(), map[string]any{ - "slug": "some-skill", + "slug": "some-skill", + "registry": "nonexistent", }) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "registry") + assert.Contains(t, result.ForLLM, "not found") } func TestInstallSkillToolParameters(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) + params := tool.Parameters() props, ok := params["properties"].(map[string]any) + assert.True(t, ok) + assert.Contains(t, props, "slug") + assert.Contains(t, props, "version") + assert.Contains(t, props, "registry") + assert.Contains(t, props, "force") required, ok := params["required"].([]string) + assert.True(t, ok) + assert.Contains(t, required, "slug") + assert.Contains(t, required, "registry") } func TestInstallSkillToolMissingRegistry(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) + result := tool.Execute(context.Background(), map[string]any{ "slug": "some-skill", }) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "invalid registry") } diff --git a/pkg/tools/skills_search.go b/pkg/tools/skills_search.go index 2b6cffd38..48baff7b6 100644 --- a/pkg/tools/skills_search.go +++ b/pkg/tools/skills_search.go @@ -9,18 +9,24 @@ import ( ) // FindSkillsTool allows the LLM agent to search for installable skills from registries. + type FindSkillsTool struct { registryMgr *skills.RegistryManager - cache *skills.SearchCache + + cache *skills.SearchCache } // NewFindSkillsTool creates a new FindSkillsTool. + // registryMgr is the shared registry manager (built from config in createToolRegistry). + // cache is the search cache for deduplicating similar queries. + func NewFindSkillsTool(registryMgr *skills.RegistryManager, cache *skills.SearchCache) *FindSkillsTool { return &FindSkillsTool{ registryMgr: registryMgr, - cache: cache, + + cache: cache, } } @@ -35,38 +41,50 @@ func (t *FindSkillsTool) Description() string { func (t *FindSkillsTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "query": map[string]any{ - "type": "string", + "type": "string", + "description": "Search query describing the desired skill capability (e.g., 'github integration', 'database management')", }, + "limit": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Maximum number of results to return (1-20, default 5)", - "minimum": 1.0, - "maximum": 20.0, + + "minimum": 1.0, + + "maximum": 20.0, }, }, + "required": []string{"query"}, } } func (t *FindSkillsTool) Execute(ctx context.Context, args map[string]any) *ToolResult { query, ok := args["query"].(string) + query = strings.ToLower(strings.TrimSpace(query)) + if !ok || query == "" { return ErrorResult("query is required and must be a non-empty string") } limit := 5 + if l, ok := args["limit"].(float64); ok { li := int(l) + if li >= 1 && li <= 20 { limit = li } } // Check cache first. + if t.cache != nil { if cached, hit := t.cache.Get(query); hit { return SilentResult(formatSearchResults(query, cached, true)) @@ -74,12 +92,14 @@ func (t *FindSkillsTool) Execute(ctx context.Context, args map[string]any) *Tool } // Search all registries. + results, err := t.registryMgr.SearchAll(ctx, query, limit) if err != nil { return ErrorResult(fmt.Sprintf("skill search failed: %v", err)) } // Cache the results. + if t.cache != nil && len(results) > 0 { t.cache.Put(query, results) } @@ -93,27 +113,36 @@ func formatSearchResults(query string, results []skills.SearchResult, cached boo } var sb strings.Builder + source := "" + if cached { source = " (cached)" } + sb.WriteString(fmt.Sprintf("Found %d skills for %q%s:\n\n", len(results), query, source)) for i, r := range results { sb.WriteString(fmt.Sprintf("%d. **%s**", i+1, r.Slug)) + if r.Version != "" { sb.WriteString(fmt.Sprintf(" v%s", r.Version)) } + sb.WriteString(fmt.Sprintf(" (score: %.3f, registry: %s)\n", r.Score, r.RegistryName)) + if r.DisplayName != "" && r.DisplayName != r.Slug { sb.WriteString(fmt.Sprintf(" Name: %s\n", r.DisplayName)) } + if r.Summary != "" { sb.WriteString(fmt.Sprintf(" %s\n", r.Summary)) } + sb.WriteString("\n") } sb.WriteString("Use install_skill with the slug to install a skill.") + return sb.String() } diff --git a/pkg/tools/skills_search_test.go b/pkg/tools/skills_search_test.go index 0e5387cf5..cb6a0e104 100644 --- a/pkg/tools/skills_search_test.go +++ b/pkg/tools/skills_search_test.go @@ -11,80 +11,110 @@ import ( func TestFindSkillsToolName(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) + assert.Equal(t, "find_skills", tool.Name()) } func TestFindSkillsToolMissingQuery(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) + result := tool.Execute(context.Background(), map[string]any{}) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "query is required") } func TestFindSkillsToolEmptyQuery(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) + result := tool.Execute(context.Background(), map[string]any{ "query": " ", }) + assert.True(t, result.IsError) } func TestFindSkillsToolCacheHit(t *testing.T) { cache := skills.NewSearchCache(10, 5*60*1000*1000*1000) // 5 min + cache.Put("github", []skills.SearchResult{ {Slug: "github", Score: 0.9, RegistryName: "clawhub"}, }) tool := NewFindSkillsTool(skills.NewRegistryManager(), cache) + result := tool.Execute(context.Background(), map[string]any{ "query": "github", }) assert.False(t, result.IsError) + assert.Contains(t, result.ForLLM, "github") + assert.Contains(t, result.ForLLM, "cached") } func TestFindSkillsToolParameters(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) + params := tool.Parameters() props, ok := params["properties"].(map[string]any) + assert.True(t, ok) + assert.Contains(t, props, "query") + assert.Contains(t, props, "limit") required, ok := params["required"].([]string) + assert.True(t, ok) + assert.Contains(t, required, "query") } func TestFindSkillsToolDescription(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) + assert.NotEmpty(t, tool.Description()) + assert.Contains(t, tool.Description(), "skill") } func TestFormatSearchResultsEmpty(t *testing.T) { result := formatSearchResults("test query", nil, false) + assert.Contains(t, result, "No skills found") } func TestFormatSearchResultsWithData(t *testing.T) { results := []skills.SearchResult{ { - Slug: "github", - Score: 0.95, - DisplayName: "GitHub", - Summary: "GitHub API integration", - Version: "1.0.0", + Slug: "github", + + Score: 0.95, + + DisplayName: "GitHub", + + Summary: "GitHub API integration", + + Version: "1.0.0", + RegistryName: "clawhub", }, } + output := formatSearchResults("github", results, false) + assert.Contains(t, output, "github") + assert.Contains(t, output, "v1.0.0") + assert.Contains(t, output, "0.950") + assert.Contains(t, output, "clawhub") + assert.Contains(t, output, "install_skill") } diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index f36542ec1..35eb2c0a9 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -8,31 +8,41 @@ import ( func TestSpawnTool_Execute_EmptyTask(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{}) + tool := NewSpawnTool(manager) ctx := context.Background() tests := []struct { name string + args map[string]any }{ {"empty string", map[string]any{"task": ""}}, + {"whitespace only", map[string]any{"task": " "}}, + {"tabs and newlines", map[string]any{"task": "\t\n "}}, + {"missing task key", map[string]any{"label": "test"}}, + {"wrong type", map[string]any{"task": 123}}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := tool.Execute(ctx, tt.args) + if result == nil { t.Fatal("Result should not be nil") } + if !result.IsError { t.Error("Expected error for invalid task parameter") } + if !strings.Contains(result.ForLLM, `"task"`) { t.Errorf("Error message should mention '\"task\"', got: %s", result.ForLLM) } @@ -42,22 +52,29 @@ func TestSpawnTool_Execute_EmptyTask(t *testing.T) { func TestSpawnTool_Execute_ValidTask(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{}) + tool := NewSpawnTool(manager) ctx := context.Background() + args := map[string]any{ - "task": "Write a haiku about coding", + "task": "Write a haiku about coding", + "label": "haiku-task", } result := tool.Execute(ctx, args) + if result == nil { t.Fatal("Result should not be nil") } + if result.IsError { t.Errorf("Expected success for valid task, got error: %s", result.ForLLM) } + if !result.Async { t.Error("SpawnTool should return async result") } @@ -67,12 +84,15 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) { tool := NewSpawnTool(nil) ctx := context.Background() + args := map[string]any{"task": "test task"} result := tool.Execute(ctx, args) + if !result.IsError { t.Error("Expected error for nil manager") } + if !strings.Contains(result.ForLLM, "spawn tool is not available") { t.Errorf("Error message should mention spawn tool not available, got: %s", result.ForLLM) } diff --git a/pkg/tools/spi.go b/pkg/tools/spi.go index 0ca17e84f..b8c9b3d74 100644 --- a/pkg/tools/spi.go +++ b/pkg/tools/spi.go @@ -10,6 +10,7 @@ import ( ) // SPITool provides SPI bus interaction for high-speed peripheral communication. + type SPITool struct{} func NewSPITool() *SPITool { @@ -27,42 +28,61 @@ func (t *SPITool) Description() string { func (t *SPITool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "action": map[string]any{ - "type": "string", - "enum": []string{"list", "transfer", "read"}, + "type": "string", + + "enum": []string{"list", "transfer", "read"}, + "description": "Action to perform: list (find available SPI devices), transfer (full-duplex send/receive), read (receive bytes by sending zeros)", }, + "device": map[string]any{ - "type": "string", + "type": "string", + "description": "SPI device identifier (e.g. \"2.0\" for /dev/spidev2.0). Required for transfer/read.", }, + "speed": map[string]any{ - "type": "integer", + "type": "integer", + "description": "SPI clock speed in Hz. Default: 1000000 (1 MHz).", }, + "mode": map[string]any{ - "type": "integer", + "type": "integer", + "description": "SPI mode (0-3). Default: 0. Mode sets CPOL and CPHA: 0=0,0 1=0,1 2=1,0 3=1,1.", }, + "bits": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Bits per word. Default: 8.", }, + "data": map[string]any{ - "type": "array", - "items": map[string]any{"type": "integer"}, + "type": "array", + + "items": map[string]any{"type": "integer"}, + "description": "Bytes to send (0-255 each). Required for transfer action.", }, + "length": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Number of bytes to read (1-4096). Required for read action.", }, + "confirm": map[string]any{ - "type": "boolean", + "type": "boolean", + "description": "Must be true for transfer operations. Safety guard to prevent accidental writes.", }, }, + "required": []string{"action"}, } } @@ -73,23 +93,32 @@ func (t *SPITool) Execute(ctx context.Context, args map[string]any) *ToolResult } action, ok := args["action"].(string) + if !ok { return ErrorResult("action is required") } switch action { case "list": + return t.list() + case "transfer": + return t.transfer(args) + case "read": + return t.readDevice(args) + default: + return ErrorResult(fmt.Sprintf("unknown action: %s (valid: list, transfer, read)", action)) } } // list finds available SPI devices by globbing /dev/spidev* + func (t *SPITool) list() *ToolResult { matches, err := filepath.Glob("/dev/spidev*") if err != nil { @@ -103,12 +132,15 @@ func (t *SPITool) list() *ToolResult { } type devInfo struct { - Path string `json:"path"` + Path string `json:"path"` + Device string `json:"device"` } devices := make([]devInfo, 0, len(matches)) + re := regexp.MustCompile(`/dev/spidev(\d+\.\d+)`) + for _, m := range matches { if sub := re.FindStringSubmatch(m); sub != nil { devices = append(devices, devInfo{Path: m, Device: sub[1]}) @@ -116,45 +148,58 @@ func (t *SPITool) list() *ToolResult { } result, _ := json.MarshalIndent(devices, "", " ") + return SilentResult(fmt.Sprintf("Found %d SPI device(s):\n%s", len(devices), string(result))) } // Helper function for SPI operations (used by platform-specific implementations) // parseSPIArgs extracts and validates common SPI parameters + // + //nolint:unused // Used by spi_linux.go + func parseSPIArgs(args map[string]any) (device string, speed uint32, mode uint8, bits uint8, errMsg string) { dev, ok := args["device"].(string) + if !ok || dev == "" { return "", 0, 0, 0, "device is required (e.g. \"2.0\" for /dev/spidev2.0)" } + matched, _ := regexp.MatchString(`^\d+\.\d+$`, dev) + if !matched { return "", 0, 0, 0, "invalid device identifier: must be in format \"X.Y\" (e.g. \"2.0\")" } speed = 1000000 // default 1 MHz + if s, ok := args["speed"].(float64); ok { if s < 1 || s > 125000000 { return "", 0, 0, 0, "speed must be between 1 Hz and 125 MHz" } + speed = uint32(s) } mode = 0 + if m, ok := args["mode"].(float64); ok { if int(m) < 0 || int(m) > 3 { return "", 0, 0, 0, "mode must be 0-3" } + mode = uint8(m) } bits = 8 + if b, ok := args["bits"].(float64); ok { if int(b) < 1 || int(b) > 32 { return "", 0, 0, 0, "bits must be between 1 and 32" } + bits = uint8(b) } diff --git a/pkg/tools/spi_linux.go b/pkg/tools/spi_linux.go index 9def73662..b023df0b3 100644 --- a/pkg/tools/spi_linux.go +++ b/pkg/tools/spi_linux.go @@ -9,190 +9,321 @@ import ( ) // SPI ioctl constants from Linux kernel headers. + // Calculated from _IOW('k', nr, size) macro: + // + // direction(1)<<30 | size<<16 | type(0x6B)<<8 | nr + const ( - spiIocWrMode = 0x40016B01 // _IOW('k', 1, __u8) + spiIocWrMode = 0x40016B01 // _IOW('k', 1, __u8) + spiIocWrBitsPerWord = 0x40016B03 // _IOW('k', 3, __u8) - spiIocWrMaxSpeedHz = 0x40046B04 // _IOW('k', 4, __u32) - spiIocMessage1 = 0x40206B00 // _IOW('k', 0, struct spi_ioc_transfer) — 32 bytes + + spiIocWrMaxSpeedHz = 0x40046B04 // _IOW('k', 4, __u32) + + spiIocMessage1 = 0x40206B00 // _IOW('k', 0, struct spi_ioc_transfer) — 32 bytes + ) // spiTransfer matches Linux kernel struct spi_ioc_transfer (32 bytes on all architectures). + type spiTransfer struct { - txBuf uint64 - rxBuf uint64 - length uint32 - speedHz uint32 - delayUsecs uint16 + txBuf uint64 + + rxBuf uint64 + + length uint32 + + speedHz uint32 + + delayUsecs uint16 + bitsPerWord uint8 - csChange uint8 - txNbits uint8 - rxNbits uint8 - wordDelay uint8 - pad uint8 + + csChange uint8 + + txNbits uint8 + + rxNbits uint8 + + wordDelay uint8 + + pad uint8 } // configureSPI opens an SPI device and sets mode, bits per word, and speed + func configureSPI(devPath string, mode uint8, bits uint8, speed uint32) (int, *ToolResult) { + fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) + if err != nil { + return -1, ErrorResult(fmt.Sprintf("failed to open %s: %v (check permissions and spidev module)", devPath, err)) + } // Set SPI mode + _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrMode, uintptr(unsafe.Pointer(&mode))) + if errno != 0 { + syscall.Close(fd) + return -1, ErrorResult(fmt.Sprintf("failed to set SPI mode %d: %v", mode, errno)) + } // Set bits per word + _, _, errno = syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrBitsPerWord, uintptr(unsafe.Pointer(&bits))) + if errno != 0 { + syscall.Close(fd) + return -1, ErrorResult(fmt.Sprintf("failed to set bits per word %d: %v", bits, errno)) + } // Set max speed + _, _, errno = syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrMaxSpeedHz, uintptr(unsafe.Pointer(&speed))) + if errno != 0 { + syscall.Close(fd) + return -1, ErrorResult(fmt.Sprintf("failed to set SPI speed %d Hz: %v", speed, errno)) + } return fd, nil + } // transfer performs a full-duplex SPI transfer + func (t *SPITool) transfer(args map[string]any) *ToolResult { + confirm, _ := args["confirm"].(bool) + if !confirm { + return ErrorResult( + "transfer operations require confirm: true. Please confirm with the user before sending data to SPI devices.", ) + } dev, speed, mode, bits, errMsg := parseSPIArgs(args) + if errMsg != "" { + return ErrorResult(errMsg) + } dataRaw, ok := args["data"].([]any) + if !ok || len(dataRaw) == 0 { + return ErrorResult("data is required for transfer (array of byte values 0-255)") + } + if len(dataRaw) > 4096 { + return ErrorResult("data too long: maximum 4096 bytes per SPI transfer") + } txBuf := make([]byte, len(dataRaw)) + for i, v := range dataRaw { + f, ok := v.(float64) + if !ok { + return ErrorResult(fmt.Sprintf("data[%d] is not a valid byte value", i)) + } + b := int(f) + if b < 0 || b > 255 { + return ErrorResult(fmt.Sprintf("data[%d] = %d is out of byte range (0-255)", i, b)) + } + txBuf[i] = byte(b) + } devPath := fmt.Sprintf("/dev/spidev%s", dev) + fd, errResult := configureSPI(devPath, mode, bits, speed) + if errResult != nil { + return errResult + } + defer syscall.Close(fd) rxBuf := make([]byte, len(txBuf)) xfer := spiTransfer{ - txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))), - rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))), - length: uint32(len(txBuf)), - speedHz: speed, + + txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))), + + rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))), + + length: uint32(len(txBuf)), + + speedHz: speed, + bitsPerWord: bits, } _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocMessage1, uintptr(unsafe.Pointer(&xfer))) + runtime.KeepAlive(txBuf) + runtime.KeepAlive(rxBuf) + if errno != 0 { + return ErrorResult(fmt.Sprintf("SPI transfer failed: %v", errno)) + } // Format received bytes + hexBytes := make([]string, len(rxBuf)) + intBytes := make([]int, len(rxBuf)) + for i, b := range rxBuf { + hexBytes[i] = fmt.Sprintf("0x%02x", b) + intBytes[i] = int(b) + } result, _ := json.MarshalIndent(map[string]any{ - "device": devPath, - "sent": len(txBuf), + + "device": devPath, + + "sent": len(txBuf), + "received": intBytes, - "hex": hexBytes, + + "hex": hexBytes, }, "", " ") + return SilentResult(string(result)) + } // readDevice reads bytes from SPI by sending zeros (read-only, no confirm needed) + func (t *SPITool) readDevice(args map[string]any) *ToolResult { + dev, speed, mode, bits, errMsg := parseSPIArgs(args) + if errMsg != "" { + return ErrorResult(errMsg) + } length := 0 + if l, ok := args["length"].(float64); ok { + length = int(l) + } + if length < 1 || length > 4096 { + return ErrorResult("length is required for read (1-4096)") + } devPath := fmt.Sprintf("/dev/spidev%s", dev) + fd, errResult := configureSPI(devPath, mode, bits, speed) + if errResult != nil { + return errResult + } + defer syscall.Close(fd) txBuf := make([]byte, length) // zeros + rxBuf := make([]byte, length) xfer := spiTransfer{ - txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))), - rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))), - length: uint32(length), - speedHz: speed, + + txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))), + + rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))), + + length: uint32(length), + + speedHz: speed, + bitsPerWord: bits, } _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocMessage1, uintptr(unsafe.Pointer(&xfer))) + runtime.KeepAlive(txBuf) + runtime.KeepAlive(rxBuf) + if errno != 0 { + return ErrorResult(fmt.Sprintf("SPI read failed: %v", errno)) + } hexBytes := make([]string, len(rxBuf)) + intBytes := make([]int, len(rxBuf)) + for i, b := range rxBuf { + hexBytes[i] = fmt.Sprintf("0x%02x", b) + intBytes[i] = int(b) + } result, _ := json.MarshalIndent(map[string]any{ + "device": devPath, - "bytes": intBytes, - "hex": hexBytes, + + "bytes": intBytes, + + "hex": hexBytes, + "length": len(rxBuf), }, "", " ") + return SilentResult(string(result)) + } diff --git a/pkg/tools/subagent_reporter_test.go b/pkg/tools/subagent_reporter_test.go index 0dcc2d6ac..923e83a9f 100644 --- a/pkg/tools/subagent_reporter_test.go +++ b/pkg/tools/subagent_reporter_test.go @@ -11,8 +11,11 @@ import ( ) // blockingProvider blocks inside Chat until the context is canceled. + // The ready channel is closed the moment Chat is entered, so callers can + // synchronize before canceling the context. + type blockingProvider struct { ready chan struct{} } @@ -23,42 +26,63 @@ func newBlockingProvider() *blockingProvider { func (p *blockingProvider) Chat( ctx context.Context, + _ []providers.Message, + _ []providers.ToolDefinition, + _ string, + _ map[string]any, ) (*providers.LLMResponse, error) { close(p.ready) // signal: we are now blocking + <-ctx.Done() + return nil, ctx.Err() } func (p *blockingProvider) GetDefaultModel() string { return "test" } // TestSubagentManager_Spawn_EmitsLifecycleEvents verifies that Spawn() fires + // the correct sequence of orchestration events through a real Broadcaster: + // + // agent_spawn → conversation(conductor→sub) → agent_state(waiting) → + // conversation(sub→conductor) → agent_gc(completed) + // + // It also verifies that the snapshot is empty after ReportGC and that the + // completion callback is invoked. + func TestSubagentManager_Spawn_EmitsLifecycleEvents(t *testing.T) { b := orch.NewBroadcaster() + sub := b.Subscribe() + defer b.Unsubscribe(sub) provider := &MockLLMProvider{} + mgr := NewSubagentManager(provider, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) var callbackCalled int32 + cb := AsyncCallback(func(_ context.Context, _ *ToolResult) { atomic.StoreInt32(&callbackCalled, 1) }) _, err := mgr.Spawn( + context.Background(), + "say hello", "hello-task", "", "cli", "direct", "", + cb, ) if err != nil { @@ -66,94 +90,131 @@ func TestSubagentManager_Spawn_EmitsLifecycleEvents(t *testing.T) { } // Collect events until agent_gc or timeout. + var events []orch.Event + deadline := time.After(3 * time.Second) + loop: + for { select { case ev := <-sub.Ch: + events = append(events, ev) + if ev.Type == "agent_gc" { break loop } + case <-deadline: + t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) } } // 1. First event must be agent_spawn with the correct label. + if len(events) == 0 || events[0].Type != "agent_spawn" { t.Fatalf("first event must be agent_spawn, got: %+v", events) } + if events[0].Label != "hello-task" { t.Errorf("agent_spawn label = %q, want %q", events[0].Label, "hello-task") } + spawnedID := events[0].ID // 2. There must be a conversation from conductor → subagent. + var hasConvToSub bool + for _, ev := range events { if ev.Type == "conversation" && ev.From == "conductor" && ev.To == spawnedID { hasConvToSub = true + break } } + if !hasConvToSub { t.Errorf("missing conversation(conductor → %s); events: %+v", spawnedID, events) } // 3. There must be at least one agent_state(waiting) for the subagent. + var hasWaiting bool + for _, ev := range events { if ev.Type == "agent_state" && ev.ID == spawnedID && ev.State == "waiting" { hasWaiting = true + break } } + if !hasWaiting { t.Errorf("missing agent_state(waiting) for %s; events: %+v", spawnedID, events) } // 4. Last event must be agent_gc with reason "completed". + last := events[len(events)-1] + if last.Type != "agent_gc" || last.ID != spawnedID || last.Reason != "completed" { t.Errorf("last event must be agent_gc(completed), got: %+v", last) } // 5. Snapshot must be empty after GC (agent removed from live map). + if snap := b.Snapshot(); len(snap) != 0 { t.Errorf("snapshot must be empty after agent_gc, got: %v", snap) } // 6. Callback must be called. The callback fires in the same goroutine + // as ReportGC (after the deferred unlock), so we poll briefly. + for i := 0; i < 100; i++ { if atomic.LoadInt32(&callbackCalled) == 1 { break } + time.Sleep(10 * time.Millisecond) } + if atomic.LoadInt32(&callbackCalled) != 1 { t.Error("completion callback was not called after agent_gc") } } // TestSubagentManager_Spawn_SnapshotLiveDuringExecution verifies that the + // Broadcaster snapshot contains the agent between agent_spawn and agent_gc. + // Because Publish() updates the agent map before dispatching to subscribers, + // the snapshot is guaranteed to be non-empty as soon as agent_spawn is + // received on the channel. + func TestSubagentManager_Spawn_SnapshotLiveDuringExecution(t *testing.T) { b := orch.NewBroadcaster() + sub := b.Subscribe() + defer b.Unsubscribe(sub) provider := &MockLLMProvider{} + mgr := NewSubagentManager(provider, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) _, err := mgr.Spawn( + context.Background(), + "any task", "live-test", "", "cli", "direct", "", + nil, ) if err != nil { @@ -161,39 +222,59 @@ func TestSubagentManager_Spawn_SnapshotLiveDuringExecution(t *testing.T) { } // Wait for agent_spawn, then immediately check snapshot. + deadline := time.After(2 * time.Second) + for { select { case ev := <-sub.Ch: + if ev.Type == "agent_spawn" { snap := b.Snapshot() + if len(snap) == 0 { t.Error("snapshot must contain the spawned agent after agent_spawn event") } + return // test complete; background goroutine drains safely } + case <-deadline: + t.Fatal("timed out waiting for agent_spawn event") } } } // TestSubagentManager_Spawn_CancelledDuringExecution verifies that when the + // context is canceled while a subagent's LLM call is in progress, the + // Broadcaster receives agent_gc with reason="canceled" and the agent is + // removed from the snapshot. + // + // Synchronization: + // 1. blockingProvider.ready is closed when Chat() is entered (goroutine is + // now blocked inside the LLM call). + // 2. Only then is the context canceled, so there is no race between spawn + // and cancellation. + func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) { b := orch.NewBroadcaster() + sub := b.Subscribe() + defer b.Unsubscribe(sub) bp := newBlockingProvider() + mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) _, err := mgr.Spawn(context.Background(), "long task", "cancel-me", "", "cli", "direct", "", nil) @@ -202,44 +283,61 @@ func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) { } // Wait until the subagent goroutine is inside Chat (blocking on ctx). + select { case <-bp.ready: + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for blockingProvider to enter Chat") } // Cancel via CancelTask — the spawned goroutine's detached context is canceled. + mgr.CancelTask("subagent-1") // Collect events until agent_gc. + var events []orch.Event + deadline := time.After(3 * time.Second) + loop: + for { select { case ev := <-sub.Ch: + events = append(events, ev) + if ev.Type == "agent_gc" { break loop } + case <-deadline: + t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) } } // Locate agent_gc and verify reason = "canceled". + var gcEv orch.Event + for _, ev := range events { if ev.Type == "agent_gc" { gcEv = ev + break } } + if gcEv.Reason != "canceled" { t.Errorf("agent_gc reason = %q, want %q; events: %+v", gcEv.Reason, "canceled", events) } // Snapshot must be empty after the GC event. + if snap := b.Snapshot(); len(snap) != 0 { t.Errorf("snapshot must be empty after agent_gc(canceled), got: %v", snap) } diff --git a/pkg/tools/subagent_tool_test.go b/pkg/tools/subagent_tool_test.go index 481ea394e..f9e1f988a 100644 --- a/pkg/tools/subagent_tool_test.go +++ b/pkg/tools/subagent_tool_test.go @@ -12,19 +12,26 @@ import ( ) // MockLLMProvider is a test implementation of LLMProvider + type MockLLMProvider struct { lastOptions map[string]any } func (m *MockLLMProvider) Chat( ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, ) (*providers.LLMResponse, error) { m.lastOptions = options + // Find the last user message to generate a response + for i := len(messages) - 1; i >= 0; i-- { if messages[i].Role == "user" { return &providers.LLMResponse{ @@ -32,6 +39,7 @@ func (m *MockLLMProvider) Chat( }, nil } } + return &providers.LLMResponse{Content: "No task provided"}, nil } @@ -49,13 +57,19 @@ func (m *MockLLMProvider) GetContextWindow() int { func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) + manager.SetLLMOptions(2048, 0.6) + tool := NewSubagentTool(manager) + tool.SetContext("cli", "direct") ctx := context.Background() + args := map[string]any{"task": "Do something"} + result := tool.Execute(ctx, args) if result == nil || result.IsError { @@ -65,18 +79,23 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { if provider.lastOptions == nil { t.Fatal("Expected LLM options to be passed, got nil") } + if provider.lastOptions["max_tokens"] != 2048 { t.Fatalf("max_tokens = %v, want %d", provider.lastOptions["max_tokens"], 2048) } + if provider.lastOptions["temperature"] != 0.6 { t.Fatalf("temperature = %v, want %v", provider.lastOptions["temperature"], 0.6) } } // TestSubagentTool_Name verifies tool name + func TestSubagentTool_Name(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) if tool.Name() != "subagent" { @@ -85,150 +104,198 @@ func TestSubagentTool_Name(t *testing.T) { } // TestSubagentTool_Description verifies tool description + func TestSubagentTool_Description(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) desc := tool.Description() + if desc == "" { t.Error("Description should not be empty") } + if !strings.Contains(desc, "BLOCK") { t.Errorf("Description should mention 'BLOCK', got: %s", desc) } + if !strings.Contains(desc, "spawn") { t.Errorf("Description should contrast with spawn, got: %s", desc) } } // TestSubagentTool_Parameters verifies tool parameters schema + func TestSubagentTool_Parameters(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) params := tool.Parameters() + if params == nil { t.Error("Parameters should not be nil") } // Check type + if params["type"] != "object" { t.Errorf("Expected type 'object', got: %v", params["type"]) } // Check properties + props, ok := params["properties"].(map[string]any) + if !ok { t.Fatal("Properties should be a map") } // Verify task parameter + task, ok := props["task"].(map[string]any) + if !ok { t.Fatal("Task parameter should exist") } + if task["type"] != "string" { t.Errorf("Task type should be 'string', got: %v", task["type"]) } // Verify label parameter + label, ok := props["label"].(map[string]any) + if !ok { t.Fatal("Label parameter should exist") } + if label["type"] != "string" { t.Errorf("Label type should be 'string', got: %v", label["type"]) } // Check required fields + required, ok := params["required"].([]string) + if !ok { t.Fatal("Required should be a string array") } + if len(required) != 1 || required[0] != "task" { t.Errorf("Required should be ['task'], got: %v", required) } } // TestSubagentTool_SetContext verifies context setting + func TestSubagentTool_SetContext(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) tool.SetContext("test-channel", "test-chat") // Verify context is set (we can't directly access private fields, + // but we can verify it doesn't crash) + // The actual context usage is tested in Execute tests } // TestSubagentTool_Execute_Success tests successful execution + func TestSubagentTool_Execute_Success(t *testing.T) { provider := &MockLLMProvider{} + msgBus := bus.NewMessageBus() + manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) + tool.SetContext("telegram", "chat-123") ctx := context.Background() + args := map[string]any{ - "task": "Write a haiku about coding", + "task": "Write a haiku about coding", + "label": "haiku-task", } result := tool.Execute(ctx, args) // Verify basic ToolResult structure + if result == nil { t.Fatal("Result should not be nil") } // Verify no error + if result.IsError { t.Errorf("Expected success, got error: %s", result.ForLLM) } // Verify not async + if result.Async { t.Error("SubagentTool should be synchronous, not async") } // Verify not silent + if result.Silent { t.Error("SubagentTool should not be silent") } // Verify ForUser contains brief summary (not empty) + if result.ForUser == "" { t.Error("ForUser should contain result summary") } + if !strings.Contains(result.ForUser, "Task completed") { t.Errorf("ForUser should contain task completion, got: %s", result.ForUser) } // Verify ForLLM contains full details + if result.ForLLM == "" { t.Error("ForLLM should contain full details") } + if !strings.Contains(result.ForLLM, "haiku-task") { t.Errorf("ForLLM should contain label 'haiku-task', got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "Task completed:") { t.Errorf("ForLLM should contain task result, got: %s", result.ForLLM) } } // TestSubagentTool_Execute_NoLabel tests execution without label + func TestSubagentTool_Execute_NoLabel(t *testing.T) { provider := &MockLLMProvider{} + msgBus := bus.NewMessageBus() + manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) ctx := context.Background() + args := map[string]any{ "task": "Test task without label", } @@ -240,18 +307,23 @@ func TestSubagentTool_Execute_NoLabel(t *testing.T) { } // ForLLM should show (unnamed) for missing label + if !strings.Contains(result.ForLLM, "(unnamed)") { t.Errorf("ForLLM should show '(unnamed)' for missing label, got: %s", result.ForLLM) } } // TestSubagentTool_Execute_MissingTask tests error handling for missing task + func TestSubagentTool_Execute_MissingTask(t *testing.T) { provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) ctx := context.Background() + args := map[string]any{ "label": "test", } @@ -259,29 +331,35 @@ func TestSubagentTool_Execute_MissingTask(t *testing.T) { result := tool.Execute(ctx, args) // Should return error + if !result.IsError { t.Error("Expected error for missing task parameter") } // ForLLM should contain helpful error with example + if !strings.Contains(result.ForLLM, `"task"`) { t.Errorf("Error message should mention '\"task\"', got: %s", result.ForLLM) } + if !strings.Contains(result.ForLLM, "Example") { t.Errorf("Error message should include usage example, got: %s", result.ForLLM) } // Err should be set + if result.Err == nil { t.Error("Err should be set for validation failure") } } // TestSubagentTool_Execute_NilManager tests error handling for nil manager + func TestSubagentTool_Execute_NilManager(t *testing.T) { tool := NewSubagentTool(nil) ctx := context.Background() + args := map[string]any{ "task": "test task", } @@ -289,6 +367,7 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) { result := tool.Execute(ctx, args) // Should return error + if !result.IsError { t.Error("Expected error for nil manager") } @@ -299,18 +378,26 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) { } // TestSubagentTool_Execute_ContextPassing verifies context is properly used + func TestSubagentTool_Execute_ContextPassing(t *testing.T) { provider := &MockLLMProvider{} + msgBus := bus.NewMessageBus() + manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) // Set context + channel := "test-channel" + chatID := "test-chat" + tool.SetContext(channel, chatID) ctx := context.Background() + args := map[string]any{ "task": "Test context passing", } @@ -318,40 +405,53 @@ func TestSubagentTool_Execute_ContextPassing(t *testing.T) { result := tool.Execute(ctx, args) // Should succeed + if result.IsError { t.Errorf("Expected success with context, got error: %s", result.ForLLM) } // The context is used internally; we can't directly test it + // but execution success indicates context was handled properly } // TestSubagentTool_ForUserTruncation verifies long content is truncated for user + func TestSubagentTool_ForUserTruncation(t *testing.T) { // Create a mock provider that returns very long content + provider := &MockLLMProvider{} + msgBus := bus.NewMessageBus() + manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) + tool := NewSubagentTool(manager) ctx := context.Background() // Create a task that will generate long response + longTask := strings.Repeat("This is a very long task description. ", 100) + args := map[string]any{ - "task": longTask, + "task": longTask, + "label": "long-test", } result := tool.Execute(ctx, args) // ForUser should be truncated to 500 chars + "..." + maxUserLen := 500 + if len(result.ForUser) > maxUserLen+3 { // +3 for "..." t.Errorf("ForUser should be truncated to ~%d chars, got: %d", maxUserLen, len(result.ForUser)) } // ForLLM should have full content + if !strings.Contains(result.ForLLM, longTask[:50]) { t.Error("ForLLM should contain reference to original task") } @@ -359,21 +459,29 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) { func TestFormatToolStats(t *testing.T) { tests := []struct { - name string + name string + stats map[string]int - want string + + want string }{ {"empty", map[string]int{}, ""}, + {"single", map[string]int{"exec": 3}, "exec:3"}, + { "multiple sorted", + map[string]int{"read_file": 5, "exec": 3, "write_file": 1}, + "exec:3,read_file:5,write_file:1", }, } + for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := formatToolStats(tt.stats) + if got != tt.want { t.Errorf("formatToolStats(%v) = %q, want %q", tt.stats, got, tt.want) } @@ -382,15 +490,22 @@ func TestFormatToolStats(t *testing.T) { } // TestSubagentManager_Spawn_SetsMetadata verifies that the bus message from a + // completed spawn includes execution statistics in Metadata. + func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) { provider := &MockLLMProvider{} + msgBus := bus.NewMessageBus() + mgr := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) _, err := mgr.Spawn( + context.Background(), + "say hello", "meta-test", "", "cli", "direct", "", + nil, ) if err != nil { @@ -398,9 +513,13 @@ func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) { } // Consume the inbound message from the bus + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + received, ok := msgBus.ConsumeInbound(ctx) + if !ok { t.Fatal("timed out waiting for bus message") } @@ -408,16 +527,21 @@ func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) { if received.Channel != "system" { t.Fatalf("expected channel 'system', got %q", received.Channel) } + if received.Metadata == nil { t.Fatal("Metadata should not be nil") } + if received.Metadata["iterations"] != "1" { t.Errorf("iterations = %q, want %q", received.Metadata["iterations"], "1") } + if received.Metadata["tool_calls"] != "0" { t.Errorf("tool_calls = %q, want %q", received.Metadata["tool_calls"], "0") } + // duration_ms should be a non-negative number + if received.Metadata["duration_ms"] == "" { t.Error("duration_ms should be present") } diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go index 793c51a3c..f463ceaf7 100644 --- a/pkg/tools/toolloop.go +++ b/pkg/tools/toolloop.go @@ -1,7 +1,11 @@ // PicoClaw - Ultra-lightweight personal AI agent + // Inspired by and based on nanobot: https://github.com/HKUDS/nanobot + // License: MIT + // + // Copyright (c) 2026 PicoClaw contributors package tools @@ -18,140 +22,207 @@ import ( ) // ToolLoopConfig configures the tool execution loop. + type ToolLoopConfig struct { - Provider providers.LLMProvider - Model string - Tools *ToolRegistry + Provider providers.LLMProvider + + Model string + + Tools *ToolRegistry + MaxIterations int - LLMOptions map[string]any + + LLMOptions map[string]any + // Reporter and AgentID replace the old OnStateChange func. + // Reporter is called with ReportStateChange("waiting","") before each LLM + // call and ReportStateChange("toolcall", toolName) when each tool starts. + // Pass nil or orch.Noop to disable. nil is treated as orch.Noop internally. + Reporter orch.AgentReporter - AgentID string + + AgentID string } // ToolLoopResult contains the result of running the tool loop. + type ToolLoopResult struct { - Content string + Content string + Iterations int - ToolCalls int // total tool call count across all iterations - ToolStats map[string]int // tool name → call count + + ToolCalls int // total tool call count across all iterations + + ToolStats map[string]int // tool name → call count } // RunToolLoop executes the LLM + tool call iteration loop. + // This is the core agent logic that can be reused by both main agent and subagents. + func RunToolLoop( ctx context.Context, + config ToolLoopConfig, + messages []providers.Message, + channel, chatID string, ) (*ToolLoopResult, error) { reporter := config.Reporter + if reporter == nil { reporter = orch.Noop } iteration := 0 + totalToolCalls := 0 + toolStats := map[string]int{} + var finalContent string for iteration < config.MaxIterations { iteration++ logger.DebugCF("toolloop", "LLM iteration", + map[string]any{ "iteration": iteration, - "max": config.MaxIterations, + + "max": config.MaxIterations, }) // 1. Build tool definitions + var providerToolDefs []providers.ToolDefinition + if config.Tools != nil { providerToolDefs = config.Tools.ToProviderDefs() } // 2. Set default LLM options + llmOpts := config.LLMOptions + if llmOpts == nil { llmOpts = map[string]any{} } + // 3. Call LLM (hook: waiting for response) + reporter.ReportStateChange(config.AgentID, orch.AgentStateWaiting, "") + response, err := config.Provider.Chat(ctx, messages, providerToolDefs, config.Model, llmOpts) if err != nil { logger.ErrorCF("toolloop", "LLM call failed", + map[string]any{ "iteration": iteration, - "error": err.Error(), + + "error": err.Error(), }) + return nil, fmt.Errorf("LLM call failed: %w", err) } // 4. If no tool calls, we're done + if len(response.ToolCalls) == 0 { finalContent = response.Content + logger.InfoCF("toolloop", "LLM response without tool calls (direct answer)", + map[string]any{ - "iteration": iteration, + "iteration": iteration, + "content_chars": len(finalContent), }) + break } normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls)) + for _, tc := range response.ToolCalls { normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc)) } // 5. Log tool calls + toolNames := make([]string, 0, len(normalizedToolCalls)) + for _, tc := range normalizedToolCalls { toolNames = append(toolNames, tc.Name) } + logger.InfoCF("toolloop", "LLM requested tool calls", + map[string]any{ - "tools": toolNames, - "count": len(normalizedToolCalls), + "tools": toolNames, + + "count": len(normalizedToolCalls), + "iteration": iteration, }) // 6. Build assistant message with tool calls + assistantMsg := providers.Message{ - Role: "assistant", + Role: "assistant", + Content: response.Content, } + for _, tc := range normalizedToolCalls { assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, providers.ToolCall{ - ID: tc.ID, - Type: "function", - Name: tc.Name, + ID: tc.ID, + + Type: "function", + + Name: tc.Name, + Arguments: tc.Arguments, + Function: &providers.FunctionCall{ - Name: tc.Name, + Name: tc.Name, + Arguments: tc.Arguments, }, }) } + messages = append(messages, assistantMsg) // 7. Execute tool calls (hook: toolcall per tool) + for _, tc := range normalizedToolCalls { argsJSON, _ := json.Marshal(tc.Arguments) + argsPreview := utils.Truncate(string(argsJSON), 200) + logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), + map[string]any{ - "tool": tc.Name, + "tool": tc.Name, + "iteration": iteration, }) + reporter.ReportStateChange(config.AgentID, orch.AgentStateToolCall, tc.Name) + totalToolCalls++ + toolStats[tc.Name]++ // 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) } else { @@ -159,25 +230,34 @@ func RunToolLoop( } // Determine content for LLM + contentForLLM := toolResult.ForLLM + if contentForLLM == "" && toolResult.Err != nil { contentForLLM = toolResult.Err.Error() } // Add tool result message + toolResultMsg := providers.Message{ - Role: "tool", - Content: contentForLLM, + Role: "tool", + + Content: contentForLLM, + ToolCallID: tc.ID, } + messages = append(messages, toolResultMsg) } } return &ToolLoopResult{ - Content: finalContent, + Content: finalContent, + Iterations: iteration, - ToolCalls: totalToolCalls, - ToolStats: toolStats, + + ToolCalls: totalToolCalls, + + ToolStats: toolStats, }, nil } diff --git a/pkg/tools/toolloop_reporter_test.go b/pkg/tools/toolloop_reporter_test.go index c4a70bccb..4755b94bd 100644 --- a/pkg/tools/toolloop_reporter_test.go +++ b/pkg/tools/toolloop_reporter_test.go @@ -10,53 +10,78 @@ import ( ) // reporterSpy records every ReportStateChange call in order. + // Spawn/Conversation/GC are not needed for toolloop tests. + type reporterSpy struct { - mu sync.Mutex + mu sync.Mutex + calls []spyCall } type spyCall struct { state orch.AgentState - tool string + + tool string } -func (r *reporterSpy) ReportSpawn(id, label, task string) {} +func (r *reporterSpy) ReportSpawn(id, label, task string) {} + func (r *reporterSpy) ReportConversation(from, to, text string) {} -func (r *reporterSpy) ReportGC(id, reason string) {} + +func (r *reporterSpy) ReportGC(id, reason string) {} + func (r *reporterSpy) ReportStateChange(id string, state orch.AgentState, tool string) { r.mu.Lock() + r.calls = append(r.calls, spyCall{state, tool}) + r.mu.Unlock() } func (r *reporterSpy) snapshot() []spyCall { r.mu.Lock() + defer r.mu.Unlock() + out := make([]spyCall, len(r.calls)) + copy(out, r.calls) + return out } // sequenceMockProvider returns a tool call on the first Chat() call and a + // plain text response on all subsequent calls. Used to exercise the + // waiting → toolcall → waiting event sequence in RunToolLoop. + type sequenceMockProvider struct { - mu sync.Mutex + mu sync.Mutex + callCount int } func (m *sequenceMockProvider) Chat( _ context.Context, + _ []providers.Message, + _ []providers.ToolDefinition, + _ string, + _ map[string]any, ) (*providers.LLMResponse, error) { m.mu.Lock() + m.callCount++ + n := m.callCount + m.mu.Unlock() + if n == 1 { return &providers.LLMResponse{ ToolCalls: []providers.ToolCall{ @@ -64,17 +89,24 @@ func (m *sequenceMockProvider) Chat( }, }, nil } + return &providers.LLMResponse{Content: "done"}, nil } + func (m *sequenceMockProvider) GetDefaultModel() string { return "test" } -func (m *sequenceMockProvider) SupportsTools() bool { return true } -func (m *sequenceMockProvider) GetContextWindow() int { return 4096 } + +func (m *sequenceMockProvider) SupportsTools() bool { return true } + +func (m *sequenceMockProvider) GetContextWindow() int { return 4096 } // echoTool is a minimal Tool stub registered as "echo_tool". + type echoTool struct{} -func (t *echoTool) Name() string { return "echo_tool" } +func (t *echoTool) Name() string { return "echo_tool" } + func (t *echoTool) Description() string { return "echo" } + func (t *echoTool) Parameters() map[string]any { return map[string]any{"type": "object", "properties": map[string]any{}} } @@ -84,13 +116,19 @@ func (t *echoTool) Execute(_ context.Context, _ map[string]any) *ToolResult { } // TestToolLoop_NilReporter_FallsBackToNoop ensures that passing nil as + // Reporter does not panic — the loop must substitute orch.Noop internally. + func TestToolLoop_NilReporter_FallsBackToNoop(t *testing.T) { _, err := RunToolLoop(context.Background(), ToolLoopConfig{ - Provider: &MockLLMProvider{}, - Model: "test", + Provider: &MockLLMProvider{}, + + Model: "test", + MaxIterations: 1, - Reporter: nil, // must not panic + + Reporter: nil, // must not panic + }, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct") if err != nil { t.Fatalf("unexpected error with nil reporter: %v", err) @@ -98,120 +136,171 @@ func TestToolLoop_NilReporter_FallsBackToNoop(t *testing.T) { } // TestToolLoop_Reporter_WaitingBeforeLLM verifies that ReportStateChange is + // called with state="waiting" before the first LLM call. The mock provider + // returns a direct text answer (no tool calls), so exactly one waiting event + // is expected. + func TestToolLoop_Reporter_WaitingBeforeLLM(t *testing.T) { rep := &reporterSpy{} + _, err := RunToolLoop(context.Background(), ToolLoopConfig{ - Provider: &MockLLMProvider{}, - Model: "test", + Provider: &MockLLMProvider{}, + + Model: "test", + MaxIterations: 1, - Reporter: rep, - AgentID: "sess-1", + + Reporter: rep, + + AgentID: "sess-1", }, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct") if err != nil { t.Fatalf("unexpected error: %v", err) } + calls := rep.snapshot() + if len(calls) == 0 { t.Fatal("expected at least one ReportStateChange call") } + if calls[0].state != orch.AgentStateWaiting { t.Fatalf("first call must be state=waiting, got %+v", calls[0]) } } // TestToolLoop_Reporter_ToolcallOrderedAfterWaiting verifies the canonical + // two-iteration sequence: + // + // waiting (before 1st LLM call) + // toolcall(echo_tool) (before tool execution) + // waiting (before 2nd LLM call) + // + // The sequenceMockProvider returns a tool call on iteration 1 and a text + // response on iteration 2, driving exactly this path. + func TestToolLoop_Reporter_ToolcallOrderedAfterWaiting(t *testing.T) { rep := &reporterSpy{} + reg := NewToolRegistry() + reg.Register(&echoTool{}) _, err := RunToolLoop(context.Background(), ToolLoopConfig{ - Provider: &sequenceMockProvider{}, - Model: "test", - Tools: reg, + Provider: &sequenceMockProvider{}, + + Model: "test", + + Tools: reg, + MaxIterations: 5, - Reporter: rep, - AgentID: "sess-1", + + Reporter: rep, + + AgentID: "sess-1", }, []providers.Message{{Role: "user", Content: "do it"}}, "cli", "direct") if err != nil { t.Fatalf("unexpected error: %v", err) } calls := rep.snapshot() + if len(calls) < 3 { t.Fatalf("expected at least 3 calls, got %d: %+v", len(calls), calls) } + if calls[0].state != orch.AgentStateWaiting { t.Fatalf("calls[0] must be waiting, got %+v", calls[0]) } + if calls[1].state != orch.AgentStateToolCall || calls[1].tool != "echo_tool" { t.Fatalf("calls[1] must be toolcall(echo_tool), got %+v", calls[1]) } + if calls[2].state != orch.AgentStateWaiting { t.Fatalf("calls[2] must be waiting (2nd LLM iteration), got %+v", calls[2]) } } // TestToolLoop_ToolCallStats verifies that ToolLoopResult.ToolCalls and + // ToolStats are populated correctly after a tool call iteration. + func TestToolLoop_ToolCallStats(t *testing.T) { reg := NewToolRegistry() + reg.Register(&echoTool{}) result, err := RunToolLoop(context.Background(), ToolLoopConfig{ - Provider: &sequenceMockProvider{}, - Model: "test", - Tools: reg, + Provider: &sequenceMockProvider{}, + + Model: "test", + + Tools: reg, + MaxIterations: 5, }, []providers.Message{{Role: "user", Content: "do it"}}, "cli", "direct") if err != nil { t.Fatalf("unexpected error: %v", err) } + if result.ToolCalls != 1 { t.Errorf("ToolCalls = %d, want 1", result.ToolCalls) } + if result.ToolStats["echo_tool"] != 1 { t.Errorf("ToolStats[echo_tool] = %d, want 1", result.ToolStats["echo_tool"]) } + if result.Iterations != 2 { t.Errorf("Iterations = %d, want 2", result.Iterations) } } // TestToolLoop_NoToolCalls_ZeroStats verifies that a direct answer (no tool + // calls) produces zero ToolCalls and an empty ToolStats map. + func TestToolLoop_NoToolCalls_ZeroStats(t *testing.T) { result, err := RunToolLoop(context.Background(), ToolLoopConfig{ - Provider: &MockLLMProvider{}, - Model: "test", + Provider: &MockLLMProvider{}, + + Model: "test", + MaxIterations: 1, }, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct") if err != nil { t.Fatalf("unexpected error: %v", err) } + if result.ToolCalls != 0 { t.Errorf("ToolCalls = %d, want 0", result.ToolCalls) } + if len(result.ToolStats) != 0 { t.Errorf("ToolStats = %v, want empty", result.ToolStats) } } // TestToolLoop_Reporter_NoopImplementsInterface is a compile-time check that + // orch.Noop satisfies the orch.AgentReporter interface accepted by + // ToolLoopConfig.Reporter. If Noop ever stops implementing the interface the + // build will fail here before any test runs. + func TestToolLoop_Reporter_NoopImplementsInterface(t *testing.T) { var _ orch.AgentReporter = orch.Noop } diff --git a/pkg/tools/web.go b/pkg/tools/web.go index 42ad6d4e6..7081424c2 100644 --- a/pkg/tools/web.go +++ b/pkg/tools/web.go @@ -17,35 +17,51 @@ const ( userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" // HTTP client timeouts for web tool providers. - searchTimeout = 10 * time.Second // Brave, Tavily, DuckDuckGo + + searchTimeout = 10 * time.Second // Brave, Tavily, DuckDuckGo + perplexityTimeout = 30 * time.Second // Perplexity (LLM-based, slower) - fetchTimeout = 60 * time.Second // WebFetchTool + + fetchTimeout = 60 * time.Second // WebFetchTool defaultMaxChars = 50000 - maxRedirects = 5 + + maxRedirects = 5 ) // Pre-compiled regexes for HTML text extraction + var ( - reScript = regexp.MustCompile(`<script[\s\S]*?</script>`) - reStyle = regexp.MustCompile(`<style[\s\S]*?</style>`) - reTags = regexp.MustCompile(`<[^>]+>`) + reScript = regexp.MustCompile(`<script[\s\S]*?</script>`) + + reStyle = regexp.MustCompile(`<style[\s\S]*?</style>`) + + reTags = regexp.MustCompile(`<[^>]+>`) + reWhitespace = regexp.MustCompile(`[^\S\n]+`) + reBlankLines = regexp.MustCompile(`\n{3,}`) // DuckDuckGo result extraction - reDDGLink = regexp.MustCompile(`<a[^>]*class="[^"]*result__a[^"]*"[^>]*href="([^"]+)"[^>]*>([\s\S]*?)</a>`) + + reDDGLink = regexp.MustCompile(`<a[^>]*class="[^"]*result__a[^"]*"[^>]*href="([^"]+)"[^>]*>([\s\S]*?)</a>`) + reDDGSnippet = regexp.MustCompile(`<a class="result__snippet[^"]*".*?>([\s\S]*?)</a>`) ) // createHTTPClient creates an HTTP client with optional proxy support + func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, error) { client := &http.Client{ Timeout: timeout, + Transport: &http.Transport{ - MaxIdleConns: 10, - IdleConnTimeout: 30 * time.Second, - DisableCompression: false, + MaxIdleConns: 10, + + IdleConnTimeout: 30 * time.Second, + + DisableCompression: false, + TLSHandshakeTimeout: 15 * time.Second, }, } @@ -55,18 +71,26 @@ func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, err if err != nil { return nil, fmt.Errorf("invalid proxy URL: %w", err) } + scheme := strings.ToLower(proxy.Scheme) + switch scheme { case "http", "https", "socks5", "socks5h": + default: + return nil, fmt.Errorf( + "unsupported proxy scheme %q (supported: http, https, socks5, socks5h)", + proxy.Scheme, ) } + if proxy.Host == "" { return nil, fmt.Errorf("invalid proxy URL: missing host") } + client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy) } else { client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment @@ -80,8 +104,10 @@ type SearchProvider interface { } type searchResultItem struct { - Title string - URL string + Title string + + URL string + Snippet string } @@ -91,32 +117,41 @@ func formatWebSearchResults(query, provider string, results []searchResultItem, } header := fmt.Sprintf("Results for: %s", query) + if provider != "" { header += " (via " + provider + ")" } var sb strings.Builder + sb.WriteString(header) + for i, item := range results { if i >= count { break } + fmt.Fprintf(&sb, "\n%d. %s\n %s", i+1, item.Title, item.URL) + if item.Snippet != "" { fmt.Fprintf(&sb, "\n %s", item.Snippet) } } + return sb.String() } type BraveSearchProvider struct { apiKey string - proxy string + + proxy string + client *http.Client } func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d", + url.QueryEscape(query), count) req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) @@ -125,12 +160,14 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in } req.Header.Set("Accept", "application/json") + req.Header.Set("X-Subscription-Token", p.apiKey) resp, err := p.client.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } + defer resp.Body.Close() body, err := io.ReadAll(resp.Body) @@ -141,8 +178,10 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in var searchResp struct { Web struct { Results []struct { - Title string `json:"title"` - URL string `json:"url"` + Title string `json:"title"` + + URL string `json:"url"` + Description string `json:"description"` } `json:"results"` } `json:"web"` @@ -150,16 +189,22 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in if err := json.Unmarshal(body, &searchResp); err != nil { // Log error body for debugging + fmt.Printf("Brave API Error Body: %s\n", string(body)) + return "", fmt.Errorf("failed to parse response: %w", err) } results := searchResp.Web.Results + items := make([]searchResultItem, 0, len(results)) + for _, item := range results { items = append(items, searchResultItem{ - Title: item.Title, - URL: item.URL, + Title: item.Title, + + URL: item.URL, + Snippet: item.Description, }) } @@ -168,26 +213,36 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in } type TavilySearchProvider struct { - apiKey string + apiKey string + baseURL string - proxy string - client *http.Client + + proxy string + + client *http.Client } func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { searchURL := p.baseURL + if searchURL == "" { searchURL = "https://api.tavily.com/search" } payload := map[string]any{ - "api_key": p.apiKey, - "query": query, - "search_depth": "advanced", - "include_answer": false, - "include_images": false, + "api_key": p.apiKey, + + "query": query, + + "search_depth": "advanced", + + "include_answer": false, + + "include_images": false, + "include_raw_content": false, - "max_results": count, + + "max_results": count, } bodyBytes, err := json.Marshal(payload) @@ -201,12 +256,14 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i } req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", userAgent) resp, err := p.client.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } + defer resp.Body.Close() body, err := io.ReadAll(resp.Body) @@ -220,8 +277,10 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i var searchResp struct { Results []struct { - Title string `json:"title"` - URL string `json:"url"` + Title string `json:"title"` + + URL string `json:"url"` + Content string `json:"content"` } `json:"results"` } @@ -231,11 +290,15 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i } results := searchResp.Results + items := make([]searchResultItem, 0, len(results)) + for _, item := range results { items = append(items, searchResultItem{ - Title: item.Title, - URL: item.URL, + Title: item.Title, + + URL: item.URL, + Snippet: item.Content, }) } @@ -244,7 +307,8 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i } type DuckDuckGoSearchProvider struct { - proxy string + proxy string + client *http.Client } @@ -262,6 +326,7 @@ func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, cou if err != nil { return "", fmt.Errorf("request failed: %w", err) } + defer resp.Body.Close() body, err := io.ReadAll(resp.Body) @@ -274,11 +339,15 @@ func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, cou func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query string) (string, error) { // Simple regex based extraction for DDG HTML + // Strategy: Find all result containers or key anchors directly // Try finding the result links directly first, as they are the most critical + // Pattern: <a class="result__a" href="...">Title</a> + // The previous regex was a bit strict. Let's make it more flexible for attributes order/content + matches := reDDGLink.FindAllStringSubmatch(html, count+5) if len(matches) == 0 { @@ -288,17 +357,22 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query snippetMatches := reDDGSnippet.FindAllStringSubmatch(html, count+5) maxItems := min(len(matches), count) + items := make([]searchResultItem, 0, maxItems) for i := range maxItems { urlStr := matches[i][1] + title := stripTags(matches[i][2]) + title = strings.TrimSpace(title) // URL decoding if needed + if strings.Contains(urlStr, "uddg=") { if u, err := url.QueryUnescape(urlStr); err == nil { _, after, ok := strings.Cut(u, "uddg=") + if ok { urlStr = after } @@ -306,15 +380,20 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query } snippet := "" + // Attempt to attach snippet if available and index aligns + if i < len(snippetMatches) { snippet = stripTags(snippetMatches[i][1]) + snippet = strings.TrimSpace(snippet) } items = append(items, searchResultItem{ - Title: title, - URL: urlStr, + Title: title, + + URL: urlStr, + Snippet: snippet, }) } @@ -328,7 +407,9 @@ func stripTags(content string) string { type PerplexitySearchProvider struct { apiKey string - proxy string + + proxy string + client *http.Client } @@ -337,16 +418,21 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou payload := map[string]any{ "model": "sonar", + "messages": []map[string]string{ { - "role": "system", + "role": "system", + "content": "You are a search assistant. Provide concise search results with titles, URLs, and brief descriptions in the following format:\n1. Title\n URL\n Description\n\nDo not add extra commentary.", }, + { - "role": "user", + "role": "user", + "content": fmt.Sprintf("Search for: %s. Provide up to %d relevant results.", query, count), }, }, + "max_tokens": 1000, } @@ -361,13 +447,16 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou } req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+p.apiKey) + req.Header.Set("User-Agent", userAgent) resp, err := p.client.Do(req) if err != nil { return "", fmt.Errorf("request failed: %w", err) } + defer resp.Body.Close() body, err := io.ReadAll(resp.Body) @@ -399,45 +488,66 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou } type WebSearchTool struct { - provider SearchProvider + provider SearchProvider + providerName string - maxResults int + + maxResults int } // ProviderName returns the name of the active search provider (e.g. "brave", "perplexity"). + func (t *WebSearchTool) ProviderName() string { return t.providerName } type WebSearchToolOptions struct { - BraveAPIKey string - BraveMaxResults int - BraveEnabled bool - TavilyAPIKey string - TavilyBaseURL string - TavilyMaxResults int - TavilyEnabled bool + BraveAPIKey string + + BraveMaxResults int + + BraveEnabled bool + + TavilyAPIKey string + + TavilyBaseURL string + + TavilyMaxResults int + + TavilyEnabled bool + DuckDuckGoMaxResults int - DuckDuckGoEnabled bool - PerplexityAPIKey string + + DuckDuckGoEnabled bool + + PerplexityAPIKey string + PerplexityMaxResults int - PerplexityEnabled bool - Proxy string + + PerplexityEnabled bool + + Proxy string } func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { var provider SearchProvider + var providerName string + maxResults := 5 // Priority: Perplexity > Brave > Tavily > DuckDuckGo + if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" { client, err := createHTTPClient(opts.Proxy, perplexityTimeout) if err != nil { return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err) } + provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy, client: client} + providerName = "perplexity" + if opts.PerplexityMaxResults > 0 { maxResults = opts.PerplexityMaxResults } @@ -446,8 +556,11 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { if err != nil { return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err) } + provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy, client: client} + providerName = "brave" + if opts.BraveMaxResults > 0 { maxResults = opts.BraveMaxResults } @@ -456,13 +569,19 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { if err != nil { return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err) } + provider = &TavilySearchProvider{ - apiKey: opts.TavilyAPIKey, + apiKey: opts.TavilyAPIKey, + baseURL: opts.TavilyBaseURL, - proxy: opts.Proxy, - client: client, + + proxy: opts.Proxy, + + client: client, } + providerName = "tavily" + if opts.TavilyMaxResults > 0 { maxResults = opts.TavilyMaxResults } @@ -471,8 +590,11 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { if err != nil { return nil, fmt.Errorf("failed to create HTTP client for DuckDuckGo: %w", err) } + provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy, client: client} + providerName = "duckduckgo" + if opts.DuckDuckGoMaxResults > 0 { maxResults = opts.DuckDuckGoMaxResults } @@ -481,9 +603,11 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { } return &WebSearchTool{ - provider: provider, + provider: provider, + providerName: providerName, - maxResults: maxResults, + + maxResults: maxResults, }, nil } @@ -498,29 +622,38 @@ func (t *WebSearchTool) Description() string { func (t *WebSearchTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "query": map[string]any{ - "type": "string", + "type": "string", + "description": "Search query", }, + "count": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Number of results (1-10)", - "minimum": 1.0, - "maximum": 10.0, + + "minimum": 1.0, + + "maximum": 10.0, }, }, + "required": []string{"query"}, } } func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { query, ok := args["query"].(string) + if !ok { return ErrorResult("query is required") } count := t.maxResults + if c, ok := args["count"].(float64); ok { if int(c) > 0 && int(c) <= 10 { count = int(c) @@ -533,20 +666,25 @@ func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolR } return &ToolResult{ - ForLLM: result, + ForLLM: result, + ForUser: result, } } type WebFetchTool struct { maxChars int - proxy string - client *http.Client + + proxy string + + client *http.Client } func NewWebFetchTool(maxChars int) *WebFetchTool { // createHTTPClient cannot fail with an empty proxy string. + tool, _ := NewWebFetchToolWithProxy(maxChars, "") + return tool } @@ -554,20 +692,26 @@ func NewWebFetchToolWithProxy(maxChars int, proxy string) (*WebFetchTool, error) if maxChars <= 0 { maxChars = defaultMaxChars } + client, err := createHTTPClient(proxy, fetchTimeout) if err != nil { return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err) } + client.CheckRedirect = func(req *http.Request, via []*http.Request) error { if len(via) >= maxRedirects { return fmt.Errorf("stopped after %d redirects", maxRedirects) } + return nil } + return &WebFetchTool{ maxChars: maxChars, - proxy: proxy, - client: client, + + proxy: proxy, + + client: client, }, nil } @@ -582,23 +726,30 @@ func (t *WebFetchTool) Description() string { func (t *WebFetchTool) Parameters() map[string]any { return map[string]any{ "type": "object", + "properties": map[string]any{ "url": map[string]any{ - "type": "string", + "type": "string", + "description": "URL to fetch", }, + "maxChars": map[string]any{ - "type": "integer", + "type": "integer", + "description": "Maximum characters to extract", - "minimum": 100.0, + + "minimum": 100.0, }, }, + "required": []string{"url"}, } } func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { urlStr, ok := args["url"].(string) + if !ok { return ErrorResult("url is required") } @@ -617,6 +768,7 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe } maxChars := t.maxChars + if mc, ok := args["maxChars"].(float64); ok { if int(mc) > 100 { maxChars = int(mc) @@ -634,6 +786,7 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe if err != nil { return ErrorResult(fmt.Sprintf("request failed: %v", err)) } + defer resp.Body.Close() body, err := io.ReadAll(resp.Body) @@ -646,71 +799,98 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe var text, extractor string bodyStr := string(body) + if strings.Contains(contentType, "application/json") { var jsonData any + if err := json.Unmarshal(body, &jsonData); err == nil { formatted, _ := json.MarshalIndent(jsonData, "", " ") + text = string(formatted) + extractor = "json" } else { text = bodyStr + extractor = "raw" } } else if strings.Contains(contentType, "text/html") || len(body) > 0 && + (strings.HasPrefix(bodyStr, "<!DOCTYPE") || strings.HasPrefix(strings.ToLower(bodyStr), "<html")) { text = t.extractText(bodyStr) + extractor = "text" } else { text = bodyStr + extractor = "raw" } truncated := len(text) > maxChars + if truncated { text = text[:maxChars] } result := map[string]any{ - "url": urlStr, - "status": resp.StatusCode, + "url": urlStr, + + "status": resp.StatusCode, + "extractor": extractor, + "truncated": truncated, - "length": len(text), - "text": text, + + "length": len(text), + + "text": text, } resultJSON, _ := json.MarshalIndent(result, "", " ") return &ToolResult{ ForLLM: fmt.Sprintf( + "Fetched %d bytes from %s (extractor: %s, truncated: %v)", + len(text), + urlStr, + extractor, + truncated, ), + ForUser: string(resultJSON), } } func (t *WebFetchTool) extractText(htmlContent string) string { result := reScript.ReplaceAllLiteralString(htmlContent, "") + result = reStyle.ReplaceAllLiteralString(result, "") + result = reTags.ReplaceAllLiteralString(result, "") result = strings.TrimSpace(result) result = reWhitespace.ReplaceAllString(result, " ") + result = reBlankLines.ReplaceAllString(result, "\n\n") lines := strings.Split(result, "\n") + var sb strings.Builder + for _, line := range lines { line = strings.TrimSpace(line) + if line != "" { if sb.Len() > 0 { sb.WriteByte('\n') } + sb.WriteString(line) } } diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index db3c08ba6..5758fed11 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -11,16 +11,22 @@ import ( ) // TestWebTool_WebFetch_Success verifies successful URL fetching + func TestWebTool_WebFetch_Success(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + w.Write([]byte("<html><body><h1>Test Page</h1><p>Content here</p></body></html>")) })) + defer server.Close() tool := NewWebFetchTool(50000) + ctx := context.Background() + args := map[string]any{ "url": server.URL, } @@ -28,35 +34,45 @@ func TestWebTool_WebFetch_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain the fetched content + if !strings.Contains(result.ForUser, "Test Page") { t.Errorf("Expected ForUser to contain 'Test Page', got: %s", result.ForUser) } // ForLLM should contain summary + if !strings.Contains(result.ForLLM, "bytes") && !strings.Contains(result.ForLLM, "extractor") { t.Errorf("Expected ForLLM to contain summary, got: %s", result.ForLLM) } } // TestWebTool_WebFetch_JSON verifies JSON content handling + func TestWebTool_WebFetch_JSON(t *testing.T) { testData := map[string]string{"key": "value", "number": "123"} + expectedJSON, _ := json.MarshalIndent(testData, "", " ") server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write(expectedJSON) })) + defer server.Close() tool := NewWebFetchTool(50000) + ctx := context.Background() + args := map[string]any{ "url": server.URL, } @@ -64,20 +80,25 @@ func TestWebTool_WebFetch_JSON(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain formatted JSON + if !strings.Contains(result.ForUser, "key") && !strings.Contains(result.ForUser, "value") { t.Errorf("Expected ForUser to contain JSON data, got: %s", result.ForUser) } } // TestWebTool_WebFetch_InvalidURL verifies error handling for invalid URL + func TestWebTool_WebFetch_InvalidURL(t *testing.T) { tool := NewWebFetchTool(50000) + ctx := context.Background() + args := map[string]any{ "url": "not-a-valid-url", } @@ -85,20 +106,25 @@ func TestWebTool_WebFetch_InvalidURL(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error for invalid URL") } // Should contain error message (either "invalid URL" or scheme error) + if !strings.Contains(result.ForLLM, "URL") && !strings.Contains(result.ForUser, "URL") { t.Errorf("Expected error message for invalid URL, got ForLLM: %s", result.ForLLM) } } // TestWebTool_WebFetch_UnsupportedScheme verifies error handling for non-http URLs + func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) { tool := NewWebFetchTool(50000) + ctx := context.Background() + args := map[string]any{ "url": "ftp://example.com/file.txt", } @@ -106,48 +132,61 @@ func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error for unsupported URL scheme") } // Should mention only http/https allowed + if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForUser, "http/https") { t.Errorf("Expected scheme error message, got ForLLM: %s", result.ForLLM) } } // TestWebTool_WebFetch_MissingURL verifies error handling for missing URL + func TestWebTool_WebFetch_MissingURL(t *testing.T) { tool := NewWebFetchTool(50000) + ctx := context.Background() + args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when URL is missing") } // Should mention URL is required + if !strings.Contains(result.ForLLM, "url is required") && !strings.Contains(result.ForUser, "url is required") { t.Errorf("Expected 'url is required' message, got ForLLM: %s", result.ForLLM) } } // TestWebTool_WebFetch_Truncation verifies content truncation + func TestWebTool_WebFetch_Truncation(t *testing.T) { longContent := strings.Repeat("x", 20000) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/plain") + w.WriteHeader(http.StatusOK) + w.Write([]byte(longContent)) })) + defer server.Close() tool := NewWebFetchTool(1000) // Limit to 1000 chars + ctx := context.Background() + args := map[string]any{ "url": server.URL, } @@ -155,13 +194,17 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain truncated content (not the full 20000 chars) + resultMap := make(map[string]any) + json.Unmarshal([]byte(result.ForUser), &resultMap) + if text, ok := resultMap["text"].(string); ok { if len(text) > 1100 { // Allow some margin t.Errorf("Expected content to be truncated to ~1000 chars, got: %d", len(text)) @@ -169,63 +212,79 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) { } // Should be marked as truncated + if truncated, ok := resultMap["truncated"].(bool); !ok || !truncated { t.Errorf("Expected 'truncated' to be true in result") } } // TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing + func TestWebTool_WebSearch_NoApiKey(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""}) if err != nil { t.Fatalf("Unexpected error: %v", err) } + if tool != nil { t.Errorf("Expected nil tool when Brave API key is empty") } // Also nil when nothing is enabled + tool, err = NewWebSearchTool(WebSearchToolOptions{}) if err != nil { t.Fatalf("Unexpected error: %v", err) } + if tool != nil { t.Errorf("Expected nil tool when no provider is enabled") } } // TestWebTool_WebSearch_MissingQuery verifies error handling for missing query + func TestWebTool_WebSearch_MissingQuery(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5}) if err != nil { t.Fatalf("Unexpected error: %v", err) } + ctx := context.Background() + args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error when query is missing") } } // TestWebTool_WebFetch_HTMLExtraction verifies HTML text extraction + func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + w.Write( + []byte( `<html><body><script>alert('test');</script><style>body{color:red;}</style><h1>Title</h1><p>Content</p></body></html>`, ), ) })) + defer server.Close() tool := NewWebFetchTool(50000) + ctx := context.Background() + args := map[string]any{ "url": server.URL, } @@ -233,80 +292,105 @@ func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain extracted text (without script/style tags) + if !strings.Contains(result.ForUser, "Title") && !strings.Contains(result.ForUser, "Content") { t.Errorf("Expected ForUser to contain extracted text, got: %s", result.ForUser) } // Should NOT contain script or style tags + if strings.Contains(result.ForUser, "<script>") || strings.Contains(result.ForUser, "<style>") { t.Errorf("Expected script/style tags to be removed, got: %s", result.ForUser) } } // TestWebFetchTool_extractText verifies text extraction preserves newlines + func TestWebFetchTool_extractText(t *testing.T) { tool := &WebFetchTool{} tests := []struct { - name string - input string + name string + + input string + wantFunc func(t *testing.T, got string) }{ { - name: "preserves newlines between block elements", + name: "preserves newlines between block elements", + input: "<html><body><h1>Title</h1>\n<p>Paragraph 1</p>\n<p>Paragraph 2</p></body></html>", + wantFunc: func(t *testing.T, got string) { lines := strings.Split(got, "\n") + if len(lines) < 2 { t.Errorf("Expected multiple lines, got %d: %q", len(lines), got) } + if !strings.Contains(got, "Title") || !strings.Contains(got, "Paragraph 1") || + !strings.Contains(got, "Paragraph 2") { t.Errorf("Missing expected text: %q", got) } }, }, + { - name: "removes script and style tags", + name: "removes script and style tags", + input: "<script>alert('x');</script><style>body{}</style><p>Keep this</p>", + wantFunc: func(t *testing.T, got string) { if strings.Contains(got, "alert") || strings.Contains(got, "body{}") { t.Errorf("Expected script/style content removed, got: %q", got) } + if !strings.Contains(got, "Keep this") { t.Errorf("Expected 'Keep this' to remain, got: %q", got) } }, }, + { - name: "collapses excessive blank lines", + name: "collapses excessive blank lines", + input: "<p>A</p>\n\n\n\n\n<p>B</p>", + wantFunc: func(t *testing.T, got string) { if strings.Contains(got, "\n\n\n") { t.Errorf("Expected excessive blank lines collapsed, got: %q", got) } }, }, + { - name: "collapses horizontal whitespace", + name: "collapses horizontal whitespace", + input: "<p>hello world</p>", + wantFunc: func(t *testing.T, got string) { if strings.Contains(got, " ") { t.Errorf("Expected spaces collapsed, got: %q", got) } + if !strings.Contains(got, "hello world") { t.Errorf("Expected 'hello world', got: %q", got) } }, }, + { - name: "empty input", + name: "empty input", + input: "", + wantFunc: func(t *testing.T, got string) { if got != "" { t.Errorf("Expected empty string, got: %q", got) @@ -318,15 +402,19 @@ func TestWebFetchTool_extractText(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := tool.extractText(tt.input) + tt.wantFunc(t, got) }) } } // TestWebTool_WebFetch_MissingDomain verifies error handling for URL without domain + func TestWebTool_WebFetch_MissingDomain(t *testing.T) { tool := NewWebFetchTool(50000) + ctx := context.Background() + args := map[string]any{ "url": "https://", } @@ -334,11 +422,13 @@ func TestWebTool_WebFetch_MissingDomain(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result + if !result.IsError { t.Errorf("Expected error for URL without domain") } // Should mention missing domain + if !strings.Contains(result.ForLLM, "domain") && !strings.Contains(result.ForUser, "domain") { t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM) } @@ -349,14 +439,17 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { if err != nil { t.Fatalf("createHTTPClient() error: %v", err) } + if client.Timeout != 12*time.Second { t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second) } tr, ok := client.Transport.(*http.Transport) + if !ok { t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) } + if tr.Proxy == nil { t.Fatal("transport.Proxy is nil, want non-nil") } @@ -365,10 +458,12 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { if err != nil { t.Fatalf("http.NewRequest() error: %v", err) } + proxyURL, err := tr.Proxy(req) if err != nil { t.Fatalf("transport.Proxy(req) error: %v", err) } + if proxyURL == nil || proxyURL.String() != "http://127.0.0.1:7890" { t.Fatalf("proxy URL = %v, want %q", proxyURL, "http://127.0.0.1:7890") } @@ -376,6 +471,7 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { func TestCreateHTTPClient_InvalidProxy(t *testing.T) { _, err := createHTTPClient("://bad-proxy", 10*time.Second) + if err == nil { t.Fatal("createHTTPClient() expected error for invalid proxy URL, got nil") } @@ -388,17 +484,21 @@ func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) { } tr, ok := client.Transport.(*http.Transport) + if !ok { t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) } + req, err := http.NewRequest("GET", "https://example.com", nil) if err != nil { t.Fatalf("http.NewRequest() error: %v", err) } + proxyURL, err := tr.Proxy(req) if err != nil { t.Fatalf("transport.Proxy(req) error: %v", err) } + if proxyURL == nil || proxyURL.String() != "socks5://127.0.0.1:1080" { t.Fatalf("proxy URL = %v, want %q", proxyURL, "socks5://127.0.0.1:1080") } @@ -406,9 +506,11 @@ func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) { func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) { _, err := createHTTPClient("ftp://127.0.0.1:21", 10*time.Second) + if err == nil { t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil") } + if !strings.Contains(err.Error(), "unsupported proxy scheme") { t.Fatalf("error = %q, want to contain %q", err.Error(), "unsupported proxy scheme") } @@ -416,12 +518,19 @@ func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) { func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { t.Setenv("HTTP_PROXY", "http://127.0.0.1:8888") + t.Setenv("http_proxy", "http://127.0.0.1:8888") + t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8888") + t.Setenv("https_proxy", "http://127.0.0.1:8888") + t.Setenv("ALL_PROXY", "") + t.Setenv("all_proxy", "") + t.Setenv("NO_PROXY", "") + t.Setenv("no_proxy", "") client, err := createHTTPClient("", 10*time.Second) @@ -430,9 +539,11 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { } tr, ok := client.Transport.(*http.Transport) + if !ok { t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) } + if tr.Proxy == nil { t.Fatal("transport.Proxy is nil, want proxy function from environment") } @@ -441,6 +552,7 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { if err != nil { t.Fatalf("http.NewRequest() error: %v", err) } + if _, err := tr.Proxy(req); err != nil { t.Fatalf("transport.Proxy(req) error: %v", err) } @@ -451,9 +563,11 @@ func TestNewWebFetchToolWithProxy(t *testing.T) { if err != nil { t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) } + if tool.maxChars != 1024 { t.Fatalf("maxChars = %d, want %d", tool.maxChars, 1024) } + if tool.proxy != "http://127.0.0.1:7890" { t.Fatalf("proxy = %q, want %q", tool.proxy, "http://127.0.0.1:7890") } @@ -462,6 +576,7 @@ func TestNewWebFetchToolWithProxy(t *testing.T) { if err != nil { t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) } + if tool.maxChars != 50000 { t.Fatalf("default maxChars = %d, want %d", tool.maxChars, 50000) } @@ -470,18 +585,24 @@ func TestNewWebFetchToolWithProxy(t *testing.T) { func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("perplexity", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ - PerplexityEnabled: true, - PerplexityAPIKey: "k", + PerplexityEnabled: true, + + PerplexityAPIKey: "k", + PerplexityMaxResults: 3, - Proxy: "http://127.0.0.1:7890", + + Proxy: "http://127.0.0.1:7890", }) if err != nil { t.Fatalf("NewWebSearchTool() error: %v", err) } + p, ok := tool.provider.(*PerplexitySearchProvider) + if !ok { t.Fatalf("provider type = %T, want *PerplexitySearchProvider", tool.provider) } + if p.proxy != "http://127.0.0.1:7890" { t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") } @@ -489,18 +610,24 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("brave", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ - BraveEnabled: true, - BraveAPIKey: "k", + BraveEnabled: true, + + BraveAPIKey: "k", + BraveMaxResults: 3, - Proxy: "http://127.0.0.1:7890", + + Proxy: "http://127.0.0.1:7890", }) if err != nil { t.Fatalf("NewWebSearchTool() error: %v", err) } + p, ok := tool.provider.(*BraveSearchProvider) + if !ok { t.Fatalf("provider type = %T, want *BraveSearchProvider", tool.provider) } + if p.proxy != "http://127.0.0.1:7890" { t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") } @@ -508,17 +635,22 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("duckduckgo", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ - DuckDuckGoEnabled: true, + DuckDuckGoEnabled: true, + DuckDuckGoMaxResults: 3, - Proxy: "http://127.0.0.1:7890", + + Proxy: "http://127.0.0.1:7890", }) if err != nil { t.Fatalf("NewWebSearchTool() error: %v", err) } + p, ok := tool.provider.(*DuckDuckGoSearchProvider) + if !ok { t.Fatalf("provider type = %T, want *DuckDuckGoSearchProvider", tool.provider) } + if p.proxy != "http://127.0.0.1:7890" { t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") } @@ -526,50 +658,69 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { } // TestWebTool_TavilySearch_Success verifies successful Tavily search + func TestWebTool_TavilySearch_Success(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != "POST" { t.Errorf("Expected POST request, got %s", r.Method) } + if r.Header.Get("Content-Type") != "application/json" { t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) } // Verify payload + var payload map[string]any + json.NewDecoder(r.Body).Decode(&payload) + if payload["api_key"] != "test-key" { t.Errorf("Expected api_key test-key, got %v", payload["api_key"]) } + if payload["query"] != "test query" { t.Errorf("Expected query 'test query', got %v", payload["query"]) } // Return mock response + response := map[string]any{ "results": []map[string]any{ { - "title": "Test Result 1", - "url": "https://example.com/1", + "title": "Test Result 1", + + "url": "https://example.com/1", + "content": "Content for result 1", }, + { - "title": "Test Result 2", - "url": "https://example.com/2", + "title": "Test Result 2", + + "url": "https://example.com/2", + "content": "Content for result 2", }, }, } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(response) })) + defer server.Close() tool, err := NewWebSearchTool(WebSearchToolOptions{ - TavilyEnabled: true, - TavilyAPIKey: "test-key", - TavilyBaseURL: server.URL, + TavilyEnabled: true, + + TavilyAPIKey: "test-key", + + TavilyBaseURL: server.URL, + TavilyMaxResults: 5, }) if err != nil { @@ -577,6 +728,7 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { } ctx := context.Background() + args := map[string]any{ "query": "test query", } @@ -584,17 +736,21 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error + if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain result titles and URLs + if !strings.Contains(result.ForUser, "Test Result 1") || + !strings.Contains(result.ForUser, "https://example.com/1") { t.Errorf("Expected results in output, got: %s", result.ForUser) } // Should mention via Tavily + if !strings.Contains(result.ForUser, "via Tavily") { t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) } diff --git a/pkg/tools/workspace_ctx.go b/pkg/tools/workspace_ctx.go index 12bc3e33d..9460431a9 100644 --- a/pkg/tools/workspace_ctx.go +++ b/pkg/tools/workspace_ctx.go @@ -8,52 +8,72 @@ import ( type ( workspaceOverrideKey struct{} - overrideFsKey struct{} + + overrideFsKey struct{} ) // WithWorkspaceOverride returns a context carrying a workspace override path + // and a pre-built sandboxFs for that workspace. Tools will resolve file + // operations against this path instead of the original workspace. + // The cached sandboxFs is reused across all resolveFS calls on the same context, + // avoiding per-operation allocation. + func WithWorkspaceOverride(ctx context.Context, workspace string) context.Context { ctx = context.WithValue(ctx, workspaceOverrideKey{}, workspace) + ctx = context.WithValue(ctx, overrideFsKey{}, &sandboxFs{workspace: workspace}) + return ctx } // WorkspaceOverrideFromCtx extracts the workspace override from context, or "". + func WorkspaceOverrideFromCtx(ctx context.Context) string { if v, ok := ctx.Value(workspaceOverrideKey{}).(string); ok { return v } + return "" } // resolveFS returns a fileSystem applying workspace override from context. + // Paths under "memory/" are excluded (always use original workspace). + // For sandboxFs: returns the cached override instance from context. + // For hostFs (unrestricted): returns as-is. + func resolveFS(ctx context.Context, fs fileSystem, path string) fileSystem { override := WorkspaceOverrideFromCtx(ctx) + if override == "" { return fs } // memory/ paths always use original workspace + if isMemoryPath(path) { return fs } // Only sandboxFs supports workspace override + if sfs, ok := fs.(*sandboxFs); ok { if sfs.workspace == override { return fs } + // Use cached sandboxFs from context + if cached, ok := ctx.Value(overrideFsKey{}).(*sandboxFs); ok { return cached } + return &sandboxFs{workspace: override} } @@ -61,16 +81,20 @@ func resolveFS(ctx context.Context, fs fileSystem, path string) fileSystem { } // isMemoryPath returns true for paths under the memory/ directory. + // Matches: "memory/MEMORY.md", "memory", "/workspace/memory/notes.md" + func isMemoryPath(path string) bool { p := filepath.ToSlash(filepath.Clean(path)) // Relative path starting with memory/ + if strings.HasPrefix(p, "memory/") || p == "memory" { return true } // Absolute path containing /memory/ or ending with /memory + if strings.Contains(p, "/memory/") || strings.HasSuffix(p, "/memory") { return true }