merge: feat/task-2-subagent-orchestration

TASKS-2 Subagent Orchestration (Container Model):
- ContainerMessage + channel-based escalation tools
- Deliberate (plan-first) and Exploratory preset workflows
- AskConductor/SubmitPlan (subagent) + Answer/Review (conductor)
- SessionRecorder DAG integration (question/plan turns)
- Environment context injection and orchestration guidance
- gci import ordering fix across pkg/{tools,agent,session}

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
dj-oyu 2026-03-05 02:03:02 +09:00
commit 5d376ce71b
78 changed files with 12603 additions and 1402 deletions

View file

@ -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 03 実装済み)
@ -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 ModeTASKS-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 03: SQLite SessionStore、LegacyAdapter、Fork/Report、CompactOldTurns、`/session` CLI コマンド、Mini App グラフ UI |
| [`todo/TASKS-4.md`](todo/TASKS-4.md) | ~~**Mini App & Static Serving**~~ ✅ 実装済み(`http.FileServer` 統合、テンプレート注入、Bun ビルド導線、frontend unit test + CI `pnpm test` |
| [`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 |

File diff suppressed because it is too large Load diff

View file

@ -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")
}

View file

@ -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)

View file

@ -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")
}

File diff suppressed because it is too large Load diff

View file

@ -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)
}

File diff suppressed because it is too large Load diff

File diff suppressed because it is too large Load diff

View file

@ -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)
}

View file

@ -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
}

View file

@ -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
}

View file

@ -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)
}

View file

@ -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
}

View file

@ -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")
}

View file

@ -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
}

View file

@ -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)
}

View file

@ -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
}

View file

@ -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 {

View file

@ -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
}
}

View file

@ -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)
}

View file

@ -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
}
}

View file

@ -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)
}

View file

@ -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,

View file

@ -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 {

View file

@ -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),
}
}

110
pkg/tools/ask_conductor.go Normal file
View file

@ -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()))
}
}

View file

@ -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")
}
}

View file

@ -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")
}

View file

@ -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)
}

View file

@ -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
}

View file

@ -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)
}

View file

@ -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"
}

View file

@ -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
}

View file

@ -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)
}

View file

@ -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
}

View file

@ -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")
}

View file

@ -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)

View file

@ -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)
}

View file

@ -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))
}

View file

@ -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)
}

View file

@ -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
}

View file

@ -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))
}

View file

@ -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
}

View file

@ -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")
}

View file

@ -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,
}
}

View file

@ -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'")
}

View file

@ -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
}

View file

@ -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)
}

View file

@ -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"])
}

View file

@ -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
}

File diff suppressed because it is too large Load diff

View file

@ -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)
}
}

View file

@ -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
}

File diff suppressed because it is too large Load diff

View file

@ -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)
}

View file

@ -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)
}

View file

@ -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")
}

View file

@ -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()
}

View file

@ -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")
}

View file

@ -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)
}

View file

@ -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)
}

View file

@ -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)
}

View file

@ -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))
}

File diff suppressed because it is too large Load diff

View file

@ -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")
}
}

View file

@ -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")
}
}

View file

@ -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
}

View file

@ -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)
}

View file

@ -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")
}

186
pkg/tools/submit_plan.go Normal file
View file

@ -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()))
}
}

View file

@ -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")
}
}

View file

@ -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
}

View file

@ -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
}

View file

@ -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)
}
}

View file

@ -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)
}

View file

@ -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
}

View file

@ -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) |
---