style: fix gci import ordering across pkg/{tools,agent,session}

Auto-fix via golangci-lint --fix to satisfy gci formatter rules
(standard → default → localmodule import grouping).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
dj-oyu 2026-03-05 01:37:13 +09:00
parent 7a5200b5ca
commit fa546ad733
60 changed files with 7487 additions and 888 deletions

View file

@ -12,70 +12,103 @@ import (
) )
// setupWorkspace creates a temporary workspace with standard directories and optional files. // setupWorkspace creates a temporary workspace with standard directories and optional files.
// Returns the tmpDir path; caller should defer os.RemoveAll(tmpDir). // Returns the tmpDir path; caller should defer os.RemoveAll(tmpDir).
func setupWorkspace(t *testing.T, files map[string]string) string { func setupWorkspace(t *testing.T, files map[string]string) string {
t.Helper() t.Helper()
tmpDir, err := os.MkdirTemp("", "picoclaw-test-*") tmpDir, err := os.MkdirTemp("", "picoclaw-test-*")
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755) os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755)
os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755) os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755)
for name, content := range files { for name, content := range files {
dir := filepath.Dir(filepath.Join(tmpDir, name)) dir := filepath.Dir(filepath.Join(tmpDir, name))
os.MkdirAll(dir, 0o755) os.MkdirAll(dir, 0o755)
if err := os.WriteFile(filepath.Join(tmpDir, name), []byte(content), 0o644); err != nil { if err := os.WriteFile(filepath.Join(tmpDir, name), []byte(content), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
} }
return tmpDir return tmpDir
} }
// TestSingleSystemMessage verifies that BuildMessages always produces exactly one // TestSingleSystemMessage verifies that BuildMessages always produces exactly one
// system message regardless of summary/history variations. // system message regardless of summary/history variations.
// Fix: multiple system messages break Anthropic (top-level system param) and // Fix: multiple system messages break Anthropic (top-level system param) and
// Codex (only reads last system message as instructions). // Codex (only reads last system message as instructions).
func TestSingleSystemMessage(t *testing.T) { func TestSingleSystemMessage(t *testing.T) {
tmpDir := setupWorkspace(t, map[string]string{ tmpDir := setupWorkspace(t, map[string]string{
"IDENTITY.md": "# Identity\nTest agent.", "IDENTITY.md": "# Identity\nTest agent.",
}) })
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
tests := []struct { tests := []struct {
name string name string
history []providers.Message history []providers.Message
summary string summary string
message string message string
}{ }{
{ {
name: "no summary, no history", name: "no summary, no history",
summary: "", summary: "",
message: "hello", message: "hello",
}, },
{ {
name: "with summary", name: "with summary",
summary: "Previous conversation discussed X", summary: "Previous conversation discussed X",
message: "hello", message: "hello",
}, },
{ {
name: "with history and summary", name: "with history and summary",
history: []providers.Message{ history: []providers.Message{
{Role: "user", Content: "hi"}, {Role: "user", Content: "hi"},
{Role: "assistant", Content: "hello"}, {Role: "assistant", Content: "hello"},
}, },
summary: strings.Repeat("Long summary text. ", 50), summary: strings.Repeat("Long summary text. ", 50),
message: "new message", message: "new message",
}, },
{ {
name: "system message in history is filtered", name: "system message in history is filtered",
history: []providers.Message{ history: []providers.Message{
{Role: "system", Content: "stale system prompt from previous session"}, {Role: "system", Content: "stale system prompt from previous session"},
{Role: "user", Content: "hi"}, {Role: "user", Content: "hi"},
{Role: "assistant", Content: "hello"}, {Role: "assistant", Content: "hello"},
}, },
summary: "", summary: "",
message: "new message", message: "new message",
}, },
} }
@ -85,35 +118,44 @@ func TestSingleSystemMessage(t *testing.T) {
msgs := cb.BuildMessages(tt.history, tt.summary, tt.message, nil, "test", "chat1") msgs := cb.BuildMessages(tt.history, tt.summary, tt.message, nil, "test", "chat1")
systemCount := 0 systemCount := 0
for _, m := range msgs { for _, m := range msgs {
if m.Role == "system" { if m.Role == "system" {
systemCount++ systemCount++
} }
} }
if systemCount != 1 { if systemCount != 1 {
t.Errorf("expected exactly 1 system message, got %d", systemCount) t.Errorf("expected exactly 1 system message, got %d", systemCount)
} }
if msgs[0].Role != "system" { if msgs[0].Role != "system" {
t.Errorf("first message should be system, got %s", msgs[0].Role) t.Errorf("first message should be system, got %s", msgs[0].Role)
} }
if msgs[len(msgs)-1].Role != "user" { if msgs[len(msgs)-1].Role != "user" {
t.Errorf("last message should be user, got %s", msgs[len(msgs)-1].Role) t.Errorf("last message should be user, got %s", msgs[len(msgs)-1].Role)
} }
// System message must contain identity (static) and time (dynamic) // System message must contain identity (static) and time (dynamic)
sys := msgs[0].Content sys := msgs[0].Content
if !strings.Contains(sys, "picoclaw") { if !strings.Contains(sys, "picoclaw") {
t.Error("system message missing identity") t.Error("system message missing identity")
} }
if !strings.Contains(sys, "Current Time") { if !strings.Contains(sys, "Current Time") {
t.Error("system message missing dynamic time context") t.Error("system message missing dynamic time context")
} }
// Summary handling // Summary handling
if tt.summary != "" { if tt.summary != "" {
if !strings.Contains(sys, "CONTEXT_SUMMARY:") { if !strings.Contains(sys, "CONTEXT_SUMMARY:") {
t.Error("summary present but CONTEXT_SUMMARY prefix missing") t.Error("summary present but CONTEXT_SUMMARY prefix missing")
} }
if !strings.Contains(sys, tt.summary[:20]) { if !strings.Contains(sys, tt.summary[:20]) {
t.Error("summary content not found in system message") 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 // TestMtimeAutoInvalidation verifies that the cache detects source file changes
// via mtime without requiring explicit InvalidateCache(). // via mtime without requiring explicit InvalidateCache().
// Fix: original implementation had no auto-invalidation — edits to bootstrap files, // Fix: original implementation had no auto-invalidation — edits to bootstrap files,
// memory, or skills were invisible until process restart. // memory, or skills were invisible until process restart.
func TestMtimeAutoInvalidation(t *testing.T) { func TestMtimeAutoInvalidation(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
file string // relative path inside workspace file string // relative path inside workspace
contentV1 string contentV1 string
contentV2 string contentV2 string
checkField string // substring to verify in rebuilt prompt checkField string // substring to verify in rebuilt prompt
}{ }{
{ {
name: "bootstrap file change", name: "bootstrap file change",
file: "IDENTITY.md", file: "IDENTITY.md",
contentV1: "# Original Identity", contentV1: "# Original Identity",
contentV2: "# Updated Identity", contentV2: "# Updated Identity",
checkField: "Updated Identity", checkField: "Updated Identity",
}, },
{ {
name: "memory file change", name: "memory file change",
file: "memory/MEMORY.md", file: "memory/MEMORY.md",
contentV1: "# Memory\nUser likes Go.", contentV1: "# Memory\nUser likes Go.",
contentV2: "# Memory\nUser likes Rust.", contentV2: "# Memory\nUser likes Rust.",
checkField: "User likes Rust", checkField: "User likes Rust",
}, },
} }
@ -157,6 +216,7 @@ func TestMtimeAutoInvalidation(t *testing.T) {
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
tmpDir := setupWorkspace(t, map[string]string{tt.file: tt.contentV1}) tmpDir := setupWorkspace(t, map[string]string{tt.file: tt.contentV1})
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
@ -164,26 +224,39 @@ func TestMtimeAutoInvalidation(t *testing.T) {
sp1 := cb.BuildSystemPromptWithCache() sp1 := cb.BuildSystemPromptWithCache()
// Overwrite file and set future mtime to ensure detection. // Overwrite file and set future mtime to ensure detection.
// Use 2s offset for filesystem mtime resolution safety (some FS // Use 2s offset for filesystem mtime resolution safety (some FS
// have 1s or coarser granularity, especially in CI containers). // have 1s or coarser granularity, especially in CI containers).
fullPath := filepath.Join(tmpDir, tt.file) fullPath := filepath.Join(tmpDir, tt.file)
os.WriteFile(fullPath, []byte(tt.contentV2), 0o644) os.WriteFile(fullPath, []byte(tt.contentV2), 0o644)
future := time.Now().Add(2 * time.Second) future := time.Now().Add(2 * time.Second)
os.Chtimes(fullPath, future, future) os.Chtimes(fullPath, future, future)
// Verify sourceFilesChangedLocked detects the mtime change // Verify sourceFilesChangedLocked detects the mtime change
cb.systemPromptMutex.RLock() cb.systemPromptMutex.RLock()
changed := cb.sourceFilesChangedLocked() changed := cb.sourceFilesChangedLocked()
cb.systemPromptMutex.RUnlock() cb.systemPromptMutex.RUnlock()
if !changed { if !changed {
t.Fatalf("sourceFilesChangedLocked() should detect %s change", tt.file) t.Fatalf("sourceFilesChangedLocked() should detect %s change", tt.file)
} }
// Should auto-rebuild without explicit InvalidateCache() // Should auto-rebuild without explicit InvalidateCache()
sp2 := cb.BuildSystemPromptWithCache() sp2 := cb.BuildSystemPromptWithCache()
if sp1 == sp2 { if sp1 == sp2 {
t.Errorf("cache not rebuilt after %s change", tt.file) t.Errorf("cache not rebuilt after %s change", tt.file)
} }
if !strings.Contains(sp2, tt.checkField) { if !strings.Contains(sp2, tt.checkField) {
t.Errorf("rebuilt prompt missing expected content %q", 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 // Skills directory mtime change
t.Run("skills dir change", func(t *testing.T) { t.Run("skills dir change", func(t *testing.T) {
tmpDir := setupWorkspace(t, nil) tmpDir := setupWorkspace(t, nil)
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
_ = cb.BuildSystemPromptWithCache() // populate cache _ = cb.BuildSystemPromptWithCache() // populate cache
// Touch skills directory (simulate new skill installed) // Touch skills directory (simulate new skill installed)
skillsDir := filepath.Join(tmpDir, "skills") skillsDir := filepath.Join(tmpDir, "skills")
future := time.Now().Add(2 * time.Second) future := time.Now().Add(2 * time.Second)
os.Chtimes(skillsDir, future, future) os.Chtimes(skillsDir, future, future)
// Verify sourceFilesChangedLocked detects it (cache is rebuilt) // Verify sourceFilesChangedLocked detects it (cache is rebuilt)
// We confirm by checking internal state: a second call should rebuild. // We confirm by checking internal state: a second call should rebuild.
cb.systemPromptMutex.RLock() cb.systemPromptMutex.RLock()
changed := cb.sourceFilesChangedLocked() changed := cb.sourceFilesChangedLocked()
cb.systemPromptMutex.RUnlock() cb.systemPromptMutex.RUnlock()
if !changed { if !changed {
t.Error("sourceFilesChangedLocked() should detect skills dir mtime change") 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 // TestExplicitInvalidateCache verifies that InvalidateCache() forces a rebuild
// even when source files haven't changed (useful for tests and reload commands). // even when source files haven't changed (useful for tests and reload commands).
func TestExplicitInvalidateCache(t *testing.T) { func TestExplicitInvalidateCache(t *testing.T) {
tmpDir := setupWorkspace(t, map[string]string{ tmpDir := setupWorkspace(t, map[string]string{
"IDENTITY.md": "# Test Identity", "IDENTITY.md": "# Test Identity",
}) })
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
sp1 := cb.BuildSystemPromptWithCache() sp1 := cb.BuildSystemPromptWithCache()
cb.InvalidateCache() cb.InvalidateCache()
sp2 := cb.BuildSystemPromptWithCache() sp2 := cb.BuildSystemPromptWithCache()
if sp1 != sp2 { if sp1 != sp2 {
@ -233,29 +322,39 @@ func TestExplicitInvalidateCache(t *testing.T) {
} }
// Verify cachedAt was reset // Verify cachedAt was reset
cb.InvalidateCache() cb.InvalidateCache()
cb.systemPromptMutex.RLock() cb.systemPromptMutex.RLock()
if !cb.cachedAt.IsZero() { if !cb.cachedAt.IsZero() {
t.Error("cachedAt should be zero after InvalidateCache()") t.Error("cachedAt should be zero after InvalidateCache()")
} }
cb.systemPromptMutex.RUnlock() cb.systemPromptMutex.RUnlock()
} }
// TestCacheStability verifies that the static prompt is stable across repeated calls // TestCacheStability verifies that the static prompt is stable across repeated calls
// when no files change (regression test for issue #607). // when no files change (regression test for issue #607).
func TestCacheStability(t *testing.T) { func TestCacheStability(t *testing.T) {
tmpDir := setupWorkspace(t, map[string]string{ tmpDir := setupWorkspace(t, map[string]string{
"IDENTITY.md": "# Identity\nContent", "IDENTITY.md": "# Identity\nContent",
"SOUL.md": "# Soul\nContent", "SOUL.md": "# Soul\nContent",
}) })
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
results := make([]string, 5) results := make([]string, 5)
for i := range results { for i := range results {
results[i] = cb.BuildSystemPromptWithCache() results[i] = cb.BuildSystemPromptWithCache()
} }
for i := 1; i < len(results); i++ { for i := 1; i < len(results); i++ {
if results[i] != results[0] { if results[i] != results[0] {
t.Errorf("cached prompt changed between call 0 and %d", i) 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 // Static prompt must NOT contain per-request data
if strings.Contains(results[0], "Current Time") { if strings.Contains(results[0], "Current Time") {
t.Error("static cached prompt should not contain time (added dynamically)") t.Error("static cached prompt should not contain time (added dynamically)")
} }
} }
// TestNewFileCreationInvalidatesCache verifies that creating a source file that // TestNewFileCreationInvalidatesCache verifies that creating a source file that
// did not exist when the cache was built triggers a cache rebuild. // did not exist when the cache was built triggers a cache rebuild.
// This catches the "from nothing to something" edge case that the old // This catches the "from nothing to something" edge case that the old
// modifiedSince (return false on stat error) would miss. // modifiedSince (return false on stat error) would miss.
func TestNewFileCreationInvalidatesCache(t *testing.T) { func TestNewFileCreationInvalidatesCache(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
file string // relative path inside workspace file string // relative path inside workspace
content string content string
checkField string // substring to verify in rebuilt prompt checkField string // substring to verify in rebuilt prompt
}{ }{
{ {
name: "new bootstrap file", name: "new bootstrap file",
file: "SOUL.md", file: "SOUL.md",
content: "# Soul\nBe kind and helpful.", content: "# Soul\nBe kind and helpful.",
checkField: "Be kind and helpful", checkField: "Be kind and helpful",
}, },
{ {
name: "new memory file", name: "new memory file",
file: "memory/MEMORY.md", file: "memory/MEMORY.md",
content: "# Memory\nUser prefers dark mode.", content: "# Memory\nUser prefers dark mode.",
checkField: "User prefers dark mode", checkField: "User prefers dark mode",
}, },
} }
@ -296,29 +410,41 @@ func TestNewFileCreationInvalidatesCache(t *testing.T) {
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
// Start with an empty workspace (no bootstrap/memory files) // Start with an empty workspace (no bootstrap/memory files)
tmpDir := setupWorkspace(t, nil) tmpDir := setupWorkspace(t, nil)
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
// Populate cache — file does not exist yet // Populate cache — file does not exist yet
sp1 := cb.BuildSystemPromptWithCache() sp1 := cb.BuildSystemPromptWithCache()
if strings.Contains(sp1, tt.checkField) { if strings.Contains(sp1, tt.checkField) {
t.Fatalf("prompt should not contain %q before file is created", tt.checkField) t.Fatalf("prompt should not contain %q before file is created", tt.checkField)
} }
// Create the file after cache was built // Create the file after cache was built
fullPath := filepath.Join(tmpDir, tt.file) fullPath := filepath.Join(tmpDir, tt.file)
os.MkdirAll(filepath.Dir(fullPath), 0o755) os.MkdirAll(filepath.Dir(fullPath), 0o755)
if err := os.WriteFile(fullPath, []byte(tt.content), 0o644); err != nil { if err := os.WriteFile(fullPath, []byte(tt.content), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
// Set future mtime to guarantee detection // Set future mtime to guarantee detection
future := time.Now().Add(2 * time.Second) future := time.Now().Add(2 * time.Second)
os.Chtimes(fullPath, future, future) os.Chtimes(fullPath, future, future)
// Cache should auto-invalidate because file went from absent -> present // Cache should auto-invalidate because file went from absent -> present
sp2 := cb.BuildSystemPromptWithCache() sp2 := cb.BuildSystemPromptWithCache()
if !strings.Contains(sp2, tt.checkField) { if !strings.Contains(sp2, tt.checkField) {
t.Errorf("cache not invalidated on new file creation: expected %q in prompt", 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 // TestSkillFileContentChange verifies that modifying a skill file's content
// (not just the directory structure) invalidates the cache. // (not just the directory structure) invalidates the cache.
// This is the scenario where directory mtime alone is insufficient — on most // This is the scenario where directory mtime alone is insufficient — on most
// filesystems, editing a file inside a directory does NOT update the parent // filesystems, editing a file inside a directory does NOT update the parent
// directory's mtime. // directory's mtime.
func TestSkillFileContentChange(t *testing.T) { func TestSkillFileContentChange(t *testing.T) {
skillMD := `--- skillMD := `---
name: test-skill name: test-skill
description: "A test skill" description: "A test skill"
--- ---
# Test Skill v1 # Test Skill v1
Original content.` Original content.`
tmpDir := setupWorkspace(t, map[string]string{ tmpDir := setupWorkspace(t, map[string]string{
"skills/test-skill/SKILL.md": skillMD, "skills/test-skill/SKILL.md": skillMD,
}) })
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
// Populate cache // Populate cache
sp1 := cb.BuildSystemPromptWithCache() sp1 := cb.BuildSystemPromptWithCache()
_ = sp1 // cache is warm _ = sp1 // cache is warm
// Modify the skill file content (without touching the skills/ directory) // Modify the skill file content (without touching the skills/ directory)
updatedSkillMD := `--- updatedSkillMD := `---
name: test-skill name: test-skill
description: "An updated test skill" description: "An updated test skill"
--- ---
# Test Skill v2 # Test Skill v2
Updated content.` Updated content.`
skillPath := filepath.Join(tmpDir, "skills", "test-skill", "SKILL.md") skillPath := filepath.Join(tmpDir, "skills", "test-skill", "SKILL.md")
if err := os.WriteFile(skillPath, []byte(updatedSkillMD), 0o644); err != nil { if err := os.WriteFile(skillPath, []byte(updatedSkillMD), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
// Set future mtime on the skill file only (NOT the directory) // Set future mtime on the skill file only (NOT the directory)
future := time.Now().Add(2 * time.Second) future := time.Now().Add(2 * time.Second)
os.Chtimes(skillPath, future, future) os.Chtimes(skillPath, future, future)
// Verify that sourceFilesChangedLocked detects the content change // Verify that sourceFilesChangedLocked detects the content change
cb.systemPromptMutex.RLock() cb.systemPromptMutex.RLock()
changed := cb.sourceFilesChangedLocked() changed := cb.sourceFilesChangedLocked()
cb.systemPromptMutex.RUnlock() cb.systemPromptMutex.RUnlock()
if !changed { if !changed {
t.Error("sourceFilesChangedLocked() should detect skill file content change") t.Error("sourceFilesChangedLocked() should detect skill file content change")
} }
// Verify cache is actually rebuilt with new content // Verify cache is actually rebuilt with new content
sp2 := cb.BuildSystemPromptWithCache() sp2 := cb.BuildSystemPromptWithCache()
if sp1 == sp2 && strings.Contains(sp1, "test-skill") { if sp1 == sp2 && strings.Contains(sp1, "test-skill") {
// If the skill appeared in the prompt and the prompt didn't change, // If the skill appeared in the prompt and the prompt didn't change,
// the cache was not invalidated. // the cache was not invalidated.
t.Error("cache should be invalidated when skill file content changes") t.Error("cache should be invalidated when skill file content changes")
} }
} }
// TestConcurrentBuildSystemPromptWithCache verifies that multiple goroutines // TestConcurrentBuildSystemPromptWithCache verifies that multiple goroutines
// can safely call BuildSystemPromptWithCache concurrently without producing // can safely call BuildSystemPromptWithCache concurrently without producing
// empty results, panics, or data races. // empty results, panics, or data races.
// Run with: go test -race ./pkg/agent/ -run TestConcurrentBuildSystemPromptWithCache // Run with: go test -race ./pkg/agent/ -run TestConcurrentBuildSystemPromptWithCache
func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { func TestConcurrentBuildSystemPromptWithCache(t *testing.T) {
tmpDir := setupWorkspace(t, map[string]string{ tmpDir := setupWorkspace(t, map[string]string{
"IDENTITY.md": "# Identity\nConcurrency test agent.", "IDENTITY.md": "# Identity\nConcurrency test agent.",
"SOUL.md": "# Soul\nBe helpful.", "SOUL.md": "# Soul\nBe helpful.",
"memory/MEMORY.md": "# Memory\nUser prefers Go.", "memory/MEMORY.md": "# Memory\nUser prefers Go.",
"skills/demo/SKILL.md": "---\nname: demo\ndescription: \"demo skill\"\n---\n# Demo", "skills/demo/SKILL.md": "---\nname: demo\ndescription: \"demo skill\"\n---\n# Demo",
}) })
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
const goroutines = 20 const goroutines = 20
const iterations = 50 const iterations = 50
var wg sync.WaitGroup var wg sync.WaitGroup
errs := make(chan string, goroutines*iterations) errs := make(chan string, goroutines*iterations)
for g := range goroutines { for g := range goroutines {
wg.Add(1) wg.Add(1)
go func(id int) { go func(id int) {
defer wg.Done() defer wg.Done()
for i := range iterations { for i := range iterations {
result := cb.BuildSystemPromptWithCache() result := cb.BuildSystemPromptWithCache()
if result == "" { if result == "" {
errs <- "empty prompt returned" errs <- "empty prompt returned"
return return
} }
if !strings.Contains(result, "picoclaw") { if !strings.Contains(result, "picoclaw") {
errs <- "prompt missing identity" errs <- "prompt missing identity"
return return
} }
// Also exercise BuildMessages concurrently // Also exercise BuildMessages concurrently
msgs := cb.BuildMessages(nil, "", "hello", nil, "test", "chat") msgs := cb.BuildMessages(nil, "", "hello", nil, "test", "chat")
if len(msgs) < 2 { if len(msgs) < 2 {
errs <- "BuildMessages returned fewer than 2 messages" errs <- "BuildMessages returned fewer than 2 messages"
return return
} }
if msgs[0].Role != "system" { if msgs[0].Role != "system" {
errs <- "first message not system" errs <- "first message not system"
return return
} }
// Occasionally invalidate to exercise the write path // Occasionally invalidate to exercise the write path
if i%10 == 0 { if i%10 == 0 {
cb.InvalidateCache() cb.InvalidateCache()
} }
@ -439,6 +618,7 @@ func TestConcurrentBuildSystemPromptWithCache(t *testing.T) {
} }
wg.Wait() wg.Wait()
close(errs) close(errs)
for errMsg := range errs { for errMsg := range errs {
@ -449,64 +629,90 @@ func TestConcurrentBuildSystemPromptWithCache(t *testing.T) {
// BenchmarkBuildMessagesWithCache measures caching performance. // BenchmarkBuildMessagesWithCache measures caching performance.
// TestEmptyWorkspaceBaselineDetectsNewFiles verifies that when the cache is // TestEmptyWorkspaceBaselineDetectsNewFiles verifies that when the cache is
// built on an empty workspace (no tracked files exist), creating a file // built on an empty workspace (no tracked files exist), creating a file
// afterwards still triggers cache invalidation. This validates the // afterwards still triggers cache invalidation. This validates the
// time.Unix(1, 0) fallback for maxMtime: any real file's mtime is after epoch, // time.Unix(1, 0) fallback for maxMtime: any real file's mtime is after epoch,
// so fileChangedSince correctly detects the absent -> present transition AND // so fileChangedSince correctly detects the absent -> present transition AND
// the mtime comparison succeeds even without artificially inflated Chtimes. // the mtime comparison succeeds even without artificially inflated Chtimes.
func TestEmptyWorkspaceBaselineDetectsNewFiles(t *testing.T) { func TestEmptyWorkspaceBaselineDetectsNewFiles(t *testing.T) {
// Empty workspace: no bootstrap files, no memory, no skills content. // Empty workspace: no bootstrap files, no memory, no skills content.
tmpDir := setupWorkspace(t, nil) tmpDir := setupWorkspace(t, nil)
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
// Build cache — all tracked files are absent, maxMtime falls back to epoch. // Build cache — all tracked files are absent, maxMtime falls back to epoch.
sp1 := cb.BuildSystemPromptWithCache() sp1 := cb.BuildSystemPromptWithCache()
// Create a bootstrap file with natural mtime (no Chtimes manipulation). // Create a bootstrap file with natural mtime (no Chtimes manipulation).
// The file's mtime should be the current wall-clock time, which is // The file's mtime should be the current wall-clock time, which is
// strictly after time.Unix(1, 0). // strictly after time.Unix(1, 0).
soulPath := filepath.Join(tmpDir, "SOUL.md") soulPath := filepath.Join(tmpDir, "SOUL.md")
if err := os.WriteFile(soulPath, []byte("# Soul\nNewly created."), 0o644); err != nil { if err := os.WriteFile(soulPath, []byte("# Soul\nNewly created."), 0o644); err != nil {
t.Fatal(err) t.Fatal(err)
} }
// Cache should detect the new file via existedAtCache (absent -> present). // Cache should detect the new file via existedAtCache (absent -> present).
cb.systemPromptMutex.RLock() cb.systemPromptMutex.RLock()
changed := cb.sourceFilesChangedLocked() changed := cb.sourceFilesChangedLocked()
cb.systemPromptMutex.RUnlock() cb.systemPromptMutex.RUnlock()
if !changed { if !changed {
t.Fatal("sourceFilesChangedLocked should detect newly created file on empty workspace") t.Fatal("sourceFilesChangedLocked should detect newly created file on empty workspace")
} }
sp2 := cb.BuildSystemPromptWithCache() sp2 := cb.BuildSystemPromptWithCache()
if !strings.Contains(sp2, "Newly created") { if !strings.Contains(sp2, "Newly created") {
t.Error("rebuilt prompt should contain new file content") t.Error("rebuilt prompt should contain new file content")
} }
if sp1 == sp2 { if sp1 == sp2 {
t.Error("cache should have been invalidated after file creation") t.Error("cache should have been invalidated after file creation")
} }
} }
// BenchmarkBuildMessagesWithCache measures caching performance. // BenchmarkBuildMessagesWithCache measures caching performance.
func BenchmarkBuildMessagesWithCache(b *testing.B) { func BenchmarkBuildMessagesWithCache(b *testing.B) {
tmpDir, _ := os.MkdirTemp("", "picoclaw-bench-*") tmpDir, _ := os.MkdirTemp("", "picoclaw-bench-*")
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755) os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755)
os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755) os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755)
for _, name := range []string{"IDENTITY.md", "SOUL.md", "USER.md"} { for _, name := range []string{"IDENTITY.md", "SOUL.md", "USER.md"} {
os.WriteFile(filepath.Join(tmpDir, name), []byte(strings.Repeat("Content.\n", 10)), 0o644) os.WriteFile(filepath.Join(tmpDir, name), []byte(strings.Repeat("Content.\n", 10)), 0o644)
} }
cb := NewContextBuilder(tmpDir) cb := NewContextBuilder(tmpDir)
history := []providers.Message{ history := []providers.Message{
{Role: "user", Content: "previous message"}, {Role: "user", Content: "previous message"},
{Role: "assistant", Content: "previous response"}, {Role: "assistant", Content: "previous response"},
} }
b.ResetTimer() b.ResetTimer()
for i := 0; i < b.N; i++ { for i := 0; i < b.N; i++ {
_ = cb.BuildMessages(history, "summary", "new message", nil, "cli", "test") _ = 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 { func assistantWithTools(toolIDs ...string) providers.Message {
calls := make([]providers.ToolCall, len(toolIDs)) calls := make([]providers.ToolCall, len(toolIDs))
for i, id := range toolIDs { for i, id := range toolIDs {
calls[i] = providers.ToolCall{ID: id, Type: "function"} calls[i] = providers.ToolCall{ID: id, Type: "function"}
} }
return providers.Message{Role: "assistant", ToolCalls: calls} return providers.Message{Role: "assistant", ToolCalls: calls}
} }
@ -24,11 +26,13 @@ func toolResult(id string) providers.Message {
func TestSanitizeHistoryForProvider_EmptyHistory(t *testing.T) { func TestSanitizeHistoryForProvider_EmptyHistory(t *testing.T) {
result := sanitizeHistoryForProvider(nil) result := sanitizeHistoryForProvider(nil)
if len(result) != 0 { if len(result) != 0 {
t.Fatalf("expected empty, got %d messages", len(result)) t.Fatalf("expected empty, got %d messages", len(result))
} }
result = sanitizeHistoryForProvider([]providers.Message{}) result = sanitizeHistoryForProvider([]providers.Message{})
if len(result) != 0 { if len(result) != 0 {
t.Fatalf("expected empty, got %d messages", len(result)) 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) { func TestSanitizeHistoryForProvider_SingleToolCall(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "hello"), msg("user", "hello"),
assistantWithTools("A"), assistantWithTools("A"),
toolResult("A"), toolResult("A"),
msg("assistant", "done"), msg("assistant", "done"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 4 { if len(result) != 4 {
t.Fatalf("expected 4 messages, got %d", len(result)) t.Fatalf("expected 4 messages, got %d", len(result))
} }
assertRoles(t, result, "user", "assistant", "tool", "assistant") assertRoles(t, result, "user", "assistant", "tool", "assistant")
} }
func TestSanitizeHistoryForProvider_MultiToolCalls(t *testing.T) { func TestSanitizeHistoryForProvider_MultiToolCalls(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "do two things"), msg("user", "do two things"),
assistantWithTools("A", "B"), assistantWithTools("A", "B"),
toolResult("A"), toolResult("A"),
toolResult("B"), toolResult("B"),
msg("assistant", "both done"), msg("assistant", "both done"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 5 { if len(result) != 5 {
t.Fatalf("expected 5 messages, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 5 messages, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant") assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant")
} }
func TestSanitizeHistoryForProvider_AssistantToolCallAfterPlainAssistant(t *testing.T) { func TestSanitizeHistoryForProvider_AssistantToolCallAfterPlainAssistant(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "hi"), msg("user", "hi"),
msg("assistant", "thinking"), msg("assistant", "thinking"),
assistantWithTools("A"), assistantWithTools("A"),
toolResult("A"), toolResult("A"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 2 { if len(result) != 2 {
t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user", "assistant") assertRoles(t, result, "user", "assistant")
} }
func TestSanitizeHistoryForProvider_OrphanedLeadingTool(t *testing.T) { func TestSanitizeHistoryForProvider_OrphanedLeadingTool(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
toolResult("A"), toolResult("A"),
msg("user", "hello"), msg("user", "hello"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 1 { if len(result) != 1 {
t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user") assertRoles(t, result, "user")
} }
func TestSanitizeHistoryForProvider_ToolAfterUserDropped(t *testing.T) { func TestSanitizeHistoryForProvider_ToolAfterUserDropped(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "hello"), msg("user", "hello"),
toolResult("A"), toolResult("A"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 1 { if len(result) != 1 {
t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user") assertRoles(t, result, "user")
} }
func TestSanitizeHistoryForProvider_ToolAfterAssistantNoToolCalls(t *testing.T) { func TestSanitizeHistoryForProvider_ToolAfterAssistantNoToolCalls(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "hello"), msg("user", "hello"),
msg("assistant", "hi"), msg("assistant", "hi"),
toolResult("A"), toolResult("A"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 2 { if len(result) != 2 {
t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user", "assistant") assertRoles(t, result, "user", "assistant")
} }
func TestSanitizeHistoryForProvider_AssistantToolCallAtStart(t *testing.T) { func TestSanitizeHistoryForProvider_AssistantToolCallAtStart(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
assistantWithTools("A"), assistantWithTools("A"),
toolResult("A"), toolResult("A"),
msg("user", "hello"), msg("user", "hello"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 1 { if len(result) != 1 {
t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user") assertRoles(t, result, "user")
} }
func TestSanitizeHistoryForProvider_MultiToolCallsThenNewRound(t *testing.T) { func TestSanitizeHistoryForProvider_MultiToolCallsThenNewRound(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "do two things"), msg("user", "do two things"),
assistantWithTools("A", "B"), assistantWithTools("A", "B"),
toolResult("A"), toolResult("A"),
toolResult("B"), toolResult("B"),
msg("assistant", "done"), msg("assistant", "done"),
msg("user", "hi"), msg("user", "hi"),
assistantWithTools("C"), assistantWithTools("C"),
toolResult("C"), toolResult("C"),
msg("assistant", "done again"), msg("assistant", "done again"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 9 { if len(result) != 9 {
t.Fatalf("expected 9 messages, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 9 messages, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "user", "assistant", "tool", "assistant") assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "user", "assistant", "tool", "assistant")
} }
func TestSanitizeHistoryForProvider_ConsecutiveMultiToolRounds(t *testing.T) { func TestSanitizeHistoryForProvider_ConsecutiveMultiToolRounds(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "start"), msg("user", "start"),
assistantWithTools("A", "B"), assistantWithTools("A", "B"),
toolResult("A"), toolResult("A"),
toolResult("B"), toolResult("B"),
assistantWithTools("C", "D"), assistantWithTools("C", "D"),
toolResult("C"), toolResult("C"),
toolResult("D"), toolResult("D"),
msg("assistant", "all done"), msg("assistant", "all done"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 8 { if len(result) != 8 {
t.Fatalf("expected 8 messages, got %d: %+v", len(result), roles(result)) t.Fatalf("expected 8 messages, got %d: %+v", len(result), roles(result))
} }
assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "tool", "tool", "assistant") assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "tool", "tool", "assistant")
} }
func TestSanitizeHistoryForProvider_PlainConversation(t *testing.T) { func TestSanitizeHistoryForProvider_PlainConversation(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
msg("user", "hello"), msg("user", "hello"),
msg("assistant", "hi"), msg("assistant", "hi"),
msg("user", "how are you"), msg("user", "how are you"),
msg("assistant", "fine"), msg("assistant", "fine"),
} }
result := sanitizeHistoryForProvider(history) result := sanitizeHistoryForProvider(history)
if len(result) != 4 { if len(result) != 4 {
t.Fatalf("expected 4 messages, got %d", len(result)) t.Fatalf("expected 4 messages, got %d", len(result))
} }
assertRoles(t, result, "user", "assistant", "user", "assistant") assertRoles(t, result, "user", "assistant", "user", "assistant")
} }
func roles(msgs []providers.Message) []string { func roles(msgs []providers.Message) []string {
r := make([]string, len(msgs)) r := make([]string, len(msgs))
for i, m := range msgs { for i, m := range msgs {
r[i] = m.Role r[i] = m.Role
} }
return r return r
} }
func assertRoles(t *testing.T, msgs []providers.Message, expected ...string) { func assertRoles(t *testing.T, msgs []providers.Message, expected ...string) {
t.Helper() t.Helper()
if len(msgs) != len(expected) { if len(msgs) != len(expected) {
t.Fatalf("role count mismatch: got %v, want %v", roles(msgs), expected) t.Fatalf("role count mismatch: got %v, want %v", roles(msgs), expected)
} }
for i, exp := range expected { for i, exp := range expected {
if msgs[i].Role != exp { if msgs[i].Role != exp {
t.Errorf("message[%d]: got role %q, want %q", i, 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 { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
} }
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cfg := &config.Config{ cfg := &config.Config{
Agents: config.AgentsConfig{ Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{ Defaults: config.AgentDefaults{
Workspace: tmpDir, Workspace: tmpDir,
Model: "test-model", Model: "test-model",
MaxTokens: 1234, MaxTokens: 1234,
MaxToolIterations: 5, MaxToolIterations: 5,
}, },
}, },
} }
configuredTemp := 1.0 configuredTemp := 1.0
cfg.Agents.Defaults.Temperature = &configuredTemp cfg.Agents.Defaults.Temperature = &configuredTemp
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
if agent.MaxTokens != 1234 { if agent.MaxTokens != 1234 {
t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234) t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234)
} }
if agent.Temperature != 1.0 { if agent.Temperature != 1.0 {
t.Fatalf("Temperature = %f, want %f", 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 { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
} }
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cfg := &config.Config{ cfg := &config.Config{
Agents: config.AgentsConfig{ Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{ Defaults: config.AgentDefaults{
Workspace: tmpDir, Workspace: tmpDir,
Model: "test-model", Model: "test-model",
MaxTokens: 1234, MaxTokens: 1234,
MaxToolIterations: 5, MaxToolIterations: 5,
}, },
}, },
} }
configuredTemp := 0.0 configuredTemp := 0.0
cfg.Agents.Defaults.Temperature = &configuredTemp cfg.Agents.Defaults.Temperature = &configuredTemp
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
if agent.Temperature != 0.0 { if agent.Temperature != 0.0 {
@ -73,20 +86,25 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
} }
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cfg := &config.Config{ cfg := &config.Config{
Agents: config.AgentsConfig{ Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{ Defaults: config.AgentDefaults{
Workspace: tmpDir, Workspace: tmpDir,
Model: "test-model", Model: "test-model",
MaxTokens: 1234, MaxTokens: 1234,
MaxToolIterations: 5, MaxToolIterations: 5,
}, },
}, },
} }
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
if agent.Temperature != 0.7 { if agent.Temperature != 0.7 {
@ -99,33 +117,41 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
} }
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cfg := &config.Config{ cfg := &config.Config{
Agents: config.AgentsConfig{ Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{ Defaults: config.AgentDefaults{
Workspace: tmpDir, Workspace: tmpDir,
Model: "step-3.5-flash", Model: "step-3.5-flash",
}, },
}, },
ModelList: []config.ModelConfig{ ModelList: []config.ModelConfig{
{ {
ModelName: "step-3.5-flash", ModelName: "step-3.5-flash",
Model: "openrouter/stepfun/step-3.5-flash:free", Model: "openrouter/stepfun/step-3.5-flash:free",
APIBase: "https://openrouter.ai/api/v1", APIBase: "https://openrouter.ai/api/v1",
}, },
}, },
} }
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
if len(agent.Candidates) != 1 { if len(agent.Candidates) != 1 {
t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates))
} }
if agent.Candidates[0].Provider != "openrouter" { if agent.Candidates[0].Provider != "openrouter" {
t.Fatalf("candidate provider = %q, want %q", 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" { 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") 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 { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
} }
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
cfg := &config.Config{ cfg := &config.Config{
Agents: config.AgentsConfig{ Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{ Defaults: config.AgentDefaults{
Workspace: tmpDir, Workspace: tmpDir,
Model: "glm-5", Model: "glm-5",
}, },
}, },
ModelList: []config.ModelConfig{ ModelList: []config.ModelConfig{
{ {
ModelName: "glm-5", ModelName: "glm-5",
Model: "glm-5", Model: "glm-5",
APIBase: "https://api.z.ai/api/coding/paas/v4", APIBase: "https://api.z.ai/api/coding/paas/v4",
}, },
}, },
} }
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
if len(agent.Candidates) != 1 { if len(agent.Candidates) != 1 {
t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates))
} }
if agent.Candidates[0].Provider != "openai" { if agent.Candidates[0].Provider != "openai" {
t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openai") t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openai")
} }
if agent.Candidates[0].Model != "glm-5" { if agent.Candidates[0].Model != "glm-5" {
t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "glm-5") t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "glm-5")
} }

View file

@ -12,63 +12,93 @@ import (
) )
// makeOrchTestLoop creates a minimal AgentLoop with a temp workspace and // makeOrchTestLoop creates a minimal AgentLoop with a temp workspace and
// a real Broadcaster wired as the reporter. // a real Broadcaster wired as the reporter.
// Returns the loop, the broadcaster, and a cleanup function. // Returns the loop, the broadcaster, and a cleanup function.
func makeOrchTestLoop(t *testing.T) (*AgentLoop, *orch.Broadcaster) { func makeOrchTestLoop(t *testing.T) (*AgentLoop, *orch.Broadcaster) {
t.Helper() t.Helper()
tmpDir, err := os.MkdirTemp("", "agent-orch-test-*") tmpDir, err := os.MkdirTemp("", "agent-orch-test-*")
if err != nil { if err != nil {
t.Fatalf("MkdirTemp: %v", err) t.Fatalf("MkdirTemp: %v", err)
} }
t.Cleanup(func() { os.RemoveAll(tmpDir) }) t.Cleanup(func() { os.RemoveAll(tmpDir) })
cfg := &config.Config{ cfg := &config.Config{
Agents: config.AgentsConfig{ Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{ Defaults: config.AgentDefaults{
Workspace: tmpDir, Workspace: tmpDir,
Model: "test-model", Model: "test-model",
MaxTokens: 512, MaxTokens: 512,
MaxToolIterations: 5, MaxToolIterations: 5,
}, },
}, },
} }
al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{})
b := orch.NewBroadcaster() b := orch.NewBroadcaster()
al.SetOrchReporter(b) al.SetOrchReporter(b)
return al, b return al, b
} }
// collectOrchEvents drains the subscriber channel until an agent_gc event // collectOrchEvents drains the subscriber channel until an agent_gc event
// arrives or the deadline is exceeded. // arrives or the deadline is exceeded.
func collectOrchEvents(t *testing.T, ch <-chan orch.Event, timeout time.Duration) []orch.Event { func collectOrchEvents(t *testing.T, ch <-chan orch.Event, timeout time.Duration) []orch.Event {
t.Helper() t.Helper()
var events []orch.Event var events []orch.Event
deadline := time.After(timeout) deadline := time.After(timeout)
for { for {
select { select {
case ev := <-ch: case ev := <-ch:
events = append(events, ev) events = append(events, ev)
if ev.Type == "agent_gc" { if ev.Type == "agent_gc" {
return events return events
} }
case <-deadline: case <-deadline:
t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events)
} }
} }
} }
// TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC verifies that a main // TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC verifies that a main
// session processed via ProcessDirect emits the full lifecycle: // session processed via ProcessDirect emits the full lifecycle:
// //
// agent_spawn(sessionKey) → agent_state(waiting) → agent_gc(completed) // agent_spawn(sessionKey) → agent_state(waiting) → agent_gc(completed)
// //
// and that the Broadcaster snapshot is empty after the call returns. // and that the Broadcaster snapshot is empty after the call returns.
func TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC(t *testing.T) { func TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC(t *testing.T) {
al, b := makeOrchTestLoop(t) al, b := makeOrchTestLoop(t)
sub := b.Subscribe() sub := b.Subscribe()
defer b.Unsubscribe(sub) defer b.Unsubscribe(sub)
const sessionKey = "orch-test-session" const sessionKey = "orch-test-session"
_, err := al.ProcessDirect(context.Background(), "hello", sessionKey) _, err := al.ProcessDirect(context.Background(), "hello", sessionKey)
if err != nil { if err != nil {
t.Fatalf("ProcessDirect: %v", err) t.Fatalf("ProcessDirect: %v", err)
@ -77,39 +107,51 @@ func TestAgentLoop_ProcessDirect_EmitsSpawnWaitingGC(t *testing.T) {
events := collectOrchEvents(t, sub.Ch, 5*time.Second) events := collectOrchEvents(t, sub.Ch, 5*time.Second)
// First event: agent_spawn with correct ID. // First event: agent_spawn with correct ID.
if events[0].Type != "agent_spawn" || events[0].ID != sessionKey { if events[0].Type != "agent_spawn" || events[0].ID != sessionKey {
t.Errorf("first event must be agent_spawn(%s), got: %+v", sessionKey, events[0]) t.Errorf("first event must be agent_spawn(%s), got: %+v", sessionKey, events[0])
} }
// At least one agent_state(waiting) for this session. // At least one agent_state(waiting) for this session.
var hasWaiting bool var hasWaiting bool
for _, ev := range events { for _, ev := range events {
if ev.Type == "agent_state" && ev.ID == sessionKey && ev.State == "waiting" { if ev.Type == "agent_state" && ev.ID == sessionKey && ev.State == "waiting" {
hasWaiting = true hasWaiting = true
break break
} }
} }
if !hasWaiting { if !hasWaiting {
t.Errorf("missing agent_state(waiting) for %s; events: %+v", sessionKey, events) t.Errorf("missing agent_state(waiting) for %s; events: %+v", sessionKey, events)
} }
// Last event: agent_gc(completed) for this session. // Last event: agent_gc(completed) for this session.
last := events[len(events)-1] last := events[len(events)-1]
if last.Type != "agent_gc" || last.ID != sessionKey || last.Reason != "completed" { 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) t.Errorf("last event must be agent_gc(completed,%s), got: %+v", sessionKey, last)
} }
// Snapshot must be empty — session removed on GC. // Snapshot must be empty — session removed on GC.
if snap := b.Snapshot(); len(snap) != 0 { if snap := b.Snapshot(); len(snap) != 0 {
t.Errorf("snapshot must be empty after GC, got: %v", snap) t.Errorf("snapshot must be empty after GC, got: %v", snap)
} }
} }
// TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC verifies that heartbeat // TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC verifies that heartbeat
// sessions appear on canvas with sessionKey = "heartbeat". // sessions appear on canvas with sessionKey = "heartbeat".
func TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC(t *testing.T) { func TestAgentLoop_ProcessHeartbeat_EmitsSpawnAndGC(t *testing.T) {
al, b := makeOrchTestLoop(t) al, b := makeOrchTestLoop(t)
sub := b.Subscribe() sub := b.Subscribe()
defer b.Unsubscribe(sub) defer b.Unsubscribe(sub)
_, err := al.ProcessHeartbeat(context.Background(), "check system", "heartbeat-chan", "none") _, 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) events := collectOrchEvents(t, sub.Ch, 5*time.Second)
// ProcessHeartbeat always uses sessionKey = "heartbeat". // ProcessHeartbeat always uses sessionKey = "heartbeat".
const want = "heartbeat" const want = "heartbeat"
if events[0].Type != "agent_spawn" || events[0].ID != want { if events[0].Type != "agent_spawn" || events[0].ID != want {
t.Errorf("first event must be agent_spawn(%s), got: %+v", want, events[0]) t.Errorf("first event must be agent_spawn(%s), got: %+v", want, events[0])
} }
last := events[len(events)-1] last := events[len(events)-1]
if last.Type != "agent_gc" || last.ID != want || last.Reason != "completed" { 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) 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()) { func newTestMemoryStore(t *testing.T) (*MemoryStore, func()) {
t.Helper() t.Helper()
tmpDir, err := os.MkdirTemp("", "memory-test-*") tmpDir, err := os.MkdirTemp("", "memory-test-*")
if err != nil { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
} }
ms := NewMemoryStore(tmpDir) ms := NewMemoryStore(tmpDir)
return ms, func() { os.RemoveAll(tmpDir) } return ms, func() { os.RemoveAll(tmpDir) }
} }
const testPlanInterviewing = `# Active Plan const testPlanInterviewing = `# Active Plan
> Task: Set up server monitoring > Task: Set up server monitoring
> Status: interviewing > Status: interviewing
> Phase: 1 > Phase: 1
` `
const testPlanExecuting = `# Active Plan const testPlanExecuting = `# Active Plan
> Task: Set up server monitoring > Task: Set up server monitoring
> Status: executing > Status: executing
> Phase: 2 > Phase: 2
## Phase 1: Prometheus Install ## Phase 1: Prometheus Install
- [x] Install Prometheus - [x] Install Prometheus
- [x] Configure node_exporter - [x] Configure node_exporter
## Phase 2: Grafana Setup ## Phase 2: Grafana Setup
- [ ] Install Grafana - [ ] Install Grafana
- [ ] Create dashboard - [ ] Create dashboard
## Phase 3: Alert Configuration ## Phase 3: Alert Configuration
- [ ] Set up alert rules - [ ] Set up alert rules
- [ ] Configure Telegram notifications - [ ] Configure Telegram notifications
## Commands ## Commands
build: go build ./... build: go build ./...
test: go test ./pkg/... -count=1 test: go test ./pkg/... -count=1
lint: golangci-lint run lint: golangci-lint run
## Context ## Context
Pi: Debian Bookworm arm64, ports: 3000/9090 Pi: Debian Bookworm arm64, ports: 3000/9090
` `
const testPlanPhase1Complete = `# Active Plan const testPlanPhase1Complete = `# Active Plan
> Task: Set up server monitoring > Task: Set up server monitoring
> Status: executing > Status: executing
> Phase: 1 > Phase: 1
## Phase 1: Prometheus Install ## Phase 1: Prometheus Install
- [x] Install Prometheus - [x] Install Prometheus
- [x] Configure node_exporter - [x] Configure node_exporter
## Phase 2: Grafana Setup ## Phase 2: Grafana Setup
- [ ] Install Grafana - [ ] Install Grafana
- [ ] Create dashboard - [ ] Create dashboard
## Context ## Context
Pi: Debian Bookworm arm64 Pi: Debian Bookworm arm64
` `
const testPlanAllComplete = `# Active Plan const testPlanAllComplete = `# Active Plan
> Task: Set up server monitoring > Task: Set up server monitoring
> Status: executing > Status: executing
> Phase: 2 > Phase: 2
## Phase 1: Prometheus Install ## Phase 1: Prometheus Install
- [x] Install Prometheus - [x] Install Prometheus
- [x] Configure node_exporter - [x] Configure node_exporter
## Phase 2: Grafana Setup ## Phase 2: Grafana Setup
- [x] Install Grafana - [x] Install Grafana
- [x] Create dashboard - [x] Create dashboard
## Context ## Context
Pi: Debian Bookworm arm64 Pi: Debian Bookworm arm64
` `
func TestHasActivePlan(t *testing.T) { func TestHasActivePlan(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// No plan // No plan
if ms.HasActivePlan() { if ms.HasActivePlan() {
t.Error("expected no active plan for empty memory") t.Error("expected no active plan for empty memory")
} }
// With regular content // With regular content
ms.WriteLongTerm("Some random notes") ms.WriteLongTerm("Some random notes")
if ms.HasActivePlan() { if ms.HasActivePlan() {
t.Error("expected no active plan for regular content") t.Error("expected no active plan for regular content")
} }
// With active plan // With active plan
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
if !ms.HasActivePlan() { if !ms.HasActivePlan() {
t.Error("expected active plan to be detected") t.Error("expected active plan to be detected")
} }
@ -111,21 +182,27 @@ func TestHasActivePlan(t *testing.T) {
func TestGetPlanStatus(t *testing.T) { func TestGetPlanStatus(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// No plan // No plan
if status := ms.GetPlanStatus(); status != "" { if status := ms.GetPlanStatus(); status != "" {
t.Errorf("expected empty status, got %q", status) t.Errorf("expected empty status, got %q", status)
} }
// Interviewing // Interviewing
ms.WriteLongTerm(testPlanInterviewing) ms.WriteLongTerm(testPlanInterviewing)
if status := ms.GetPlanStatus(); status != "interviewing" { if status := ms.GetPlanStatus(); status != "interviewing" {
t.Errorf("expected 'interviewing', got %q", status) t.Errorf("expected 'interviewing', got %q", status)
} }
// Executing // Executing
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
if status := ms.GetPlanStatus(); status != "executing" { if status := ms.GetPlanStatus(); status != "executing" {
t.Errorf("expected 'executing', got %q", status) t.Errorf("expected 'executing', got %q", status)
} }
@ -133,15 +210,19 @@ func TestGetPlanStatus(t *testing.T) {
func TestGetCurrentPhase(t *testing.T) { func TestGetCurrentPhase(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// No plan // No plan
if phase := ms.GetCurrentPhase(); phase != 0 { if phase := ms.GetCurrentPhase(); phase != 0 {
t.Errorf("expected phase 0, got %d", phase) t.Errorf("expected phase 0, got %d", phase)
} }
// Phase 2 // Phase 2
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
if phase := ms.GetCurrentPhase(); phase != 2 { if phase := ms.GetCurrentPhase(); phase != 2 {
t.Errorf("expected phase 2, got %d", phase) t.Errorf("expected phase 2, got %d", phase)
} }
@ -149,15 +230,19 @@ func TestGetCurrentPhase(t *testing.T) {
func TestGetTotalPhases(t *testing.T) { func TestGetTotalPhases(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// No plan // No plan
if total := ms.GetTotalPhases(); total != 0 { if total := ms.GetTotalPhases(); total != 0 {
t.Errorf("expected 0 phases, got %d", total) t.Errorf("expected 0 phases, got %d", total)
} }
// 3 phases // 3 phases
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
if total := ms.GetTotalPhases(); total != 3 { if total := ms.GetTotalPhases(); total != 3 {
t.Errorf("expected 3 phases, got %d", total) t.Errorf("expected 3 phases, got %d", total)
} }
@ -165,22 +250,29 @@ func TestGetTotalPhases(t *testing.T) {
func TestIsPlanComplete(t *testing.T) { func TestIsPlanComplete(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// Not complete // Not complete
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
if ms.IsPlanComplete() { if ms.IsPlanComplete() {
t.Error("expected plan to be incomplete") t.Error("expected plan to be incomplete")
} }
// All complete // All complete
ms.WriteLongTerm(testPlanAllComplete) ms.WriteLongTerm(testPlanAllComplete)
if !ms.IsPlanComplete() { if !ms.IsPlanComplete() {
t.Error("expected plan to be complete") t.Error("expected plan to be complete")
} }
// No plan // No plan
ms.ClearLongTerm() ms.ClearLongTerm()
if ms.IsPlanComplete() { if ms.IsPlanComplete() {
t.Error("expected false when no plan exists") t.Error("expected false when no plan exists")
} }
@ -188,16 +280,21 @@ func TestIsPlanComplete(t *testing.T) {
func TestIsCurrentPhaseComplete(t *testing.T) { func TestIsCurrentPhaseComplete(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// Phase 2 not complete // Phase 2 not complete
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
if ms.IsCurrentPhaseComplete() { if ms.IsCurrentPhaseComplete() {
t.Error("expected current phase to be incomplete") t.Error("expected current phase to be incomplete")
} }
// Phase 1 complete (current=1) // Phase 1 complete (current=1)
ms.WriteLongTerm(testPlanPhase1Complete) ms.WriteLongTerm(testPlanPhase1Complete)
if !ms.IsCurrentPhaseComplete() { if !ms.IsCurrentPhaseComplete() {
t.Error("expected phase 1 to be complete") t.Error("expected phase 1 to be complete")
} }
@ -205,12 +302,15 @@ func TestIsCurrentPhaseComplete(t *testing.T) {
func TestSetStatus(t *testing.T) { func TestSetStatus(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanInterviewing) ms.WriteLongTerm(testPlanInterviewing)
if err := ms.SetStatus("executing"); err != nil { if err := ms.SetStatus("executing"); err != nil {
t.Fatalf("SetStatus failed: %v", err) t.Fatalf("SetStatus failed: %v", err)
} }
if status := ms.GetPlanStatus(); status != "executing" { if status := ms.GetPlanStatus(); status != "executing" {
t.Errorf("expected 'executing', got %q", status) t.Errorf("expected 'executing', got %q", status)
} }
@ -218,12 +318,15 @@ func TestSetStatus(t *testing.T) {
func TestAdvancePhase(t *testing.T) { func TestAdvancePhase(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanPhase1Complete) ms.WriteLongTerm(testPlanPhase1Complete)
if err := ms.AdvancePhase(); err != nil { if err := ms.AdvancePhase(); err != nil {
t.Fatalf("AdvancePhase failed: %v", err) t.Fatalf("AdvancePhase failed: %v", err)
} }
if phase := ms.GetCurrentPhase(); phase != 2 { if phase := ms.GetCurrentPhase(); phase != 2 {
t.Errorf("expected phase 2 after advance, got %d", phase) t.Errorf("expected phase 2 after advance, got %d", phase)
} }
@ -231,37 +334,49 @@ func TestAdvancePhase(t *testing.T) {
func TestMarkStep(t *testing.T) { func TestMarkStep(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
// Mark step 1 in phase 2 // Mark step 1 in phase 2
if err := ms.MarkStep(2, 1); err != nil { if err := ms.MarkStep(2, 1); err != nil {
t.Fatalf("MarkStep failed: %v", err) t.Fatalf("MarkStep failed: %v", err)
} }
content := ms.ReadLongTerm() content := ms.ReadLongTerm()
// Phase 2 should have first step checked // Phase 2 should have first step checked
lines := strings.Split(content, "\n") lines := strings.Split(content, "\n")
foundChecked := false foundChecked := false
inPhase2 := false inPhase2 := false
for _, line := range lines { for _, line := range lines {
if strings.HasPrefix(line, "## Phase 2:") { if strings.HasPrefix(line, "## Phase 2:") {
inPhase2 = true inPhase2 = true
continue continue
} }
if inPhase2 && strings.HasPrefix(line, "## ") { if inPhase2 && strings.HasPrefix(line, "## ") {
break break
} }
if inPhase2 && strings.HasPrefix(line, "- [x] Install Grafana") { if inPhase2 && strings.HasPrefix(line, "- [x] Install Grafana") {
foundChecked = true foundChecked = true
} }
} }
if !foundChecked { if !foundChecked {
t.Error("expected 'Install Grafana' to be marked [x]") t.Error("expected 'Install Grafana' to be marked [x]")
} }
// Error case: invalid step // Error case: invalid step
if err := ms.MarkStep(2, 99); err == nil { if err := ms.MarkStep(2, 99); err == nil {
t.Error("expected error for invalid step number") t.Error("expected error for invalid step number")
} }
@ -269,23 +384,29 @@ func TestMarkStep(t *testing.T) {
func TestAddStep(t *testing.T) { func TestAddStep(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
// Add step to phase 2 // Add step to phase 2
if err := ms.AddStep(2, "Test dashboard"); err != nil { if err := ms.AddStep(2, "Test dashboard"); err != nil {
t.Fatalf("AddStep failed: %v", err) t.Fatalf("AddStep failed: %v", err)
} }
content := ms.ReadLongTerm() content := ms.ReadLongTerm()
if !strings.Contains(content, "- [ ] Test dashboard") { if !strings.Contains(content, "- [ ] Test dashboard") {
t.Error("expected new step to be added") t.Error("expected new step to be added")
} }
// Verify it's in the right place (before Phase 3) // Verify it's in the right place (before Phase 3)
idx := strings.Index(content, "- [ ] Test dashboard") idx := strings.Index(content, "- [ ] Test dashboard")
phase3Idx := strings.Index(content, "## Phase 3:") phase3Idx := strings.Index(content, "## Phase 3:")
if idx > phase3Idx { if idx > phase3Idx {
t.Error("expected new step to be before Phase 3") t.Error("expected new step to be before Phase 3")
} }
@ -293,17 +414,21 @@ func TestAddStep(t *testing.T) {
func TestClearLongTerm(t *testing.T) { func TestClearLongTerm(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
if err := ms.ClearLongTerm(); err != nil { if err := ms.ClearLongTerm(); err != nil {
t.Fatalf("ClearLongTerm failed: %v", err) t.Fatalf("ClearLongTerm failed: %v", err)
} }
if content := ms.ReadLongTerm(); content != "" { if content := ms.ReadLongTerm(); content != "" {
t.Errorf("expected empty memory after clear, got %q", content) t.Errorf("expected empty memory after clear, got %q", content)
} }
// Clearing again should not error // Clearing again should not error
if err := ms.ClearLongTerm(); err != nil { if err := ms.ClearLongTerm(); err != nil {
t.Fatalf("ClearLongTerm (idempotent) failed: %v", err) t.Fatalf("ClearLongTerm (idempotent) failed: %v", err)
} }
@ -311,34 +436,45 @@ func TestClearLongTerm(t *testing.T) {
func TestGetInterviewContext(t *testing.T) { func TestGetInterviewContext(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanInterviewing) ms.WriteLongTerm(testPlanInterviewing)
ctx := ms.GetInterviewContext() ctx := ms.GetInterviewContext()
if !strings.Contains(ctx, "Active Plan (interviewing)") { if !strings.Contains(ctx, "Active Plan (interviewing)") {
t.Error("expected 'Active Plan (interviewing)' header") t.Error("expected 'Active Plan (interviewing)' header")
} }
if !strings.Contains(ctx, "Interview Guide") { if !strings.Contains(ctx, "Interview Guide") {
t.Error("expected 'Interview Guide' section") t.Error("expected 'Interview Guide' section")
} }
if !strings.Contains(ctx, "Target Format") { if !strings.Contains(ctx, "Target Format") {
t.Error("expected 'Target Format' section") t.Error("expected 'Target Format' section")
} }
if !strings.Contains(ctx, "Set up server monitoring") { if !strings.Contains(ctx, "Set up server monitoring") {
t.Error("expected task description in context") t.Error("expected task description in context")
} }
// Should guide AI to ask about tooling // Should guide AI to ask about tooling
if !strings.Contains(ctx, "test framework") || !strings.Contains(ctx, "linter") { if !strings.Contains(ctx, "test framework") || !strings.Contains(ctx, "linter") {
t.Error("expected interview guide to mention test framework and linter") t.Error("expected interview guide to mention test framework and linter")
} }
// Target format should include Commands section example // Target format should include Commands section example
if !strings.Contains(ctx, "## Commands") { if !strings.Contains(ctx, "## Commands") {
t.Error("expected target format to include ## Commands section") t.Error("expected target format to include ## Commands section")
} }
if !strings.Contains(ctx, "project-specific test command") { if !strings.Contains(ctx, "project-specific test command") {
t.Error("expected target format Commands to include test command placeholder") t.Error("expected target format Commands to include test command placeholder")
} }
if !strings.Contains(ctx, "project-specific lint command") { if !strings.Contains(ctx, "project-specific lint command") {
t.Error("expected target format Commands to include lint command placeholder") 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) { func TestGetPlanContext(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
ctx := ms.GetPlanContext() ctx := ms.GetPlanContext()
// Should have task summary // Should have task summary
if !strings.Contains(ctx, "Phase 2/3") { if !strings.Contains(ctx, "Phase 2/3") {
t.Error("expected 'Phase 2/3' in plan context") t.Error("expected 'Phase 2/3' in plan context")
} }
// Completed phase should be summarized // Completed phase should be summarized
if !strings.Contains(ctx, "Done: Phase 1") { if !strings.Contains(ctx, "Done: Phase 1") {
t.Error("expected completed phase summary") t.Error("expected completed phase summary")
} }
// Current phase should have full detail // Current phase should have full detail
if !strings.Contains(ctx, "Current: Phase 2") { if !strings.Contains(ctx, "Current: Phase 2") {
t.Error("expected current phase detail") t.Error("expected current phase detail")
} }
if !strings.Contains(ctx, "Install Grafana") { if !strings.Contains(ctx, "Install Grafana") {
t.Error("expected current phase steps") t.Error("expected current phase steps")
} }
// Future phases should NOT appear // Future phases should NOT appear
if strings.Contains(ctx, "Phase 3") { if strings.Contains(ctx, "Phase 3") {
t.Error("expected future phases to be omitted") t.Error("expected future phases to be omitted")
} }
// Commands should be included // Commands should be included
if !strings.Contains(ctx, "### Commands") { if !strings.Contains(ctx, "### Commands") {
t.Error("expected Commands section in plan context") t.Error("expected Commands section in plan context")
} }
if !strings.Contains(ctx, "go test") { if !strings.Contains(ctx, "go test") {
t.Error("expected test command in Commands section") t.Error("expected test command in Commands section")
} }
if !strings.Contains(ctx, "golangci-lint") { if !strings.Contains(ctx, "golangci-lint") {
t.Error("expected lint command in Commands section") t.Error("expected lint command in Commands section")
} }
// Context should be included // Context should be included
if !strings.Contains(ctx, "Debian Bookworm") { if !strings.Contains(ctx, "Debian Bookworm") {
t.Error("expected Context section") t.Error("expected Context section")
} }
@ -393,20 +540,27 @@ func TestGetPlanContext(t *testing.T) {
func TestGetMemoryContext_PlanActive_SuppressesDailyNotes(t *testing.T) { func TestGetMemoryContext_PlanActive_SuppressesDailyNotes(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// Write a daily note // Write a daily note
ms.AppendToday("Today's note") ms.AppendToday("Today's note")
// Without plan, daily notes should appear // Without plan, daily notes should appear
ctx := ms.GetMemoryContext() ctx := ms.GetMemoryContext()
if !strings.Contains(ctx, "Recent Daily Notes") { if !strings.Contains(ctx, "Recent Daily Notes") {
t.Error("expected daily notes when no plan active") t.Error("expected daily notes when no plan active")
} }
// With plan, daily notes should be suppressed // With plan, daily notes should be suppressed
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
ctx = ms.GetMemoryContext() ctx = ms.GetMemoryContext()
if strings.Contains(ctx, "Recent Daily Notes") { if strings.Contains(ctx, "Recent Daily Notes") {
t.Error("expected daily notes to be suppressed when plan is active") 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) { func TestGetMemoryContext_InterviewingMode(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanInterviewing) ms.WriteLongTerm(testPlanInterviewing)
ctx := ms.GetMemoryContext() ctx := ms.GetMemoryContext()
if !strings.Contains(ctx, "interviewing") { if !strings.Contains(ctx, "interviewing") {
t.Error("expected interviewing context") t.Error("expected interviewing context")
} }
if !strings.Contains(ctx, "Interview Guide") { if !strings.Contains(ctx, "Interview Guide") {
t.Error("expected interview guide in context") t.Error("expected interview guide in context")
} }
@ -429,14 +586,17 @@ func TestGetMemoryContext_InterviewingMode(t *testing.T) {
func TestGetMemoryContext_ExecutingMode(t *testing.T) { func TestGetMemoryContext_ExecutingMode(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
ctx := ms.GetMemoryContext() ctx := ms.GetMemoryContext()
if !strings.Contains(ctx, "Active Plan") { if !strings.Contains(ctx, "Active Plan") {
t.Error("expected active plan in context") t.Error("expected active plan in context")
} }
if !strings.Contains(ctx, "Current: Phase 2") { if !strings.Contains(ctx, "Current: Phase 2") {
t.Error("expected current phase in context") t.Error("expected current phase in context")
} }
@ -444,14 +604,17 @@ func TestGetMemoryContext_ExecutingMode(t *testing.T) {
func TestGetMemoryContext_RegularMemory(t *testing.T) { func TestGetMemoryContext_RegularMemory(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm("Some notes about projects") ms.WriteLongTerm("Some notes about projects")
ctx := ms.GetMemoryContext() ctx := ms.GetMemoryContext()
if !strings.Contains(ctx, "Long-term Memory") { if !strings.Contains(ctx, "Long-term Memory") {
t.Error("expected regular long-term memory section") t.Error("expected regular long-term memory section")
} }
if !strings.Contains(ctx, "Some notes about projects") { if !strings.Contains(ctx, "Some notes about projects") {
t.Error("expected memory content") t.Error("expected memory content")
} }
@ -463,12 +626,15 @@ func TestBuildInterviewSeed(t *testing.T) {
if !strings.Contains(seed, "# Active Plan") { if !strings.Contains(seed, "# Active Plan") {
t.Error("expected '# Active Plan' header") t.Error("expected '# Active Plan' header")
} }
if !strings.Contains(seed, "Deploy monitoring stack") { if !strings.Contains(seed, "Deploy monitoring stack") {
t.Error("expected task description") t.Error("expected task description")
} }
if !strings.Contains(seed, "interviewing") { if !strings.Contains(seed, "interviewing") {
t.Error("expected interviewing status") t.Error("expected interviewing status")
} }
if !strings.Contains(seed, "> Phase: 1") { if !strings.Contains(seed, "> Phase: 1") {
t.Error("expected Phase: 1") t.Error("expected Phase: 1")
} }
@ -476,28 +642,37 @@ func TestBuildInterviewSeed(t *testing.T) {
func TestFormatPlanDisplay(t *testing.T) { func TestFormatPlanDisplay(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
// No plan // No plan
display := ms.FormatPlanDisplay() display := ms.FormatPlanDisplay()
if display != "No active plan." { if display != "No active plan." {
t.Errorf("expected 'No active plan.', got %q", display) t.Errorf("expected 'No active plan.', got %q", display)
} }
// With plan // With plan
ms.WriteLongTerm(testPlanExecuting) ms.WriteLongTerm(testPlanExecuting)
display = ms.FormatPlanDisplay() display = ms.FormatPlanDisplay()
if !strings.Contains(display, "Set up server monitoring") { if !strings.Contains(display, "Set up server monitoring") {
t.Error("expected task name in display") t.Error("expected task name in display")
} }
if !strings.Contains(display, "Phase 2/3") { if !strings.Contains(display, "Phase 2/3") {
t.Error("expected phase count in display") t.Error("expected phase count in display")
} }
// Commands section should be visible // Commands section should be visible
if !strings.Contains(display, "Commands:") { if !strings.Contains(display, "Commands:") {
t.Error("expected Commands section in display") t.Error("expected Commands section in display")
} }
if !strings.Contains(display, "go test") { if !strings.Contains(display, "go test") {
t.Error("expected test command in display") t.Error("expected test command in display")
} }
@ -506,117 +681,210 @@ func TestFormatPlanDisplay(t *testing.T) {
func TestValidatePlanStructure(t *testing.T) { func TestValidatePlanStructure(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
content string content string
wantErr string // "" means nil error expected wantErr string // "" means nil error expected
}{ }{
{ {
name: "valid plan with 1 phase and 1 step", name: "valid plan with 1 phase and 1 step",
content: `# Active Plan content: `# Active Plan
> Task: Do something > Task: Do something
> Status: executing > Status: executing
> Phase: 1 > Phase: 1
## Phase 1: Setup ## Phase 1: Setup
- [ ] Install deps - [ ] Install deps
`, `,
wantErr: "", wantErr: "",
}, },
{ {
name: "missing Active Plan header", name: "missing Active Plan header",
content: `> Status: executing`, content: `> Status: executing`,
wantErr: "missing '# Active Plan' header", wantErr: "missing '# Active Plan' header",
}, },
{ {
name: "missing Status line", name: "missing Status line",
content: `# Active Plan content: `# Active Plan
> Phase: 1 > Phase: 1
## Phase 1: Setup ## Phase 1: Setup
- [ ] Install deps - [ ] Install deps
`, `,
wantErr: "missing '> Status:' line", wantErr: "missing '> Status:' line",
}, },
{ {
name: "missing Phase line", name: "missing Phase line",
content: `# Active Plan content: `# Active Plan
> Status: executing > Status: executing
## Phase 1: Setup ## Phase 1: Setup
- [ ] Install deps - [ ] Install deps
`, `,
wantErr: "missing '> Phase:' line", wantErr: "missing '> Phase:' line",
}, },
{ {
name: "no Phase sections", name: "no Phase sections",
content: `# Active Plan content: `# Active Plan
> Task: Do something > Task: Do something
> Status: executing > Status: executing
> Phase: 1 > Phase: 1
`, `,
wantErr: "no '## Phase N:' sections found", wantErr: "no '## Phase N:' sections found",
}, },
{ {
name: "phase with no checkbox steps", name: "phase with no checkbox steps",
content: `# Active Plan content: `# Active Plan
> Task: Do something > Task: Do something
> Status: executing > Status: executing
> Phase: 1 > Phase: 1
## Phase 1: Setup ## Phase 1: Setup
Some description without checkboxes Some description without checkboxes
`, `,
wantErr: "Phase 1 has no checkbox steps", wantErr: "Phase 1 has no checkbox steps",
}, },
{ {
name: "all steps done is valid", name: "all steps done is valid",
content: `# Active Plan content: `# Active Plan
> Task: Do something > Task: Do something
> Status: executing > Status: executing
> Phase: 1 > Phase: 1
## Phase 1: Setup ## Phase 1: Setup
- [x] Install deps - [x] Install deps
- [x] Configure - [x] Configure
`, `,
wantErr: "", wantErr: "",
}, },
{ {
name: "multi-phase valid", name: "multi-phase valid",
content: `# Active Plan content: `# Active Plan
> Task: Do something > Task: Do something
> Status: executing > Status: executing
> Phase: 1 > Phase: 1
## Phase 1: Setup ## Phase 1: Setup
- [ ] Install deps - [ ] Install deps
## Phase 2: Build ## Phase 2: Build
- [ ] Compile - [ ] Compile
- [ ] Test - [ ] Test
`, `,
wantErr: "", wantErr: "",
}, },
{ {
name: "second phase empty steps", name: "second phase empty steps",
content: `# Active Plan content: `# Active Plan
> Task: Do something > Task: Do something
> Status: executing > Status: executing
> Phase: 1 > Phase: 1
## Phase 1: Setup ## Phase 1: Setup
- [ ] Install deps - [ ] Install deps
## Phase 2: Build ## Phase 2: Build
No checkboxes here No checkboxes here
`, `,
wantErr: "Phase 2 has no checkbox steps", wantErr: "Phase 2 has no checkbox steps",
}, },
} }
@ -624,9 +892,11 @@ No checkboxes here
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
ms, cleanup := newTestMemoryStore(t) ms, cleanup := newTestMemoryStore(t)
defer cleanup() defer cleanup()
ms.WriteLongTerm(tt.content) ms.WriteLongTerm(tt.content)
err := ms.ValidatePlanStructure() err := ms.ValidatePlanStructure()
if tt.wantErr == "" { if tt.wantErr == "" {
@ -649,18 +919,23 @@ func TestMemoryStoreCreation(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("Failed to create temp dir: %v", err) t.Fatalf("Failed to create temp dir: %v", err)
} }
defer os.RemoveAll(tmpDir) defer os.RemoveAll(tmpDir)
ms := NewMemoryStore(tmpDir) ms := NewMemoryStore(tmpDir)
// Verify memory directory was created // Verify memory directory was created
memoryDir := filepath.Join(tmpDir, "memory") memoryDir := filepath.Join(tmpDir, "memory")
if _, err := os.Stat(memoryDir); os.IsNotExist(err) { if _, err := os.Stat(memoryDir); os.IsNotExist(err) {
t.Error("expected memory directory to be created") t.Error("expected memory directory to be created")
} }
// Verify memory file path // Verify memory file path
expectedFile := filepath.Join(memoryDir, "MEMORY.md") expectedFile := filepath.Join(memoryDir, "MEMORY.md")
if ms.memoryFile != expectedFile { if ms.memoryFile != expectedFile {
t.Errorf("expected memory file %q, got %q", expectedFile, ms.memoryFile) t.Errorf("expected memory file %q, got %q", expectedFile, ms.memoryFile)
} }

View file

@ -10,13 +10,18 @@ type mockProvider struct{}
func (m *mockProvider) Chat( func (m *mockProvider) Chat(
ctx context.Context, ctx context.Context,
messages []providers.Message, messages []providers.Message,
tools []providers.ToolDefinition, tools []providers.ToolDefinition,
model string, model string,
opts map[string]any, opts map[string]any,
) (*providers.LLMResponse, error) { ) (*providers.LLMResponse, error) {
return &providers.LLMResponse{ return &providers.LLMResponse{
Content: "Mock response", Content: "Mock response",
ToolCalls: []providers.ToolCall{}, ToolCalls: []providers.ToolCall{},
}, nil }, nil
} }

View file

@ -10,42 +10,61 @@ import (
) )
// AgentRegistry manages multiple agent instances and routes messages to them. // AgentRegistry manages multiple agent instances and routes messages to them.
type AgentRegistry struct { type AgentRegistry struct {
agents map[string]*AgentInstance agents map[string]*AgentInstance
resolver *routing.RouteResolver resolver *routing.RouteResolver
mu sync.RWMutex mu sync.RWMutex
} }
// NewAgentRegistry creates a registry from config, instantiating all agents. // NewAgentRegistry creates a registry from config, instantiating all agents.
func NewAgentRegistry( func NewAgentRegistry(
cfg *config.Config, cfg *config.Config,
provider providers.LLMProvider, provider providers.LLMProvider,
) *AgentRegistry { ) *AgentRegistry {
registry := &AgentRegistry{ registry := &AgentRegistry{
agents: make(map[string]*AgentInstance), agents: make(map[string]*AgentInstance),
resolver: routing.NewRouteResolver(cfg), resolver: routing.NewRouteResolver(cfg),
} }
agentConfigs := cfg.Agents.List agentConfigs := cfg.Agents.List
if len(agentConfigs) == 0 { if len(agentConfigs) == 0 {
implicitAgent := &config.AgentConfig{ implicitAgent := &config.AgentConfig{
ID: "main", ID: "main",
Default: true, Default: true,
} }
instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider) instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider)
registry.agents["main"] = instance registry.agents["main"] = instance
logger.InfoCF("agent", "Created implicit main agent (no agents.list configured)", nil) logger.InfoCF("agent", "Created implicit main agent (no agents.list configured)", nil)
} else { } else {
for i := range agentConfigs { for i := range agentConfigs {
ac := &agentConfigs[i] ac := &agentConfigs[i]
id := routing.NormalizeAgentID(ac.ID) id := routing.NormalizeAgentID(ac.ID)
instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider) instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider)
registry.agents[id] = instance registry.agents[id] = instance
logger.InfoCF("agent", "Registered agent", logger.InfoCF("agent", "Registered agent",
map[string]any{ map[string]any{
"agent_id": id, "agent_id": id,
"name": ac.Name, "name": ac.Name,
"workspace": instance.Workspace, "workspace": instance.Workspace,
"model": instance.Model, "model": instance.Model,
}) })
} }
@ -55,60 +74,83 @@ func NewAgentRegistry(
} }
// GetAgent returns the agent instance for a given ID. // GetAgent returns the agent instance for a given ID.
func (r *AgentRegistry) GetAgent(agentID string) (*AgentInstance, bool) { func (r *AgentRegistry) GetAgent(agentID string) (*AgentInstance, bool) {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
id := routing.NormalizeAgentID(agentID) id := routing.NormalizeAgentID(agentID)
agent, ok := r.agents[id] agent, ok := r.agents[id]
return agent, ok return agent, ok
} }
// ResolveRoute determines which agent handles the message. // ResolveRoute determines which agent handles the message.
func (r *AgentRegistry) ResolveRoute(input routing.RouteInput) routing.ResolvedRoute { func (r *AgentRegistry) ResolveRoute(input routing.RouteInput) routing.ResolvedRoute {
return r.resolver.ResolveRoute(input) return r.resolver.ResolveRoute(input)
} }
// ListAgentIDs returns all registered agent IDs. // ListAgentIDs returns all registered agent IDs.
func (r *AgentRegistry) ListAgentIDs() []string { func (r *AgentRegistry) ListAgentIDs() []string {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
ids := make([]string, 0, len(r.agents)) ids := make([]string, 0, len(r.agents))
for id := range r.agents { for id := range r.agents {
ids = append(ids, id) ids = append(ids, id)
} }
return ids return ids
} }
// CanSpawnSubagent checks if parentAgentID is allowed to spawn targetAgentID. // CanSpawnSubagent checks if parentAgentID is allowed to spawn targetAgentID.
func (r *AgentRegistry) CanSpawnSubagent(parentAgentID, targetAgentID string) bool { func (r *AgentRegistry) CanSpawnSubagent(parentAgentID, targetAgentID string) bool {
parent, ok := r.GetAgent(parentAgentID) parent, ok := r.GetAgent(parentAgentID)
if !ok { if !ok {
return false return false
} }
if parent.Subagents == nil || parent.Subagents.AllowAgents == nil { if parent.Subagents == nil || parent.Subagents.AllowAgents == nil {
return false return false
} }
targetNorm := routing.NormalizeAgentID(targetAgentID) targetNorm := routing.NormalizeAgentID(targetAgentID)
for _, allowed := range parent.Subagents.AllowAgents { for _, allowed := range parent.Subagents.AllowAgents {
if allowed == "*" { if allowed == "*" {
return true return true
} }
if routing.NormalizeAgentID(allowed) == targetNorm { if routing.NormalizeAgentID(allowed) == targetNorm {
return true return true
} }
} }
return false return false
} }
// GetDefaultAgent returns the default agent instance. // GetDefaultAgent returns the default agent instance.
func (r *AgentRegistry) GetDefaultAgent() *AgentInstance { func (r *AgentRegistry) GetDefaultAgent() *AgentInstance {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
if agent, ok := r.agents["main"]; ok { if agent, ok := r.agents["main"]; ok {
return agent return agent
} }
for _, agent := range r.agents { for _, agent := range r.agents {
return agent return agent
} }
return nil return nil
} }

View file

@ -12,9 +12,13 @@ type mockRegistryProvider struct{}
func (m *mockRegistryProvider) Chat( func (m *mockRegistryProvider) Chat(
ctx context.Context, ctx context.Context,
messages []providers.Message, messages []providers.Message,
tools []providers.ToolDefinition, tools []providers.ToolDefinition,
model string, model string,
options map[string]any, options map[string]any,
) (*providers.LLMResponse, error) { ) (*providers.LLMResponse, error) {
return &providers.LLMResponse{Content: "mock", FinishReason: "stop"}, nil return &providers.LLMResponse{Content: "mock", FinishReason: "stop"}, nil
@ -29,10 +33,14 @@ func testCfg(agents []config.AgentConfig) *config.Config {
Agents: config.AgentsConfig{ Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{ Defaults: config.AgentDefaults{
Workspace: "/tmp/picoclaw-test-registry", Workspace: "/tmp/picoclaw-test-registry",
Model: "gpt-4", Model: "gpt-4",
MaxTokens: 8192, MaxTokens: 8192,
MaxToolIterations: 10, MaxToolIterations: 10,
}, },
List: agents, List: agents,
}, },
} }
@ -40,17 +48,21 @@ func testCfg(agents []config.AgentConfig) *config.Config {
func TestNewAgentRegistry_ImplicitMain(t *testing.T) { func TestNewAgentRegistry_ImplicitMain(t *testing.T) {
cfg := testCfg(nil) cfg := testCfg(nil)
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
ids := registry.ListAgentIDs() ids := registry.ListAgentIDs()
if len(ids) != 1 || ids[0] != "main" { if len(ids) != 1 || ids[0] != "main" {
t.Errorf("expected implicit main agent, got %v", ids) t.Errorf("expected implicit main agent, got %v", ids)
} }
agent, ok := registry.GetAgent("main") agent, ok := registry.GetAgent("main")
if !ok || agent == nil { if !ok || agent == nil {
t.Fatal("expected to find 'main' agent") t.Fatal("expected to find 'main' agent")
} }
if agent.ID != "main" { if agent.ID != "main" {
t.Errorf("agent.ID = %q, want 'main'", agent.ID) 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) { func TestNewAgentRegistry_ExplicitAgents(t *testing.T) {
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ID: "sales", Default: true, Name: "Sales Bot"}, {ID: "sales", Default: true, Name: "Sales Bot"},
{ID: "support", Name: "Support Bot"}, {ID: "support", Name: "Support Bot"},
}) })
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
ids := registry.ListAgentIDs() ids := registry.ListAgentIDs()
if len(ids) != 2 { if len(ids) != 2 {
t.Fatalf("expected 2 agents, got %d: %v", len(ids), ids) t.Fatalf("expected 2 agents, got %d: %v", len(ids), ids)
} }
sales, ok := registry.GetAgent("sales") sales, ok := registry.GetAgent("sales")
if !ok || sales == nil { if !ok || sales == nil {
t.Fatal("expected to find 'sales' agent") t.Fatal("expected to find 'sales' agent")
} }
if sales.Name != "Sales Bot" { if sales.Name != "Sales Bot" {
t.Errorf("sales.Name = %q, want 'Sales Bot'", sales.Name) t.Errorf("sales.Name = %q, want 'Sales Bot'", sales.Name)
} }
support, ok := registry.GetAgent("support") support, ok := registry.GetAgent("support")
if !ok || support == nil { if !ok || support == nil {
t.Fatal("expected to find 'support' agent") t.Fatal("expected to find 'support' agent")
} }
@ -86,12 +104,15 @@ func TestAgentRegistry_GetAgent_Normalize(t *testing.T) {
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ID: "my-agent", Default: true}, {ID: "my-agent", Default: true},
}) })
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
agent, ok := registry.GetAgent("My-Agent") agent, ok := registry.GetAgent("My-Agent")
if !ok || agent == nil { if !ok || agent == nil {
t.Fatal("expected to find agent with normalized ID") t.Fatal("expected to find agent with normalized ID")
} }
if agent.ID != "my-agent" { if agent.ID != "my-agent" {
t.Errorf("agent.ID = %q, want 'my-agent'", agent.ID) 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) { func TestAgentRegistry_GetDefaultAgent(t *testing.T) {
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ID: "alpha"}, {ID: "alpha"},
{ID: "beta", Default: true}, {ID: "beta", Default: true},
}) })
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
// GetDefaultAgent first checks for "main", then returns any // GetDefaultAgent first checks for "main", then returns any
agent := registry.GetDefaultAgent() agent := registry.GetDefaultAgent()
if agent == nil { if agent == nil {
t.Fatal("expected a default agent") t.Fatal("expected a default agent")
} }
@ -115,26 +140,35 @@ func TestAgentRegistry_CanSpawnSubagent(t *testing.T) {
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ {
ID: "parent", ID: "parent",
Default: true, Default: true,
Subagents: &config.SubagentsConfig{ Subagents: &config.SubagentsConfig{
AllowAgents: []string{"child1", "child2"}, AllowAgents: []string{"child1", "child2"},
}, },
}, },
{ID: "child1"}, {ID: "child1"},
{ID: "child2"}, {ID: "child2"},
{ID: "restricted"}, {ID: "restricted"},
}) })
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
if !registry.CanSpawnSubagent("parent", "child1") { if !registry.CanSpawnSubagent("parent", "child1") {
t.Error("expected parent to be allowed to spawn child1") t.Error("expected parent to be allowed to spawn child1")
} }
if !registry.CanSpawnSubagent("parent", "child2") { if !registry.CanSpawnSubagent("parent", "child2") {
t.Error("expected parent to be allowed to spawn child2") t.Error("expected parent to be allowed to spawn child2")
} }
if registry.CanSpawnSubagent("parent", "restricted") { if registry.CanSpawnSubagent("parent", "restricted") {
t.Error("expected parent to NOT be allowed to spawn restricted") t.Error("expected parent to NOT be allowed to spawn restricted")
} }
if registry.CanSpawnSubagent("child1", "child2") { if registry.CanSpawnSubagent("child1", "child2") {
t.Error("expected child1 to NOT be allowed to spawn (no subagents config)") t.Error("expected child1 to NOT be allowed to spawn (no subagents config)")
} }
@ -144,18 +178,23 @@ func TestAgentRegistry_CanSpawnSubagent_Wildcard(t *testing.T) {
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ {
ID: "admin", ID: "admin",
Default: true, Default: true,
Subagents: &config.SubagentsConfig{ Subagents: &config.SubagentsConfig{
AllowAgents: []string{"*"}, AllowAgents: []string{"*"},
}, },
}, },
{ID: "any-agent"}, {ID: "any-agent"},
}) })
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
if !registry.CanSpawnSubagent("admin", "any-agent") { if !registry.CanSpawnSubagent("admin", "any-agent") {
t.Error("expected wildcard to allow spawning any agent") t.Error("expected wildcard to allow spawning any agent")
} }
if !registry.CanSpawnSubagent("admin", "nonexistent") { if !registry.CanSpawnSubagent("admin", "nonexistent") {
t.Error("expected wildcard to allow spawning even nonexistent agents") 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) { func TestAgentInstance_Model(t *testing.T) {
model := &config.AgentModelConfig{Primary: "claude-opus"} model := &config.AgentModelConfig{Primary: "claude-opus"}
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ID: "custom", Default: true, Model: model}, {ID: "custom", Default: true, Model: model},
}) })
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
agent, _ := registry.GetAgent("custom") agent, _ := registry.GetAgent("custom")
if agent.Model != "claude-opus" { if agent.Model != "claude-opus" {
t.Errorf("agent.Model = %q, want 'claude-opus'", agent.Model) t.Errorf("agent.Model = %q, want 'claude-opus'", agent.Model)
} }
@ -178,10 +220,13 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) {
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ID: "inherit", Default: true}, {ID: "inherit", Default: true},
}) })
cfg.Agents.Defaults.ModelFallbacks = []string{"openai/gpt-4o-mini", "anthropic/haiku"} cfg.Agents.Defaults.ModelFallbacks = []string{"openai/gpt-4o-mini", "anthropic/haiku"}
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
agent, _ := registry.GetAgent("inherit") agent, _ := registry.GetAgent("inherit")
if len(agent.Fallbacks) != 2 { if len(agent.Fallbacks) != 2 {
t.Errorf("expected 2 fallbacks inherited from defaults, got %d", len(agent.Fallbacks)) t.Errorf("expected 2 fallbacks inherited from defaults, got %d", len(agent.Fallbacks))
} }
@ -190,15 +235,21 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) {
func TestAgentInstance_FallbackExplicitEmpty(t *testing.T) { func TestAgentInstance_FallbackExplicitEmpty(t *testing.T) {
model := &config.AgentModelConfig{ model := &config.AgentModelConfig{
Primary: "gpt-4", Primary: "gpt-4",
Fallbacks: []string{}, // explicitly empty = disable Fallbacks: []string{}, // explicitly empty = disable
} }
cfg := testCfg([]config.AgentConfig{ cfg := testCfg([]config.AgentConfig{
{ID: "no-fallback", Default: true, Model: model}, {ID: "no-fallback", Default: true, Model: model},
}) })
cfg.Agents.Defaults.ModelFallbacks = []string{"should-not-inherit"} cfg.Agents.Defaults.ModelFallbacks = []string{"should-not-inherit"}
registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) registry := NewAgentRegistry(cfg, &mockRegistryProvider{})
agent, _ := registry.GetAgent("no-fallback") agent, _ := registry.GetAgent("no-fallback")
if len(agent.Fallbacks) != 0 { if len(agent.Fallbacks) != 0 {
t.Errorf("expected 0 fallbacks (explicit empty), got %d: %v", len(agent.Fallbacks), agent.Fallbacks) t.Errorf("expected 0 fallbacks (explicit empty), got %d: %v", len(agent.Fallbacks), agent.Fallbacks)
} }

View file

@ -8,38 +8,55 @@ import (
) )
// SessionEntry represents an active or recently-active session. // SessionEntry represents an active or recently-active session.
type SessionEntry struct { type SessionEntry struct {
SessionKey string `json:"session_key"` SessionKey string `json:"session_key"`
Channel string `json:"channel"` Channel string `json:"channel"`
ChatID string `json:"chat_id"` ChatID string `json:"chat_id"`
TouchDir string `json:"touch_dir"` TouchDir string `json:"touch_dir"`
ProjectPath string `json:"project_path,omitempty"` // canonical project path ProjectPath string `json:"project_path,omitempty"` // canonical project path
Purpose string `json:"purpose,omitempty"` // 1-line task description Purpose string `json:"purpose,omitempty"` // 1-line task description
Branch string `json:"branch,omitempty"` // git branch name Branch string `json:"branch,omitempty"` // git branch name
LastSeenAt time.Time `json:"last_seen_at"` LastSeenAt time.Time `json:"last_seen_at"`
} }
// TouchMeta carries optional metadata for Touch calls. // TouchMeta carries optional metadata for Touch calls.
type TouchMeta struct { type TouchMeta struct {
ProjectPath string // canonical project path (always original workspace-relative) ProjectPath string // canonical project path (always original workspace-relative)
Purpose string // 1-line task description Purpose string // 1-line task description
Branch string // git branch name Branch string // git branch name
} }
// PeerInfo is the minimal info shared between sessions on the same project. // PeerInfo is the minimal info shared between sessions on the same project.
type PeerInfo struct { type PeerInfo struct {
SessionKey string SessionKey string
Purpose string Purpose string
Branch string Branch string
} }
// SessionTracker tracks per-session tool-call activity. // SessionTracker tracks per-session tool-call activity.
// Thread-safe; used by AgentLoop for plan coordination and by the mini app API for observability. // Thread-safe; used by AgentLoop for plan coordination and by the mini app API for observability.
type SessionTracker struct { type SessionTracker struct {
entries sync.Map // sessionKey → *SessionEntry entries sync.Map // sessionKey → *SessionEntry
} }
// NewSessionTracker creates a new tracker. // NewSessionTracker creates a new tracker.
func NewSessionTracker() *SessionTracker { func NewSessionTracker() *SessionTracker {
return &SessionTracker{} return &SessionTracker{}
} }
@ -47,121 +64,177 @@ func NewSessionTracker() *SessionTracker {
const sessionActivityTimeout = 15 * time.Minute const sessionActivityTimeout = 15 * time.Minute
// Touch records a tool-call activity for a session. // Touch records a tool-call activity for a session.
// dir is the workspace-relative directory the tool call targeted. // dir is the workspace-relative directory the tool call targeted.
// If dir is empty, only LastSeenAt is updated. // If dir is empty, only LastSeenAt is updated.
// meta is optional and carries project coordination metadata. // meta is optional and carries project coordination metadata.
func (st *SessionTracker) Touch(sessionKey, channel, chatID, dir string, meta *TouchMeta) { func (st *SessionTracker) Touch(sessionKey, channel, chatID, dir string, meta *TouchMeta) {
now := time.Now() now := time.Now()
val, loaded := st.entries.Load(sessionKey) val, loaded := st.entries.Load(sessionKey)
if loaded { if loaded {
entry := val.(*SessionEntry) entry := val.(*SessionEntry)
entry.LastSeenAt = now entry.LastSeenAt = now
if dir != "" { if dir != "" {
entry.TouchDir = dir entry.TouchDir = dir
} }
if channel != "" { if channel != "" {
entry.Channel = channel entry.Channel = channel
} }
if chatID != "" { if chatID != "" {
entry.ChatID = chatID entry.ChatID = chatID
} }
if meta != nil { if meta != nil {
if meta.ProjectPath != "" { if meta.ProjectPath != "" {
entry.ProjectPath = meta.ProjectPath entry.ProjectPath = meta.ProjectPath
} }
if meta.Purpose != "" { if meta.Purpose != "" {
entry.Purpose = meta.Purpose entry.Purpose = meta.Purpose
} }
if meta.Branch != "" { if meta.Branch != "" {
entry.Branch = meta.Branch entry.Branch = meta.Branch
} }
} }
return return
} }
entry := &SessionEntry{ entry := &SessionEntry{
SessionKey: sessionKey, SessionKey: sessionKey,
Channel: channel, Channel: channel,
ChatID: chatID, ChatID: chatID,
TouchDir: dir, TouchDir: dir,
LastSeenAt: now, LastSeenAt: now,
} }
if meta != nil { if meta != nil {
entry.ProjectPath = meta.ProjectPath entry.ProjectPath = meta.ProjectPath
entry.Purpose = meta.Purpose entry.Purpose = meta.Purpose
entry.Branch = meta.Branch entry.Branch = meta.Branch
} }
st.entries.Store(sessionKey, entry) st.entries.Store(sessionKey, entry)
} }
// IsActiveInDir returns true if any session (excluding those matching excludeKey) // IsActiveInDir returns true if any session (excluding those matching excludeKey)
// has touched a directory overlapping with dir within sessionActivityTimeout. // has touched a directory overlapping with dir within sessionActivityTimeout.
// Overlap = either is a prefix of the other (parent/child relationship). // Overlap = either is a prefix of the other (parent/child relationship).
func (st *SessionTracker) IsActiveInDir(dir, excludeKey string) bool { func (st *SessionTracker) IsActiveInDir(dir, excludeKey string) bool {
cutoff := time.Now().Add(-sessionActivityTimeout) cutoff := time.Now().Add(-sessionActivityTimeout)
active := false active := false
st.entries.Range(func(key, val any) bool { st.entries.Range(func(key, val any) bool {
if key.(string) == excludeKey { if key.(string) == excludeKey {
return true return true
} }
entry := val.(*SessionEntry) entry := val.(*SessionEntry)
if entry.LastSeenAt.After(cutoff) && entry.TouchDir != "" && if entry.LastSeenAt.After(cutoff) && entry.TouchDir != "" &&
(strings.HasPrefix(entry.TouchDir, dir) || strings.HasPrefix(dir, entry.TouchDir)) { (strings.HasPrefix(entry.TouchDir, dir) || strings.HasPrefix(dir, entry.TouchDir)) {
active = true active = true
return false return false
} }
return true return true
}) })
return active return active
} }
// ListActive returns all sessions seen within sessionActivityTimeout, // ListActive returns all sessions seen within sessionActivityTimeout,
// sorted by LastSeenAt descending (most recent first). // sorted by LastSeenAt descending (most recent first).
func (st *SessionTracker) ListActive() []SessionEntry { func (st *SessionTracker) ListActive() []SessionEntry {
cutoff := time.Now().Add(-sessionActivityTimeout) cutoff := time.Now().Add(-sessionActivityTimeout)
var result []SessionEntry var result []SessionEntry
st.entries.Range(func(key, val any) bool { st.entries.Range(func(key, val any) bool {
entry := val.(*SessionEntry) entry := val.(*SessionEntry)
if entry.LastSeenAt.After(cutoff) { if entry.LastSeenAt.After(cutoff) {
result = append(result, *entry) // copy result = append(result, *entry) // copy
} }
return true return true
}) })
sort.Slice(result, func(i, j int) bool { sort.Slice(result, func(i, j int) bool {
return result[i].LastSeenAt.After(result[j].LastSeenAt) return result[i].LastSeenAt.After(result[j].LastSeenAt)
}) })
return result return result
} }
// GetTouchDir returns the TouchDir for a given session key, or "" if not found. // GetTouchDir returns the TouchDir for a given session key, or "" if not found.
func (st *SessionTracker) GetTouchDir(sessionKey string) string { func (st *SessionTracker) GetTouchDir(sessionKey string) string {
val, ok := st.entries.Load(sessionKey) val, ok := st.entries.Load(sessionKey)
if !ok { if !ok {
return "" return ""
} }
return val.(*SessionEntry).TouchDir return val.(*SessionEntry).TouchDir
} }
// GetPeerPurposes returns purposes of other active sessions targeting the same project. // GetPeerPurposes returns purposes of other active sessions targeting the same project.
// Used for lightweight coordination without context pollution. // Used for lightweight coordination without context pollution.
func (st *SessionTracker) GetPeerPurposes(sessionKey, projectPath string) []PeerInfo { func (st *SessionTracker) GetPeerPurposes(sessionKey, projectPath string) []PeerInfo {
if projectPath == "" { if projectPath == "" {
return nil return nil
} }
cutoff := time.Now().Add(-sessionActivityTimeout) cutoff := time.Now().Add(-sessionActivityTimeout)
var result []PeerInfo var result []PeerInfo
st.entries.Range(func(key, val any) bool { st.entries.Range(func(key, val any) bool {
if key.(string) == sessionKey { if key.(string) == sessionKey {
return true return true
} }
entry := val.(*SessionEntry) entry := val.(*SessionEntry)
if entry.LastSeenAt.After(cutoff) && entry.ProjectPath == projectPath { if entry.LastSeenAt.After(cutoff) && entry.ProjectPath == projectPath {
result = append(result, PeerInfo{ result = append(result, PeerInfo{
SessionKey: entry.SessionKey, SessionKey: entry.SessionKey,
Purpose: entry.Purpose, Purpose: entry.Purpose,
Branch: entry.Branch, Branch: entry.Branch,
}) })
} }
return true return true
}) })
return result return result
} }

View file

@ -9,38 +9,53 @@ func TestTouch(t *testing.T) {
st := NewSessionTracker() st := NewSessionTracker()
// Basic touch creates entry // Basic touch creates entry
st.Touch("sess1", "telegram", "123", "projects/myapp", nil) st.Touch("sess1", "telegram", "123", "projects/myapp", nil)
entries := st.ListActive() entries := st.ListActive()
if len(entries) != 1 { if len(entries) != 1 {
t.Fatalf("expected 1 entry, got %d", len(entries)) t.Fatalf("expected 1 entry, got %d", len(entries))
} }
if entries[0].SessionKey != "sess1" { if entries[0].SessionKey != "sess1" {
t.Errorf("expected session_key=sess1, got %s", entries[0].SessionKey) t.Errorf("expected session_key=sess1, got %s", entries[0].SessionKey)
} }
if entries[0].Channel != "telegram" { if entries[0].Channel != "telegram" {
t.Errorf("expected channel=telegram, got %s", entries[0].Channel) t.Errorf("expected channel=telegram, got %s", entries[0].Channel)
} }
if entries[0].TouchDir != "projects/myapp" { if entries[0].TouchDir != "projects/myapp" {
t.Errorf("expected touch_dir=projects/myapp, got %s", entries[0].TouchDir) t.Errorf("expected touch_dir=projects/myapp, got %s", entries[0].TouchDir)
} }
// Touch again with new dir overwrites TouchDir // Touch again with new dir overwrites TouchDir
st.Touch("sess1", "", "", "projects/other", nil) st.Touch("sess1", "", "", "projects/other", nil)
entries = st.ListActive() entries = st.ListActive()
if len(entries) != 1 { if len(entries) != 1 {
t.Fatalf("expected 1 entry, got %d", len(entries)) t.Fatalf("expected 1 entry, got %d", len(entries))
} }
if entries[0].TouchDir != "projects/other" { if entries[0].TouchDir != "projects/other" {
t.Errorf("expected touch_dir=projects/other, got %s", entries[0].TouchDir) t.Errorf("expected touch_dir=projects/other, got %s", entries[0].TouchDir)
} }
// Channel should remain from first touch // Channel should remain from first touch
if entries[0].Channel != "telegram" { if entries[0].Channel != "telegram" {
t.Errorf("expected channel=telegram (unchanged), got %s", entries[0].Channel) t.Errorf("expected channel=telegram (unchanged), got %s", entries[0].Channel)
} }
// Touch with empty dir does not overwrite TouchDir // Touch with empty dir does not overwrite TouchDir
st.Touch("sess1", "", "", "", nil) st.Touch("sess1", "", "", "", nil)
entries = st.ListActive() entries = st.ListActive()
if entries[0].TouchDir != "projects/other" { if entries[0].TouchDir != "projects/other" {
t.Errorf("expected touch_dir unchanged, got %s", entries[0].TouchDir) t.Errorf("expected touch_dir unchanged, got %s", entries[0].TouchDir)
} }
@ -50,36 +65,45 @@ func TestIsActiveInDir(t *testing.T) {
st := NewSessionTracker() st := NewSessionTracker()
// Setup: sess1 touches "projects/myapp" // Setup: sess1 touches "projects/myapp"
st.Touch("sess1", "telegram", "123", "projects/myapp", nil) st.Touch("sess1", "telegram", "123", "projects/myapp", nil)
// Same dir, excluding sess1 → false // Same dir, excluding sess1 → false
if st.IsActiveInDir("projects/myapp", "sess1") { if st.IsActiveInDir("projects/myapp", "sess1") {
t.Error("expected false when excluding the only active session") t.Error("expected false when excluding the only active session")
} }
// Same dir, excluding different key → true // Same dir, excluding different key → true
if !st.IsActiveInDir("projects/myapp", "heartbeat") { if !st.IsActiveInDir("projects/myapp", "heartbeat") {
t.Error("expected true for exact dir match") t.Error("expected true for exact dir match")
} }
// Parent dir match: "projects" is prefix of "projects/myapp" // Parent dir match: "projects" is prefix of "projects/myapp"
if !st.IsActiveInDir("projects", "heartbeat") { if !st.IsActiveInDir("projects", "heartbeat") {
t.Error("expected true for parent dir match") t.Error("expected true for parent dir match")
} }
// Child dir match: "projects/myapp/src" has prefix "projects/myapp" // Child dir match: "projects/myapp/src" has prefix "projects/myapp"
if !st.IsActiveInDir("projects/myapp/src", "heartbeat") { if !st.IsActiveInDir("projects/myapp/src", "heartbeat") {
t.Error("expected true for child dir match") t.Error("expected true for child dir match")
} }
// Unrelated dir → false // Unrelated dir → false
if st.IsActiveInDir("other/stuff", "heartbeat") { if st.IsActiveInDir("other/stuff", "heartbeat") {
t.Error("expected false for unrelated dir") t.Error("expected false for unrelated dir")
} }
// Stale entry (manually set LastSeenAt to past) // Stale entry (manually set LastSeenAt to past)
val, _ := st.entries.Load("sess1") val, _ := st.entries.Load("sess1")
entry := val.(*SessionEntry) entry := val.(*SessionEntry)
entry.LastSeenAt = time.Now().Add(-sessionActivityTimeout - time.Minute) entry.LastSeenAt = time.Now().Add(-sessionActivityTimeout - time.Minute)
if st.IsActiveInDir("projects/myapp", "heartbeat") { if st.IsActiveInDir("projects/myapp", "heartbeat") {
@ -91,32 +115,43 @@ func TestListActive(t *testing.T) {
st := NewSessionTracker() st := NewSessionTracker()
// Add two sessions // Add two sessions
st.Touch("sess1", "telegram", "123", "projects/a", nil) st.Touch("sess1", "telegram", "123", "projects/a", nil)
time.Sleep(5 * time.Millisecond) // ensure different timestamps time.Sleep(5 * time.Millisecond) // ensure different timestamps
st.Touch("sess2", "discord", "456", "projects/b", nil) st.Touch("sess2", "discord", "456", "projects/b", nil)
entries := st.ListActive() entries := st.ListActive()
if len(entries) != 2 { if len(entries) != 2 {
t.Fatalf("expected 2 entries, got %d", len(entries)) t.Fatalf("expected 2 entries, got %d", len(entries))
} }
// Most recent first // Most recent first
if entries[0].SessionKey != "sess2" { if entries[0].SessionKey != "sess2" {
t.Errorf("expected sess2 first (most recent), got %s", entries[0].SessionKey) t.Errorf("expected sess2 first (most recent), got %s", entries[0].SessionKey)
} }
if entries[1].SessionKey != "sess1" { if entries[1].SessionKey != "sess1" {
t.Errorf("expected sess1 second, got %s", entries[1].SessionKey) t.Errorf("expected sess1 second, got %s", entries[1].SessionKey)
} }
// Make sess1 stale // Make sess1 stale
val, _ := st.entries.Load("sess1") val, _ := st.entries.Load("sess1")
entry := val.(*SessionEntry) entry := val.(*SessionEntry)
entry.LastSeenAt = time.Now().Add(-sessionActivityTimeout - time.Minute) entry.LastSeenAt = time.Now().Add(-sessionActivityTimeout - time.Minute)
entries = st.ListActive() entries = st.ListActive()
if len(entries) != 1 { if len(entries) != 1 {
t.Fatalf("expected 1 active entry after stale, got %d", len(entries)) t.Fatalf("expected 1 active entry after stale, got %d", len(entries))
} }
if entries[0].SessionKey != "sess2" { if entries[0].SessionKey != "sess2" {
t.Errorf("expected only sess2, got %s", entries[0].SessionKey) 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 // SessionGraph is a thin wrapper around SessionStore that provides
// structured turn-writing via BeginTurn/TurnWriter. // structured turn-writing via BeginTurn/TurnWriter.
// It does NOT replace LegacyAdapter — existing call sites remain unchanged. // It does NOT replace LegacyAdapter — existing call sites remain unchanged.
// Future phases will migrate callers to use SessionGraph directly. // Future phases will migrate callers to use SessionGraph directly.
type SessionGraph struct { type SessionGraph struct {
store SessionStore store SessionStore
} }
// NewSessionGraph creates a SessionGraph backed by the given store. // NewSessionGraph creates a SessionGraph backed by the given store.
func NewSessionGraph(store SessionStore) *SessionGraph { func NewSessionGraph(store SessionStore) *SessionGraph {
return &SessionGraph{store: store} return &SessionGraph{store: store}
} }
// Messages returns all messages for the session by reading turns from the store. // Messages returns all messages for the session by reading turns from the store.
func (g *SessionGraph) Messages(sessionKey string) ([]providers.Message, error) { func (g *SessionGraph) Messages(sessionKey string) ([]providers.Message, error) {
turns, err := g.store.Turns(sessionKey, 0) turns, err := g.store.Turns(sessionKey, 0)
if err != nil { if err != nil {
return nil, err return nil, err
} }
var msgs []providers.Message var msgs []providers.Message
for _, t := range turns { for _, t := range turns {
msgs = append(msgs, t.Messages...) msgs = append(msgs, t.Messages...)
} }
if msgs == nil { if msgs == nil {
msgs = []providers.Message{} msgs = []providers.Message{}
} }
return msgs, nil return msgs, nil
} }
// BeginTurn starts a new turn that can be built up incrementally // BeginTurn starts a new turn that can be built up incrementally
// and committed atomically. // and committed atomically.
func (g *SessionGraph) BeginTurn(sessionKey string, kind TurnKind) *TurnWriter { func (g *SessionGraph) BeginTurn(sessionKey string, kind TurnKind) *TurnWriter {
return &TurnWriter{ return &TurnWriter{
store: g.store, store: g.store,
sessionKey: sessionKey, sessionKey: sessionKey,
turn: Turn{ turn: Turn{
SessionKey: sessionKey, SessionKey: sessionKey,
Kind: kind, Kind: kind,
}, },
} }
} }
// TurnWriter accumulates messages for a single turn and commits them atomically. // TurnWriter accumulates messages for a single turn and commits them atomically.
type TurnWriter struct { type TurnWriter struct {
mu sync.Mutex mu sync.Mutex
store SessionStore store SessionStore
sessionKey string sessionKey string
turn Turn turn Turn
committed bool committed bool
discarded bool discarded bool
} }
// Add appends a message to the pending turn. // Add appends a message to the pending turn.
func (tw *TurnWriter) Add(msg providers.Message) { func (tw *TurnWriter) Add(msg providers.Message) {
tw.mu.Lock() tw.mu.Lock()
defer tw.mu.Unlock() defer tw.mu.Unlock()
tw.turn.Messages = append(tw.turn.Messages, msg) tw.turn.Messages = append(tw.turn.Messages, msg)
} }
// SetOrigin sets the origin session key for this turn (e.g. subagent source). // SetOrigin sets the origin session key for this turn (e.g. subagent source).
func (tw *TurnWriter) SetOrigin(sessionKey string) { func (tw *TurnWriter) SetOrigin(sessionKey string) {
tw.mu.Lock() tw.mu.Lock()
defer tw.mu.Unlock() defer tw.mu.Unlock()
tw.turn.OriginKey = sessionKey tw.turn.OriginKey = sessionKey
} }
// SetAuthor sets the author field for this turn. // SetAuthor sets the author field for this turn.
func (tw *TurnWriter) SetAuthor(author string) { func (tw *TurnWriter) SetAuthor(author string) {
tw.mu.Lock() tw.mu.Lock()
defer tw.mu.Unlock() defer tw.mu.Unlock()
tw.turn.Author = author tw.turn.Author = author
} }
// Commit writes the accumulated turn to the store. // Commit writes the accumulated turn to the store.
// Returns an error if already committed or discarded. // Returns an error if already committed or discarded.
func (tw *TurnWriter) Commit() error { func (tw *TurnWriter) Commit() error {
tw.mu.Lock() tw.mu.Lock()
defer tw.mu.Unlock() defer tw.mu.Unlock()
if tw.committed { if tw.committed {
return errors.New("turn already committed") return errors.New("turn already committed")
} }
if tw.discarded { if tw.discarded {
return errors.New("turn already discarded") return errors.New("turn already discarded")
} }
tw.committed = true tw.committed = true
return tw.store.Append(tw.sessionKey, &tw.turn) return tw.store.Append(tw.sessionKey, &tw.turn)
} }
// Discard marks the turn as abandoned — nothing is written. // Discard marks the turn as abandoned — nothing is written.
func (tw *TurnWriter) Discard() { func (tw *TurnWriter) Discard() {
tw.mu.Lock() tw.mu.Lock()
defer tw.mu.Unlock() defer tw.mu.Unlock()
tw.discarded = true tw.discarded = true
} }

View file

@ -8,19 +8,25 @@ import (
func TestSessionGraph_Messages(t *testing.T) { func TestSessionGraph_Messages(t *testing.T) {
store := newTestStore(t) store := newTestStore(t)
if err := store.Create("g1", nil); err != nil { if err := store.Create("g1", nil); err != nil {
t.Fatal(err) t.Fatal(err)
} }
if err := store.Append("g1", &Turn{ if err := store.Append("g1", &Turn{
Kind: TurnNormal, Kind: TurnNormal,
Messages: []providers.Message{{Role: "user", Content: "hello"}}, Messages: []providers.Message{{Role: "user", Content: "hello"}},
}); err != nil { }); err != nil {
t.Fatal(err) t.Fatal(err)
} }
if err := store.Append("g1", &Turn{ if err := store.Append("g1", &Turn{
Kind: TurnNormal, Kind: TurnNormal,
Messages: []providers.Message{ Messages: []providers.Message{
{Role: "assistant", Content: "hi"}, {Role: "assistant", Content: "hi"},
{Role: "user", Content: "how are you"}, {Role: "user", Content: "how are you"},
}, },
}); err != nil { }); err != nil {
@ -28,13 +34,16 @@ func TestSessionGraph_Messages(t *testing.T) {
} }
g := NewSessionGraph(store) g := NewSessionGraph(store)
msgs, err := g.Messages("g1") msgs, err := g.Messages("g1")
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if len(msgs) != 3 { if len(msgs) != 3 {
t.Fatalf("expected 3 messages, got %d", len(msgs)) t.Fatalf("expected 3 messages, got %d", len(msgs))
} }
if msgs[0].Content != "hello" || msgs[1].Content != "hi" || msgs[2].Content != "how are you" { if msgs[0].Content != "hello" || msgs[1].Content != "hi" || msgs[2].Content != "how are you" {
t.Errorf("unexpected messages: %+v", msgs) t.Errorf("unexpected messages: %+v", msgs)
} }
@ -42,14 +51,18 @@ func TestSessionGraph_Messages(t *testing.T) {
func TestSessionGraph_Messages_Empty(t *testing.T) { func TestSessionGraph_Messages_Empty(t *testing.T) {
store := newTestStore(t) store := newTestStore(t)
if err := store.Create("empty", nil); err != nil { if err := store.Create("empty", nil); err != nil {
t.Fatal(err) t.Fatal(err)
} }
g := NewSessionGraph(store) g := NewSessionGraph(store)
msgs, err := g.Messages("empty") msgs, err := g.Messages("empty")
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if msgs == nil || len(msgs) != 0 { if msgs == nil || len(msgs) != 0 {
t.Errorf("expected empty slice, got %v", msgs) 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) { func TestTurnWriter_Commit(t *testing.T) {
store := newTestStore(t) store := newTestStore(t)
if err := store.Create("tw1", nil); err != nil { if err := store.Create("tw1", nil); err != nil {
t.Fatal(err) t.Fatal(err)
} }
g := NewSessionGraph(store) g := NewSessionGraph(store)
tw := g.BeginTurn("tw1", TurnNormal) tw := g.BeginTurn("tw1", TurnNormal)
tw.Add(providers.Message{Role: "user", Content: "msg1"}) tw.Add(providers.Message{Role: "user", Content: "msg1"})
tw.Add(providers.Message{Role: "assistant", Content: "msg2"}) tw.Add(providers.Message{Role: "assistant", Content: "msg2"})
tw.SetOrigin("parent-key") tw.SetOrigin("parent-key")
tw.SetAuthor("agent-1") tw.SetAuthor("agent-1")
if err := tw.Commit(); err != nil { if err := tw.Commit(); err != nil {
@ -76,15 +95,19 @@ func TestTurnWriter_Commit(t *testing.T) {
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if len(turns) != 1 { if len(turns) != 1 {
t.Fatalf("expected 1 turn, got %d", len(turns)) t.Fatalf("expected 1 turn, got %d", len(turns))
} }
if len(turns[0].Messages) != 2 { if len(turns[0].Messages) != 2 {
t.Fatalf("expected 2 messages, got %d", len(turns[0].Messages)) t.Fatalf("expected 2 messages, got %d", len(turns[0].Messages))
} }
if turns[0].OriginKey != "parent-key" { if turns[0].OriginKey != "parent-key" {
t.Errorf("expected origin 'parent-key', got %q", turns[0].OriginKey) t.Errorf("expected origin 'parent-key', got %q", turns[0].OriginKey)
} }
if turns[0].Author != "agent-1" { if turns[0].Author != "agent-1" {
t.Errorf("expected author 'agent-1', got %q", turns[0].Author) 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) { func TestTurnWriter_Discard(t *testing.T) {
store := newTestStore(t) store := newTestStore(t)
if err := store.Create("tw2", nil); err != nil { if err := store.Create("tw2", nil); err != nil {
t.Fatal(err) t.Fatal(err)
} }
g := NewSessionGraph(store) g := NewSessionGraph(store)
tw := g.BeginTurn("tw2", TurnNormal) tw := g.BeginTurn("tw2", TurnNormal)
tw.Add(providers.Message{Role: "user", Content: "should not persist"}) tw.Add(providers.Message{Role: "user", Content: "should not persist"})
tw.Discard() tw.Discard()
turns, err := store.Turns("tw2", 0) turns, err := store.Turns("tw2", 0)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if len(turns) != 0 { if len(turns) != 0 {
t.Errorf("expected 0 turns after discard, got %d", len(turns)) 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) { func TestTurnWriter_DoubleCommit(t *testing.T) {
store := newTestStore(t) store := newTestStore(t)
if err := store.Create("tw3", nil); err != nil { if err := store.Create("tw3", nil); err != nil {
t.Fatal(err) t.Fatal(err)
} }
g := NewSessionGraph(store) g := NewSessionGraph(store)
tw := g.BeginTurn("tw3", TurnNormal) tw := g.BeginTurn("tw3", TurnNormal)
tw.Add(providers.Message{Role: "user", Content: "once"}) tw.Add(providers.Message{Role: "user", Content: "once"})
if err := tw.Commit(); err != nil { if err := tw.Commit(); err != nil {
t.Fatal(err) t.Fatal(err)
} }
if err := tw.Commit(); err == nil { if err := tw.Commit(); err == nil {
t.Error("expected error on double commit") t.Error("expected error on double commit")
} }
@ -130,13 +162,17 @@ func TestTurnWriter_DoubleCommit(t *testing.T) {
func TestTurnWriter_CommitAfterDiscard(t *testing.T) { func TestTurnWriter_CommitAfterDiscard(t *testing.T) {
store := newTestStore(t) store := newTestStore(t)
if err := store.Create("tw4", nil); err != nil { if err := store.Create("tw4", nil); err != nil {
t.Fatal(err) t.Fatal(err)
} }
g := NewSessionGraph(store) g := NewSessionGraph(store)
tw := g.BeginTurn("tw4", TurnNormal) tw := g.BeginTurn("tw4", TurnNormal)
tw.Add(providers.Message{Role: "user", Content: "x"}) tw.Add(providers.Message{Role: "user", Content: "x"})
tw.Discard() tw.Discard()
if err := tw.Commit(); err == nil { 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. // DefaultPruneTTL is the default time-to-live for session pruning.
const DefaultPruneTTL = 7 * 24 * time.Hour const DefaultPruneTTL = 7 * 24 * time.Hour
// CompactOldTurns flushes pending writes, then compacts SQLite turns // CompactOldTurns flushes pending writes, then compacts SQLite turns
// keeping only the last keepLast messages. Sets session summary to the given value. // keeping only the last keepLast messages. Sets session summary to the given value.
func (la *LegacyAdapter) CompactOldTurns(key string, keepLast int, summary string) error { func (la *LegacyAdapter) CompactOldTurns(key string, keepLast int, summary string) error {
// 1. Flush pending messages to SQLite // 1. Flush pending messages to SQLite
if err := la.Save(key); err != nil { if err := la.Save(key); err != nil {
return err return err
} }
// 2. Query all turns // 2. Query all turns
turns, err := la.store.Turns(key, 0) turns, err := la.store.Turns(key, 0)
if err != nil { if err != nil {
return err return err
} }
// 3. Count total messages, find cut point // 3. Count total messages, find cut point
totalMsgs := 0 totalMsgs := 0
for _, t := range turns { for _, t := range turns {
totalMsgs += len(t.Messages) totalMsgs += len(t.Messages)
} }
if keepLast >= totalMsgs { if keepLast >= totalMsgs {
// Nothing to compact, just update summary // Nothing to compact, just update summary
if err := la.store.SetSummary(key, summary); err != nil { if err := la.store.SetSummary(key, summary); err != nil {
return err return err
} }
la.mu.Lock() la.mu.Lock()
if c, ok := la.cache[key]; ok { if c, ok := la.cache[key]; ok {
c.summary = summary c.summary = summary
} }
la.mu.Unlock() la.mu.Unlock()
return nil return nil
} }
dropCount := totalMsgs - keepLast dropCount := totalMsgs - keepLast
accumulated := 0 accumulated := 0
cutSeq := 0 cutSeq := 0
for _, t := range turns { for _, t := range turns {
accumulated += len(t.Messages) accumulated += len(t.Messages)
if accumulated <= dropCount { if accumulated <= dropCount {
cutSeq = t.Seq cutSeq = t.Seq
} else { } else {
break break
} }
} }
if cutSeq == 0 { if cutSeq == 0 {
if err := la.store.SetSummary(key, summary); err != nil { if err := la.store.SetSummary(key, summary); err != nil {
return err return err
} }
la.mu.Lock() la.mu.Lock()
if c, ok := la.cache[key]; ok { if c, ok := la.cache[key]; ok {
c.summary = summary c.summary = summary
} }
la.mu.Unlock() la.mu.Unlock()
return nil return nil
} }
// 4. Compact in SQLite // 4. Compact in SQLite
if err := la.store.Compact(key, cutSeq, summary); err != nil { if err := la.store.Compact(key, cutSeq, summary); err != nil {
return err return err
} }
// 5. Update in-memory cache // 5. Update in-memory cache
la.mu.Lock() la.mu.Lock()
defer la.mu.Unlock() defer la.mu.Unlock()
if c, ok := la.cache[key]; ok { if c, ok := la.cache[key]; ok {
if keepLast < len(c.messages) { if keepLast < len(c.messages) {
c.messages = c.messages[len(c.messages)-keepLast:] c.messages = c.messages[len(c.messages)-keepLast:]
} }
c.stored = len(c.messages) c.stored = len(c.messages)
c.replaced = false c.replaced = false
c.dirty = false c.dirty = false
c.summary = summary c.summary = summary
} }
return nil return nil
} }
// Store returns the underlying SessionStore for direct DAG operations. // Store returns the underlying SessionStore for direct DAG operations.
func (la *LegacyAdapter) Store() SessionStore { func (la *LegacyAdapter) Store() SessionStore {
return la.store return la.store
} }
// Graph returns a SessionGraph backed by the underlying store. // Graph returns a SessionGraph backed by the underlying store.
func (la *LegacyAdapter) Graph() *SessionGraph { func (la *LegacyAdapter) Graph() *SessionGraph {
return NewSessionGraph(la.store) return NewSessionGraph(la.store)
} }
// AdvanceStored increments the stored counter for a session by delta, // AdvanceStored increments the stored counter for a session by delta,
// preventing the flush loop from re-persisting messages already written // preventing the flush loop from re-persisting messages already written
// directly to the store (e.g. TurnReport). // directly to the store (e.g. TurnReport).
func (la *LegacyAdapter) AdvanceStored(key string, delta int) { func (la *LegacyAdapter) AdvanceStored(key string, delta int) {
la.mu.Lock() la.mu.Lock()
defer la.mu.Unlock() defer la.mu.Unlock()
if c, ok := la.cache[key]; ok { if c, ok := la.cache[key]; ok {
c.stored += delta c.stored += delta
} }
@ -573,16 +616,25 @@ func (la *LegacyAdapter) Close() {
func (la *LegacyAdapter) flushLoop() { func (la *LegacyAdapter) flushLoop() {
flushTicker := time.NewTicker(5 * time.Minute) flushTicker := time.NewTicker(5 * time.Minute)
pruneTicker := time.NewTicker(6 * time.Hour) pruneTicker := time.NewTicker(6 * time.Hour)
defer flushTicker.Stop() defer flushTicker.Stop()
defer pruneTicker.Stop() defer pruneTicker.Stop()
for { for {
select { select {
case <-flushTicker.C: case <-flushTicker.C:
la.FlushDirty() la.FlushDirty()
case <-pruneTicker.C: case <-pruneTicker.C:
_, _ = la.store.Prune(DefaultPruneTTL) _, _ = la.store.Prune(DefaultPruneTTL)
case <-la.done: case <-la.done:
return return
} }
} }

View file

@ -126,7 +126,11 @@ func TestBackend_AddFullMessage(t *testing.T) {
Content: "sure", Content: "sure",
ToolCalls: []providers.ToolCall{ 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) { func TestCompactOldTurns(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "test.db") dbPath := filepath.Join(t.TempDir(), "test.db")
store, err := OpenSQLiteStore(dbPath) store, err := OpenSQLiteStore(dbPath)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
la := NewLegacyAdapter(store) la := NewLegacyAdapter(store)
defer la.Close() defer la.Close()
la.GetOrCreate("k1") la.GetOrCreate("k1")
// Turn 1: 2 messages // Turn 1: 2 messages
la.AddMessage("k1", "user", "a") la.AddMessage("k1", "user", "a")
la.AddMessage("k1", "assistant", "b") la.AddMessage("k1", "assistant", "b")
la.Save("k1") la.Save("k1")
// Turn 2: 3 messages // Turn 2: 3 messages
la.AddMessage("k1", "user", "c") la.AddMessage("k1", "user", "c")
la.AddMessage("k1", "assistant", "d") la.AddMessage("k1", "assistant", "d")
la.AddMessage("k1", "user", "e") la.AddMessage("k1", "user", "e")
la.Save("k1") la.Save("k1")
// Turn 3: 2 messages // Turn 3: 2 messages
la.AddMessage("k1", "user", "f") la.AddMessage("k1", "user", "f")
la.AddMessage("k1", "assistant", "g") la.AddMessage("k1", "assistant", "g")
la.Save("k1") la.Save("k1")
// Total: 7 messages across 3 turns. keepLast=2 → drop 5 → compact turns 1+2 (5 msgs) // 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 { if err := la.CompactOldTurns("k1", 2, "test summary"); err != nil {
t.Fatalf("CompactOldTurns: %v", err) t.Fatalf("CompactOldTurns: %v", err)
} }
h := la.GetHistory("k1") h := la.GetHistory("k1")
if len(h) != 2 { if len(h) != 2 {
t.Fatalf("expected 2 messages in cache, got %d", len(h)) t.Fatalf("expected 2 messages in cache, got %d", len(h))
} }
if h[0].Content != "f" || h[1].Content != "g" { if h[0].Content != "f" || h[1].Content != "g" {
t.Errorf("unexpected messages: %+v", h) t.Errorf("unexpected messages: %+v", h)
} }
if s := la.GetSummary("k1"); s != "test summary" { if s := la.GetSummary("k1"); s != "test summary" {
t.Errorf("expected summary 'test summary', got %q", s) t.Errorf("expected summary 'test summary', got %q", s)
} }
// Verify in SQLite: only turn 3 remains // Verify in SQLite: only turn 3 remains
turns, _ := store.Turns("k1", 0) turns, _ := store.Turns("k1", 0)
if len(turns) != 1 { if len(turns) != 1 {
t.Fatalf("expected 1 turn in SQLite, got %d", len(turns)) t.Fatalf("expected 1 turn in SQLite, got %d", len(turns))
} }
if len(turns[0].Messages) != 2 { if len(turns[0].Messages) != 2 {
t.Errorf("expected 2 messages in remaining turn, got %d", len(turns[0].Messages)) 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) { func TestCompactOldTurns_NothingToCompact(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "test.db") dbPath := filepath.Join(t.TempDir(), "test.db")
store, err := OpenSQLiteStore(dbPath) store, err := OpenSQLiteStore(dbPath)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
la := NewLegacyAdapter(store) la := NewLegacyAdapter(store)
defer la.Close() defer la.Close()
la.GetOrCreate("k1") la.GetOrCreate("k1")
la.AddMessage("k1", "user", "a") la.AddMessage("k1", "user", "a")
la.AddMessage("k1", "assistant", "b") la.AddMessage("k1", "assistant", "b")
la.Save("k1") la.Save("k1")
// keepLast=10 >= total 2 → nothing compacted, summary still updated // keepLast=10 >= total 2 → nothing compacted, summary still updated
if err := la.CompactOldTurns("k1", 10, "new summary"); err != nil { if err := la.CompactOldTurns("k1", 10, "new summary"); err != nil {
t.Fatalf("CompactOldTurns: %v", err) t.Fatalf("CompactOldTurns: %v", err)
} }
h := la.GetHistory("k1") h := la.GetHistory("k1")
if len(h) != 2 { if len(h) != 2 {
t.Fatalf("expected 2 messages, got %d", len(h)) t.Fatalf("expected 2 messages, got %d", len(h))
} }
if s := la.GetSummary("k1"); s != "new summary" { if s := la.GetSummary("k1"); s != "new summary" {
t.Errorf("expected 'new summary', got %q", s) t.Errorf("expected 'new summary', got %q", s)
} }
@ -563,29 +599,40 @@ func TestCompactOldTurns_NothingToCompact(t *testing.T) {
func TestCompactOldTurns_SingleTurn(t *testing.T) { func TestCompactOldTurns_SingleTurn(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "test.db") dbPath := filepath.Join(t.TempDir(), "test.db")
store, err := OpenSQLiteStore(dbPath) store, err := OpenSQLiteStore(dbPath)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
la := NewLegacyAdapter(store) la := NewLegacyAdapter(store)
defer la.Close() defer la.Close()
la.GetOrCreate("k1") la.GetOrCreate("k1")
la.AddMessage("k1", "user", "a") la.AddMessage("k1", "user", "a")
la.AddMessage("k1", "assistant", "b") la.AddMessage("k1", "assistant", "b")
la.AddMessage("k1", "user", "c") la.AddMessage("k1", "user", "c")
la.Save("k1") la.Save("k1")
// Single turn with 3 messages, keepLast=2 → dropCount=1, but first turn has 3 msgs // 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 // accumulated(3) > dropCount(1) on first turn → cutSeq=0 → no compaction
if err := la.CompactOldTurns("k1", 2, "sum"); err != nil { if err := la.CompactOldTurns("k1", 2, "sum"); err != nil {
t.Fatalf("CompactOldTurns: %v", err) t.Fatalf("CompactOldTurns: %v", err)
} }
h := la.GetHistory("k1") h := la.GetHistory("k1")
if len(h) != 3 { if len(h) != 3 {
t.Fatalf("expected 3 messages (no compaction), got %d", len(h)) t.Fatalf("expected 3 messages (no compaction), got %d", len(h))
} }
if s := la.GetSummary("k1"); s != "sum" { if s := la.GetSummary("k1"); s != "sum" {
t.Errorf("expected 'sum', got %q", s) t.Errorf("expected 'sum', got %q", s)
} }
@ -593,22 +640,29 @@ func TestCompactOldTurns_SingleTurn(t *testing.T) {
func TestCompactOldTurns_Graph(t *testing.T) { func TestCompactOldTurns_Graph(t *testing.T) {
dbPath := filepath.Join(t.TempDir(), "test.db") dbPath := filepath.Join(t.TempDir(), "test.db")
store, err := OpenSQLiteStore(dbPath) store, err := OpenSQLiteStore(dbPath)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
la := NewLegacyAdapter(store) la := NewLegacyAdapter(store)
defer la.Close() defer la.Close()
la.GetOrCreate("k1") la.GetOrCreate("k1")
la.AddMessage("k1", "user", "hello") la.AddMessage("k1", "user", "hello")
la.Save("k1") la.Save("k1")
g := la.Graph() g := la.Graph()
msgs, err := g.Messages("k1") msgs, err := g.Messages("k1")
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }
if len(msgs) != 1 || msgs[0].Content != "hello" { if len(msgs) != 1 || msgs[0].Content != "hello" {
t.Errorf("unexpected graph messages: %+v", msgs) t.Errorf("unexpected graph messages: %+v", msgs)
} }

View file

@ -13,55 +13,75 @@ import (
type Session struct { type Session struct {
Key string `json:"key"` Key string `json:"key"`
Messages []providers.Message `json:"messages"` Messages []providers.Message `json:"messages"`
Summary string `json:"summary,omitempty"` Summary string `json:"summary,omitempty"`
Created time.Time `json:"created"` Created time.Time `json:"created"`
Updated time.Time `json:"updated"` Updated time.Time `json:"updated"`
} }
type SessionManager struct { type SessionManager struct {
sessions map[string]*Session sessions map[string]*Session
mu sync.RWMutex mu sync.RWMutex
storage string storage string
// Write-behind: dirty keys are flushed periodically to reduce disk writes. // Write-behind: dirty keys are flushed periodically to reduce disk writes.
dirtyMu sync.Mutex dirtyMu sync.Mutex
dirtyKeys map[string]bool dirtyKeys map[string]bool
done chan struct{} done chan struct{}
} }
func NewSessionManager(storage string) *SessionManager { func NewSessionManager(storage string) *SessionManager {
sm := &SessionManager{ sm := &SessionManager{
sessions: make(map[string]*Session), sessions: make(map[string]*Session),
storage: storage, storage: storage,
dirtyKeys: make(map[string]bool), dirtyKeys: make(map[string]bool),
done: make(chan struct{}), done: make(chan struct{}),
} }
if storage != "" { if storage != "" {
os.MkdirAll(storage, 0o755) os.MkdirAll(storage, 0o755)
sm.loadSessions() sm.loadSessions()
} }
go sm.flushLoop() go sm.flushLoop()
return sm return sm
} }
func (sm *SessionManager) GetOrCreate(key string) *Session { func (sm *SessionManager) GetOrCreate(key string) *Session {
sm.mu.Lock() sm.mu.Lock()
defer sm.mu.Unlock() defer sm.mu.Unlock()
session, ok := sm.sessions[key] session, ok := sm.sessions[key]
if ok { if ok {
return session return session
} }
session = &Session{ session = &Session{
Key: key, Key: key,
Messages: []providers.Message{}, Messages: []providers.Message{},
Created: time.Now(), Created: time.Now(),
Updated: time.Now(), Updated: time.Now(),
} }
sm.sessions[key] = session sm.sessions[key] = session
return session return session
@ -70,78 +90,101 @@ func (sm *SessionManager) GetOrCreate(key string) *Session {
func (sm *SessionManager) AddMessage(sessionKey, role, content string) { func (sm *SessionManager) AddMessage(sessionKey, role, content string) {
sm.AddFullMessage(sessionKey, providers.Message{ sm.AddFullMessage(sessionKey, providers.Message{
Role: role, Role: role,
Content: content, Content: content,
}) })
} }
// AddFullMessage adds a complete message with tool calls and tool call ID to the session. // 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. // This is used to save the full conversation flow including tool calls and tool results.
func (sm *SessionManager) AddFullMessage(sessionKey string, msg providers.Message) { func (sm *SessionManager) AddFullMessage(sessionKey string, msg providers.Message) {
sm.mu.Lock() sm.mu.Lock()
defer sm.mu.Unlock() defer sm.mu.Unlock()
session, ok := sm.sessions[sessionKey] session, ok := sm.sessions[sessionKey]
if !ok { if !ok {
session = &Session{ session = &Session{
Key: sessionKey, Key: sessionKey,
Messages: []providers.Message{}, Messages: []providers.Message{},
Created: time.Now(), Created: time.Now(),
} }
sm.sessions[sessionKey] = session sm.sessions[sessionKey] = session
} }
session.Messages = append(session.Messages, msg) session.Messages = append(session.Messages, msg)
session.Updated = time.Now() session.Updated = time.Now()
} }
func (sm *SessionManager) GetHistory(key string) []providers.Message { func (sm *SessionManager) GetHistory(key string) []providers.Message {
sm.mu.RLock() sm.mu.RLock()
defer sm.mu.RUnlock() defer sm.mu.RUnlock()
session, ok := sm.sessions[key] session, ok := sm.sessions[key]
if !ok { if !ok {
return []providers.Message{} return []providers.Message{}
} }
history := make([]providers.Message, len(session.Messages)) history := make([]providers.Message, len(session.Messages))
copy(history, session.Messages) copy(history, session.Messages)
return history return history
} }
func (sm *SessionManager) GetSummary(key string) string { func (sm *SessionManager) GetSummary(key string) string {
sm.mu.RLock() sm.mu.RLock()
defer sm.mu.RUnlock() defer sm.mu.RUnlock()
session, ok := sm.sessions[key] session, ok := sm.sessions[key]
if !ok { if !ok {
return "" return ""
} }
return session.Summary return session.Summary
} }
func (sm *SessionManager) SetSummary(key string, summary string) { func (sm *SessionManager) SetSummary(key string, summary string) {
sm.mu.Lock() sm.mu.Lock()
defer sm.mu.Unlock() defer sm.mu.Unlock()
session, ok := sm.sessions[key] session, ok := sm.sessions[key]
if ok { if ok {
session.Summary = summary session.Summary = summary
session.Updated = time.Now() session.Updated = time.Now()
} }
} }
func (sm *SessionManager) TruncateHistory(key string, keepLast int) { func (sm *SessionManager) TruncateHistory(key string, keepLast int) {
sm.mu.Lock() sm.mu.Lock()
defer sm.mu.Unlock() defer sm.mu.Unlock()
session, ok := sm.sessions[key] session, ok := sm.sessions[key]
if !ok { if !ok {
return return
} }
if keepLast <= 0 { if keepLast <= 0 {
session.Messages = []providers.Message{} session.Messages = []providers.Message{}
session.Updated = time.Now() session.Updated = time.Now()
return return
} }
@ -150,14 +193,20 @@ func (sm *SessionManager) TruncateHistory(key string, keepLast int) {
} }
session.Messages = session.Messages[len(session.Messages)-keepLast:] session.Messages = session.Messages[len(session.Messages)-keepLast:]
session.Updated = time.Now() session.Updated = time.Now()
} }
// sanitizeFilename converts a session key into a cross-platform safe filename. // sanitizeFilename converts a session key into a cross-platform safe filename.
// Session keys use "channel:chatID" (e.g. "telegram:123456") but ':' is the // Session keys use "channel:chatID" (e.g. "telegram:123456") but ':' is the
// volume separator on Windows, so filepath.Base would misinterpret the key. // volume separator on Windows, so filepath.Base would misinterpret the key.
// We replace it with '_'. The original key is preserved inside the JSON file, // We replace it with '_'. The original key is preserved inside the JSON file,
// so loadSessions still maps back to the right in-memory key. // so loadSessions still maps back to the right in-memory key.
func sanitizeFilename(key string) string { func sanitizeFilename(key string) string {
return strings.ReplaceAll(key, ":", "_") return strings.ReplaceAll(key, ":", "_")
} }
@ -170,33 +219,47 @@ func (sm *SessionManager) Save(key string) error {
filename := sanitizeFilename(key) filename := sanitizeFilename(key)
// filepath.IsLocal rejects empty names, "..", absolute paths, and // filepath.IsLocal rejects empty names, "..", absolute paths, and
// OS-reserved device names (NUL, COM1 … on Windows). // OS-reserved device names (NUL, COM1 … on Windows).
// The extra checks reject "." and any directory separators so that // The extra checks reject "." and any directory separators so that
// the session file is always written directly inside sm.storage. // the session file is always written directly inside sm.storage.
if filename == "." || !filepath.IsLocal(filename) || strings.ContainsAny(filename, `/\`) { if filename == "." || !filepath.IsLocal(filename) || strings.ContainsAny(filename, `/\`) {
return os.ErrInvalid return os.ErrInvalid
} }
// Snapshot under read lock, then perform slow file I/O after unlock. // Snapshot under read lock, then perform slow file I/O after unlock.
sm.mu.RLock() sm.mu.RLock()
stored, ok := sm.sessions[key] stored, ok := sm.sessions[key]
if !ok { if !ok {
sm.mu.RUnlock() sm.mu.RUnlock()
return nil return nil
} }
snapshot := Session{ snapshot := Session{
Key: stored.Key, Key: stored.Key,
Summary: stored.Summary, Summary: stored.Summary,
Created: stored.Created, Created: stored.Created,
Updated: stored.Updated, Updated: stored.Updated,
} }
if len(stored.Messages) > 0 { if len(stored.Messages) > 0 {
snapshot.Messages = make([]providers.Message, len(stored.Messages)) snapshot.Messages = make([]providers.Message, len(stored.Messages))
copy(snapshot.Messages, stored.Messages) copy(snapshot.Messages, stored.Messages)
} else { } else {
snapshot.Messages = []providers.Message{} snapshot.Messages = []providers.Message{}
} }
sm.mu.RUnlock() sm.mu.RUnlock()
data, err := json.MarshalIndent(snapshot, "", " ") data, err := json.MarshalIndent(snapshot, "", " ")
@ -205,13 +268,16 @@ func (sm *SessionManager) Save(key string) error {
} }
sessionPath := filepath.Join(sm.storage, filename+".json") sessionPath := filepath.Join(sm.storage, filename+".json")
tmpFile, err := os.CreateTemp(sm.storage, "session-*.tmp") tmpFile, err := os.CreateTemp(sm.storage, "session-*.tmp")
if err != nil { if err != nil {
return err return err
} }
tmpPath := tmpFile.Name() tmpPath := tmpFile.Name()
cleanup := true cleanup := true
defer func() { defer func() {
if cleanup { if cleanup {
_ = os.Remove(tmpPath) _ = os.Remove(tmpPath)
@ -220,16 +286,22 @@ func (sm *SessionManager) Save(key string) error {
if _, err := tmpFile.Write(data); err != nil { if _, err := tmpFile.Write(data); err != nil {
_ = tmpFile.Close() _ = tmpFile.Close()
return err return err
} }
if err := tmpFile.Chmod(0o644); err != nil { if err := tmpFile.Chmod(0o644); err != nil {
_ = tmpFile.Close() _ = tmpFile.Close()
return err return err
} }
if err := tmpFile.Sync(); err != nil { if err := tmpFile.Sync(); err != nil {
_ = tmpFile.Close() _ = tmpFile.Close()
return err return err
} }
if err := tmpFile.Close(); err != nil { if err := tmpFile.Close(); err != nil {
return err return err
} }
@ -237,7 +309,9 @@ func (sm *SessionManager) Save(key string) error {
if err := os.Rename(tmpPath, sessionPath); err != nil { if err := os.Rename(tmpPath, sessionPath); err != nil {
return err return err
} }
cleanup = false cleanup = false
return nil return nil
} }
@ -257,12 +331,14 @@ func (sm *SessionManager) loadSessions() error {
} }
sessionPath := filepath.Join(sm.storage, file.Name()) sessionPath := filepath.Join(sm.storage, file.Name())
data, err := os.ReadFile(sessionPath) data, err := os.ReadFile(sessionPath)
if err != nil { if err != nil {
continue continue
} }
var session Session var session Session
if err := json.Unmarshal(data, &session); err != nil { if err := json.Unmarshal(data, &session); err != nil {
continue continue
} }
@ -274,57 +350,84 @@ func (sm *SessionManager) loadSessions() error {
} }
// SanitizeHistory rebuilds session history to ensure valid tool-call ordering. // SanitizeHistory rebuilds session history to ensure valid tool-call ordering.
// LLM APIs require that every assistant message with ToolCalls is immediately // LLM APIs require that every assistant message with ToolCalls is immediately
// followed by exactly the matching tool-result messages (role="tool"), with no // followed by exactly the matching tool-result messages (role="tool"), with no
// other messages in between. Violations can happen from session collisions or // other messages in between. Violations can happen from session collisions or
// mid-execution crashes. // mid-execution crashes.
// //
// The function walks the full history and copies only well-formed groups: // The function walks the full history and copies only well-formed groups:
// - user/system messages are always kept // - user/system messages are always kept
// - assistant messages without tool calls are always kept // - assistant messages without tool calls are always kept
// - assistant messages WITH tool calls are kept only if the immediately // - assistant messages WITH tool calls are kept only if the immediately
// following messages are the complete set of matching tool results // following messages are the complete set of matching tool results
// //
// Returns the sanitized history and the number of messages removed. // Returns the sanitized history and the number of messages removed.
func SanitizeHistory(history []providers.Message) ([]providers.Message, int) { func SanitizeHistory(history []providers.Message) ([]providers.Message, int) {
if len(history) == 0 { if len(history) == 0 {
return history, 0 return history, 0
} }
result := make([]providers.Message, 0, len(history)) result := make([]providers.Message, 0, len(history))
i := 0 i := 0
for i < len(history) { for i < len(history) {
msg := history[i] msg := history[i]
// Non-assistant messages or assistant without tool calls: keep // Non-assistant messages or assistant without tool calls: keep
if msg.Role != "assistant" || len(msg.ToolCalls) == 0 { if msg.Role != "assistant" || len(msg.ToolCalls) == 0 {
// Skip stray tool results not preceded by their assistant // Skip stray tool results not preceded by their assistant
if msg.Role == "tool" { if msg.Role == "tool" {
i++ i++
continue continue
} }
result = append(result, msg) result = append(result, msg)
i++ i++
continue continue
} }
// Assistant with tool calls: validate the immediately following messages // Assistant with tool calls: validate the immediately following messages
expectedIDs := make(map[string]bool, len(msg.ToolCalls)) expectedIDs := make(map[string]bool, len(msg.ToolCalls))
for _, tc := range msg.ToolCalls { for _, tc := range msg.ToolCalls {
expectedIDs[tc.ID] = true expectedIDs[tc.ID] = true
} }
needed := len(expectedIDs) needed := len(expectedIDs)
// Peek ahead: the next `needed` messages must all be tool results with matching IDs // Peek ahead: the next `needed` messages must all be tool results with matching IDs
groupOK := true groupOK := true
if i+needed >= len(history) { if i+needed >= len(history) {
groupOK = false groupOK = false
} else { } else {
for j := 0; j < needed; j++ { for j := 0; j < needed; j++ {
next := history[i+1+j] next := history[i+1+j]
if next.Role != "tool" || !expectedIDs[next.ToolCallID] { if next.Role != "tool" || !expectedIDs[next.ToolCallID] {
groupOK = false groupOK = false
break break
} }
} }
@ -332,14 +435,19 @@ func SanitizeHistory(history []providers.Message) ([]providers.Message, int) {
if groupOK { if groupOK {
// Copy assistant + all tool results // Copy assistant + all tool results
result = append(result, msg) result = append(result, msg)
for j := 0; j < needed; j++ { for j := 0; j < needed; j++ {
result = append(result, history[i+1+j]) result = append(result, history[i+1+j])
} }
i += 1 + needed i += 1 + needed
} else { } else {
// Skip the broken assistant message; tool results will be skipped // Skip the broken assistant message; tool results will be skipped
// individually when encountered (the "stray tool result" check above) // individually when encountered (the "stray tool result" check above)
i++ i++
} }
} }
@ -348,37 +456,54 @@ func SanitizeHistory(history []providers.Message) ([]providers.Message, int) {
} }
// SetHistory updates the messages of a session. // SetHistory updates the messages of a session.
func (sm *SessionManager) SetHistory(key string, history []providers.Message) { func (sm *SessionManager) SetHistory(key string, history []providers.Message) {
sm.mu.Lock() sm.mu.Lock()
defer sm.mu.Unlock() defer sm.mu.Unlock()
session, ok := sm.sessions[key] session, ok := sm.sessions[key]
if ok { if ok {
// Create a deep copy to strictly isolate internal state // Create a deep copy to strictly isolate internal state
// from the caller's slice. // from the caller's slice.
msgs := make([]providers.Message, len(history)) msgs := make([]providers.Message, len(history))
copy(msgs, history) copy(msgs, history)
session.Messages = msgs session.Messages = msgs
session.Updated = time.Now() session.Updated = time.Now()
} }
} }
// MarkDirty marks a session key for deferred persistence. // MarkDirty marks a session key for deferred persistence.
// The session will be written to disk on the next periodic flush or on Close(). // The session will be written to disk on the next periodic flush or on Close().
func (sm *SessionManager) MarkDirty(key string) { func (sm *SessionManager) MarkDirty(key string) {
sm.dirtyMu.Lock() sm.dirtyMu.Lock()
sm.dirtyKeys[key] = true sm.dirtyKeys[key] = true
sm.dirtyMu.Unlock() sm.dirtyMu.Unlock()
} }
// FlushDirty writes all dirty sessions to disk. // FlushDirty writes all dirty sessions to disk.
func (sm *SessionManager) FlushDirty() { func (sm *SessionManager) FlushDirty() {
sm.dirtyMu.Lock() sm.dirtyMu.Lock()
keys := make([]string, 0, len(sm.dirtyKeys)) keys := make([]string, 0, len(sm.dirtyKeys))
for k := range sm.dirtyKeys { for k := range sm.dirtyKeys {
keys = append(keys, k) keys = append(keys, k)
} }
sm.dirtyKeys = make(map[string]bool) sm.dirtyKeys = make(map[string]bool)
sm.dirtyMu.Unlock() sm.dirtyMu.Unlock()
for _, k := range keys { for _, k := range keys {
@ -387,24 +512,34 @@ func (sm *SessionManager) FlushDirty() {
} }
// Close stops the background flush goroutine and writes all dirty sessions. // Close stops the background flush goroutine and writes all dirty sessions.
func (sm *SessionManager) Close() { func (sm *SessionManager) Close() {
select { select {
case <-sm.done: case <-sm.done:
return // already closed return // already closed
default: default:
} }
close(sm.done) close(sm.done)
sm.FlushDirty() sm.FlushDirty()
} }
func (sm *SessionManager) flushLoop() { func (sm *SessionManager) flushLoop() {
ticker := time.NewTicker(5 * time.Minute) ticker := time.NewTicker(5 * time.Minute)
defer ticker.Stop() defer ticker.Stop()
for { for {
select { select {
case <-ticker.C: case <-ticker.C:
sm.FlushDirty() sm.FlushDirty()
case <-sm.done: case <-sm.done:
return return
} }
} }

View file

@ -11,19 +11,26 @@ import (
func TestSanitizeFilename(t *testing.T) { func TestSanitizeFilename(t *testing.T) {
tests := []struct { tests := []struct {
input string input string
expected string expected string
}{ }{
{"simple", "simple"}, {"simple", "simple"},
{"telegram:123456", "telegram_123456"}, {"telegram:123456", "telegram_123456"},
{"discord:987654321", "discord_987654321"}, {"discord:987654321", "discord_987654321"},
{"slack:C01234", "slack_C01234"}, {"slack:C01234", "slack_C01234"},
{"no-colons-here", "no-colons-here"}, {"no-colons-here", "no-colons-here"},
{"multiple:colons:here", "multiple_colons_here"}, {"multiple:colons:here", "multiple_colons_here"},
} }
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.input, func(t *testing.T) { t.Run(tt.input, func(t *testing.T) {
got := sanitizeFilename(tt.input) got := sanitizeFilename(tt.input)
if got != tt.expected { if got != tt.expected {
t.Errorf("sanitizeFilename(%q) = %q, want %q", tt.input, 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) { func TestSave_WithColonInKey(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
sm := NewSessionManager(tmpDir) sm := NewSessionManager(tmpDir)
// Create a session with a key containing colon (typical channel session key). // Create a session with a key containing colon (typical channel session key).
key := "telegram:123456" key := "telegram:123456"
sm.GetOrCreate(key) sm.GetOrCreate(key)
sm.AddMessage(key, "user", "hello") sm.AddMessage(key, "user", "hello")
// Save should succeed even though the key contains ':' // Save should succeed even though the key contains ':'
if err := sm.Save(key); err != nil { if err := sm.Save(key); err != nil {
t.Fatalf("Save(%q) failed: %v", key, err) t.Fatalf("Save(%q) failed: %v", key, err)
} }
// The file on disk should use sanitized name. // The file on disk should use sanitized name.
expectedFile := filepath.Join(tmpDir, "telegram_123456.json") expectedFile := filepath.Join(tmpDir, "telegram_123456.json")
if _, err := os.Stat(expectedFile); os.IsNotExist(err) { if _, err := os.Stat(expectedFile); os.IsNotExist(err) {
t.Fatalf("expected session file %s to exist", expectedFile) t.Fatalf("expected session file %s to exist", expectedFile)
} }
// Load into a fresh manager and verify the session round-trips. // Load into a fresh manager and verify the session round-trips.
sm2 := NewSessionManager(tmpDir) sm2 := NewSessionManager(tmpDir)
history := sm2.GetHistory(key) history := sm2.GetHistory(key)
if len(history) != 1 { if len(history) != 1 {
t.Fatalf("expected 1 message after reload, got %d", len(history)) t.Fatalf("expected 1 message after reload, got %d", len(history))
} }
if history[0].Content != "hello" { if history[0].Content != "hello" {
t.Errorf("expected message content %q, got %q", "hello", history[0].Content) 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) { func TestSanitizeHistory_OrphanedToolCall(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
{Role: "user", Content: "hello"}, {Role: "user", Content: "hello"},
{Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{
{ID: "call_1", Name: "exec"}, {ID: "call_1", Name: "exec"},
{ID: "call_2", Name: "list_dir"}, {ID: "call_2", Name: "list_dir"},
}}, }},
{Role: "tool", Content: "ok", ToolCallID: "call_1"}, {Role: "tool", Content: "ok", ToolCallID: "call_1"},
// Missing tool result for call_2 → orphaned // Missing tool result for call_2 → orphaned
} }
sanitized, removed := SanitizeHistory(history) sanitized, removed := SanitizeHistory(history)
if removed == 0 { if removed == 0 {
t.Fatal("expected orphaned messages to be removed") t.Fatal("expected orphaned messages to be removed")
} }
// After sanitization, only the user message should remain // After sanitization, only the user message should remain
if len(sanitized) != 1 || sanitized[0].Role != "user" { if len(sanitized) != 1 || sanitized[0].Role != "user" {
t.Errorf("expected [user], got %d messages", len(sanitized)) 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) { func TestSanitizeHistory_InterleavedMessages(t *testing.T) {
// Simulates session collision: a user message got interleaved between // Simulates session collision: a user message got interleaved between
// an assistant tool call and its tool result // an assistant tool call and its tool result
history := []providers.Message{ history := []providers.Message{
{Role: "user", Content: "first"}, {Role: "user", Content: "first"},
{Role: "assistant", Content: "ok", ToolCalls: []providers.ToolCall{ {Role: "assistant", Content: "ok", ToolCalls: []providers.ToolCall{
{ID: "call_1", Name: "exec"}, {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: "tool", Content: "ok", ToolCallID: "call_1"}, // ← out of order
{Role: "assistant", Content: "done"}, {Role: "assistant", Content: "done"},
} }
sanitized, removed := SanitizeHistory(history) sanitized, removed := SanitizeHistory(history)
if removed == 0 { if removed == 0 {
t.Fatal("expected interleaved messages to be removed") t.Fatal("expected interleaved messages to be removed")
} }
// Should keep: user("first"), user("collision!"), assistant("done") // Should keep: user("first"), user("collision!"), assistant("done")
// Should remove: assistant(call_1), tool(call_1) // Should remove: assistant(call_1), tool(call_1)
if len(sanitized) != 3 { if len(sanitized) != 3 {
t.Errorf("expected 3 messages, got %d", len(sanitized)) t.Errorf("expected 3 messages, got %d", len(sanitized))
for i, m := range sanitized { for i, m := range sanitized {
t.Logf(" [%d] role=%s content=%q", i, m.Role, m.Content) 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) { func TestSanitizeHistory_CleanHistory(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
{Role: "user", Content: "hello"}, {Role: "user", Content: "hello"},
{Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{
{ID: "call_1", Name: "exec"}, {ID: "call_1", Name: "exec"},
}}, }},
{Role: "tool", Content: "ok", ToolCallID: "call_1"}, {Role: "tool", Content: "ok", ToolCallID: "call_1"},
{Role: "assistant", Content: "done"}, {Role: "assistant", Content: "done"},
} }
sanitized, removed := SanitizeHistory(history) sanitized, removed := SanitizeHistory(history)
if removed != 0 { if removed != 0 {
t.Errorf("expected 0 removed, got %d", removed) t.Errorf("expected 0 removed, got %d", removed)
} }
if len(sanitized) != 4 { if len(sanitized) != 4 {
t.Errorf("expected 4 messages, got %d", len(sanitized)) 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) { func TestSanitizeHistory_MultipleToolCalls(t *testing.T) {
history := []providers.Message{ history := []providers.Message{
{Role: "user", Content: "hello"}, {Role: "user", Content: "hello"},
{Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{
{ID: "call_1", Name: "exec"}, {ID: "call_1", Name: "exec"},
{ID: "call_2", Name: "read_file"}, {ID: "call_2", Name: "read_file"},
}}, }},
{Role: "tool", Content: "ok", ToolCallID: "call_1"}, {Role: "tool", Content: "ok", ToolCallID: "call_1"},
{Role: "tool", Content: "content", ToolCallID: "call_2"}, {Role: "tool", Content: "content", ToolCallID: "call_2"},
{Role: "assistant", Content: "all done"}, {Role: "assistant", Content: "all done"},
} }
sanitized, removed := SanitizeHistory(history) sanitized, removed := SanitizeHistory(history)
if removed != 0 { if removed != 0 {
t.Errorf("expected 0 removed, got %d", removed) t.Errorf("expected 0 removed, got %d", removed)
} }
if len(sanitized) != 5 { if len(sanitized) != 5 {
t.Errorf("expected 5 messages, got %d", len(sanitized)) 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) { func TestSanitizeHistory_Empty(t *testing.T) {
sanitized, removed := SanitizeHistory(nil) sanitized, removed := SanitizeHistory(nil)
if removed != 0 || sanitized != nil { if removed != 0 || sanitized != nil {
t.Errorf("expected nil/0, got %v/%d", sanitized, removed) 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) { func TestSave_RejectsPathTraversal(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
sm := NewSessionManager(tmpDir) sm := NewSessionManager(tmpDir)
badKeys := []string{"", ".", "..", "foo/bar", "foo\\bar"} badKeys := []string{"", ".", "..", "foo/bar", "foo\\bar"}
for _, key := range badKeys { for _, key := range badKeys {
sm.GetOrCreate(key) sm.GetOrCreate(key)
if err := sm.Save(key); err == nil { if err := sm.Save(key); err == nil {
t.Errorf("Save(%q) should have failed but didn't", key) t.Errorf("Save(%q) should have failed but didn't", key)
} }

View file

@ -14,54 +14,104 @@ const sqliteDriver = "sqlite"
const schema = ` const schema = `
CREATE TABLE IF NOT EXISTS sessions ( CREATE TABLE IF NOT EXISTS sessions (
key TEXT PRIMARY KEY, key TEXT PRIMARY KEY,
parent_key TEXT NOT NULL DEFAULT '', parent_key TEXT NOT NULL DEFAULT '',
fork_turn_id TEXT NOT NULL DEFAULT '', fork_turn_id TEXT NOT NULL DEFAULT '',
status TEXT NOT NULL DEFAULT 'active', status TEXT NOT NULL DEFAULT 'active',
label TEXT NOT NULL DEFAULT '', label TEXT NOT NULL DEFAULT '',
summary TEXT NOT NULL DEFAULT '', summary TEXT NOT NULL DEFAULT '',
created_at TEXT NOT NULL, created_at TEXT NOT NULL,
updated_at TEXT NOT NULL updated_at TEXT NOT NULL
); );
CREATE TABLE IF NOT EXISTS turns ( CREATE TABLE IF NOT EXISTS turns (
id TEXT PRIMARY KEY, id TEXT PRIMARY KEY,
session_key TEXT NOT NULL REFERENCES sessions(key) ON DELETE CASCADE, session_key TEXT NOT NULL REFERENCES sessions(key) ON DELETE CASCADE,
seq INTEGER NOT NULL, seq INTEGER NOT NULL,
kind INTEGER NOT NULL DEFAULT 0, kind INTEGER NOT NULL DEFAULT 0,
messages TEXT NOT NULL DEFAULT '[]', messages TEXT NOT NULL DEFAULT '[]',
origin_key TEXT NOT NULL DEFAULT '', origin_key TEXT NOT NULL DEFAULT '',
summary TEXT NOT NULL DEFAULT '', summary TEXT NOT NULL DEFAULT '',
author TEXT NOT NULL DEFAULT '', author TEXT NOT NULL DEFAULT '',
created_at TEXT NOT NULL, created_at TEXT NOT NULL,
meta TEXT NOT NULL DEFAULT '{}' 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_turns_session_seq ON turns(session_key, seq);
CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_key); CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_key);
` `
// SQLiteStore implements SessionStore backed by a single SQLite file. // 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) `INSERT INTO sessions (key, parent_key, fork_turn_id, label, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?)`, VALUES (?, ?, ?, ?, ?, ?)`,
key, parentKey, forkTurnID, label, now, now, 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 `SELECT key, parent_key, fork_turn_id, status, label, summary, created_at, updated_at
FROM sessions WHERE key = ?`, key, 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) `INSERT INTO turns (id, session_key, seq, kind, messages, origin_key, summary, author, created_at, meta)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
turn.ID, sessionKey, turn.Seq, int(turn.Kind), 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 `SELECT id, session_key, seq, kind, messages, origin_key, summary, author, created_at, meta
FROM turns WHERE session_key = ? AND seq > ? ORDER BY seq`, FROM turns WHERE session_key = ? AND seq > ? ORDER BY seq`,
sessionKey, sinceSeq, 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 `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`, FROM turns WHERE session_key = ? ORDER BY seq DESC LIMIT 1`,
sessionKey, sessionKey,

View file

@ -11,16 +11,20 @@ import (
const ( const (
bgWatchPollInterval = 100 * time.Millisecond bgWatchPollInterval = 100 * time.Millisecond
bgWatchDefaultTimeout = 30 * time.Second bgWatchDefaultTimeout = 30 * time.Second
bgTailDefaultLines = 20 bgTailDefaultLines = 20
) )
// BgMonitorTool monitors and inspects background processes managed by ExecTool. // BgMonitorTool monitors and inspects background processes managed by ExecTool.
type BgMonitorTool struct { type BgMonitorTool struct {
exec *ExecTool exec *ExecTool
} }
// NewBgMonitorTool creates a new BgMonitorTool that accesses bg processes from the given ExecTool. // NewBgMonitorTool creates a new BgMonitorTool that accesses bg processes from the given ExecTool.
func NewBgMonitorTool(exec *ExecTool) *BgMonitorTool { func NewBgMonitorTool(exec *ExecTool) *BgMonitorTool {
return &BgMonitorTool{exec: exec} return &BgMonitorTool{exec: exec}
} }
@ -36,77 +40,109 @@ func (t *BgMonitorTool) Description() string {
func (t *BgMonitorTool) Parameters() map[string]any { func (t *BgMonitorTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"action": map[string]any{ "action": map[string]any{
"type": "string", "type": "string",
"enum": []string{"list", "watch", "tail"}, "enum": []string{"list", "watch", "tail"},
"description": "Action: 'list' all bg processes, 'watch' for a pattern in output, 'tail' recent output lines.", "description": "Action: 'list' all bg processes, 'watch' for a pattern in output, 'tail' recent output lines.",
}, },
"bg_id": map[string]any{ "bg_id": map[string]any{
"type": "string", "type": "string",
"description": "Background process ID (e.g. 'bg-1'). Required for watch and tail.", "description": "Background process ID (e.g. 'bg-1'). Required for watch and tail.",
}, },
"pattern": map[string]any{ "pattern": map[string]any{
"type": "string", "type": "string",
"description": "Regex pattern to watch for in output (used with action='watch').", "description": "Regex pattern to watch for in output (used with action='watch').",
}, },
"lines": map[string]any{ "lines": map[string]any{
"type": "number", "type": "number",
"description": "Number of recent lines to return (used with action='tail', default 20).", "description": "Number of recent lines to return (used with action='tail', default 20).",
}, },
"watch_timeout": map[string]any{ "watch_timeout": map[string]any{
"type": "number", "type": "number",
"description": "Timeout in seconds for watch action (default 30).", "description": "Timeout in seconds for watch action (default 30).",
}, },
}, },
"required": []string{"action"}, "required": []string{"action"},
} }
} }
func (t *BgMonitorTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *BgMonitorTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
action, _ := args["action"].(string) action, _ := args["action"].(string)
switch action { switch action {
case "list": case "list":
return t.actionList() return t.actionList()
case "watch": case "watch":
return t.actionWatch(ctx, args) return t.actionWatch(ctx, args)
case "tail": case "tail":
return t.actionTail(args) return t.actionTail(args)
default: default:
return ErrorResult(fmt.Sprintf("unknown action %q (use 'list', 'watch', or 'tail')", action)) return ErrorResult(fmt.Sprintf("unknown action %q (use 'list', 'watch', or 'tail')", action))
} }
} }
func (t *BgMonitorTool) actionList() *ToolResult { func (t *BgMonitorTool) actionList() *ToolResult {
procs := t.exec.BgProcesses() procs := t.exec.BgProcesses()
if len(procs) == 0 { if len(procs) == 0 {
return &ToolResult{ return &ToolResult{
ForLLM: "No background processes.", ForLLM: "No background processes.",
ForUser: "No background processes.", ForUser: "No background processes.",
} }
} }
ids := make([]string, 0, len(procs)) ids := make([]string, 0, len(procs))
for id := range procs { for id := range procs {
ids = append(ids, id) ids = append(ids, id)
} }
sort.Strings(ids) sort.Strings(ids)
var sb strings.Builder var sb strings.Builder
sb.WriteString("Background Processes:\n\n") sb.WriteString("Background Processes:\n\n")
for _, id := range ids { for _, id := range ids {
bp := procs[id] bp := procs[id]
if bp.isRunning() { if bp.isRunning() {
uptime := time.Since(bp.startedAt).Truncate(time.Second) uptime := time.Since(bp.startedAt).Truncate(time.Second)
fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n", fmt.Fprintf(&sb, " [%s] pid=%d running (uptime: %s, max: %s) %s\n",
id, bp.pid, uptime, getBgMaxLifetime(), bp.command) id, bp.pid, uptime, getBgMaxLifetime(), bp.command)
} else { } else {
ran := time.Since(bp.startedAt).Truncate(time.Second) ran := time.Since(bp.startedAt).Truncate(time.Second)
if bp.exitErr != nil { if bp.exitErr != nil {
fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n", fmt.Fprintf(&sb, " [%s] pid=%d exited=err (ran: %s) %s\n",
id, bp.pid, ran, bp.command) id, bp.pid, ran, bp.command)
} else { } else {
fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n", fmt.Fprintf(&sb, " [%s] pid=%d exited=0 (ran: %s) %s\n",
id, bp.pid, ran, bp.command) id, bp.pid, ran, bp.command)
} }
} }
@ -114,17 +150,20 @@ func (t *BgMonitorTool) actionList() *ToolResult {
return &ToolResult{ return &ToolResult{
ForLLM: sb.String(), ForLLM: sb.String(),
ForUser: sb.String(), ForUser: sb.String(),
} }
} }
func (t *BgMonitorTool) actionWatch(ctx context.Context, args map[string]any) *ToolResult { func (t *BgMonitorTool) actionWatch(ctx context.Context, args map[string]any) *ToolResult {
bgID, _ := args["bg_id"].(string) bgID, _ := args["bg_id"].(string)
if bgID == "" { if bgID == "" {
return ErrorResult("bg_id is required for watch action") return ErrorResult("bg_id is required for watch action")
} }
patternStr, _ := args["pattern"].(string) patternStr, _ := args["pattern"].(string)
if patternStr == "" { if patternStr == "" {
return ErrorResult("pattern is required for watch action") 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 timeout := bgWatchDefaultTimeout
if t, ok := args["watch_timeout"].(float64); ok && t > 0 { if t, ok := args["watch_timeout"].(float64); ok && t > 0 {
timeout = time.Duration(t) * time.Second timeout = time.Duration(t) * time.Second
} }
procs := t.exec.BgProcesses() procs := t.exec.BgProcesses()
bp, ok := procs[bgID] bp, ok := procs[bgID]
if !ok { if !ok {
return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) return ErrorResult(fmt.Sprintf("background process %q not found", bgID))
} }
deadline := time.After(timeout) deadline := time.After(timeout)
ticker := time.NewTicker(bgWatchPollInterval) ticker := time.NewTicker(bgWatchPollInterval)
defer ticker.Stop() defer ticker.Stop()
for { for {
// Check for pattern match // Check for pattern match
if match := bp.output.Match(pattern); match != "" { if match := bp.output.Match(pattern); match != "" {
return &ToolResult{ 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), ForUser: fmt.Sprintf("Match found in [%s]: %s", bgID, match),
} }
} }
// Check if process exited // Check if process exited
if !bp.isRunning() { if !bp.isRunning() {
output := bp.output.String() output := bp.output.String()
tail := lastNLines(output, 10) tail := lastNLines(output, 10)
var sb strings.Builder var sb strings.Builder
fmt.Fprintf(&sb, "Process %s exited before pattern matched.\n", bgID) fmt.Fprintf(&sb, "Process %s exited before pattern matched.\n", bgID)
if bp.exitErr != nil { if bp.exitErr != nil {
fmt.Fprintf(&sb, "Exit: %v\n", bp.exitErr) fmt.Fprintf(&sb, "Exit: %v\n", bp.exitErr)
} else { } else {
fmt.Fprintf(&sb, "Exit: 0\n") fmt.Fprintf(&sb, "Exit: 0\n")
} }
fmt.Fprintf(&sb, "\nLast output:\n%s", tail) fmt.Fprintf(&sb, "\nLast output:\n%s", tail)
return &ToolResult{ return &ToolResult{
ForLLM: sb.String(), ForLLM: sb.String(),
ForUser: sb.String(), ForUser: sb.String(),
IsError: true, IsError: true,
} }
} }
select { select {
case <-deadline: case <-deadline:
// Timeout // Timeout
output := bp.output.String() output := bp.output.String()
tail := lastNLines(output, 10) tail := lastNLines(output, 10)
var sb strings.Builder 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, "Watch timed out after %s waiting for pattern %q in [%s].\n", timeout, patternStr, bgID)
fmt.Fprintf(&sb, "\nLast output:\n%s", tail) fmt.Fprintf(&sb, "\nLast output:\n%s", tail)
return &ToolResult{ return &ToolResult{
ForLLM: sb.String(), ForLLM: sb.String(),
ForUser: sb.String(), ForUser: sb.String(),
IsError: true, IsError: true,
} }
case <-ctx.Done(): case <-ctx.Done():
return ErrorResult("watch canceled") return ErrorResult("watch canceled")
case <-ticker.C: case <-ticker.C:
// Continue polling // 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 { func (t *BgMonitorTool) actionTail(args map[string]any) *ToolResult {
bgID, _ := args["bg_id"].(string) bgID, _ := args["bg_id"].(string)
if bgID == "" { if bgID == "" {
return ErrorResult("bg_id is required for tail action") return ErrorResult("bg_id is required for tail action")
} }
n := bgTailDefaultLines n := bgTailDefaultLines
if lines, ok := args["lines"].(float64); ok && lines > 0 { if lines, ok := args["lines"].(float64); ok && lines > 0 {
n = int(lines) n = int(lines)
} }
procs := t.exec.BgProcesses() procs := t.exec.BgProcesses()
bp, ok := procs[bgID] bp, ok := procs[bgID]
if !ok { if !ok {
return ErrorResult(fmt.Sprintf("background process %q not found", bgID)) 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) lines := bp.output.Lines(n)
var sb strings.Builder var sb strings.Builder
fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command) fmt.Fprintf(&sb, "[%s] pid=%d %s\n", bp.id, bp.pid, bp.command)
if bp.isRunning() { if bp.isRunning() {
fmt.Fprintf(&sb, "Status: running\n") fmt.Fprintf(&sb, "Status: running\n")
} else { } else {
@ -228,7 +302,9 @@ func (t *BgMonitorTool) actionTail(args map[string]any) *ToolResult {
fmt.Fprintf(&sb, "Status: exited=0\n") fmt.Fprintf(&sb, "Status: exited=0\n")
} }
} }
fmt.Fprintf(&sb, "\nLast %d lines:\n", n) fmt.Fprintf(&sb, "\nLast %d lines:\n", n)
for _, line := range lines { for _, line := range lines {
fmt.Fprintf(&sb, "%s\n", line) fmt.Fprintf(&sb, "%s\n", line)
} }
@ -239,18 +315,23 @@ func (t *BgMonitorTool) actionTail(args map[string]any) *ToolResult {
return &ToolResult{ return &ToolResult{
ForLLM: sb.String(), ForLLM: sb.String(),
ForUser: sb.String(), ForUser: sb.String(),
} }
} }
// lastNLines returns the last n lines from a string. // lastNLines returns the last n lines from a string.
func lastNLines(s string, n int) string { func lastNLines(s string, n int) string {
lines := strings.Split(s, "\n") lines := strings.Split(s, "\n")
if len(lines) > 0 && lines[len(lines)-1] == "" { if len(lines) > 0 && lines[len(lines)-1] == "" {
lines = lines[:len(lines)-1] lines = lines[:len(lines)-1]
} }
if n >= len(lines) { if n >= len(lines) {
return strings.Join(lines, "\n") return strings.Join(lines, "\n")
} }
return strings.Join(lines[len(lines)-n:], "\n") return strings.Join(lines[len(lines)-n:], "\n")
} }

View file

@ -10,64 +10,83 @@ import (
func TestBgMonitor_List(t *testing.T) { func TestBgMonitor_List(t *testing.T) {
tool, _ := NewExecTool("", false) tool, _ := NewExecTool("", false)
monitor := NewBgMonitorTool(tool) monitor := NewBgMonitorTool(tool)
// List with no processes // List with no processes
result := monitor.Execute(context.Background(), map[string]any{"action": "list"}) result := monitor.Execute(context.Background(), map[string]any{"action": "list"})
if result.IsError { if result.IsError {
t.Fatalf("unexpected error: %s", result.ForLLM) t.Fatalf("unexpected error: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "No background") { if !strings.Contains(result.ForLLM, "No background") {
t.Errorf("expected 'No background' message, got: %s", result.ForLLM) t.Errorf("expected 'No background' message, got: %s", result.ForLLM)
} }
// Start two bg processes // Start two bg processes
var cmd1, cmd2 string var cmd1, cmd2 string
if runtime.GOOS == "windows" { if runtime.GOOS == "windows" {
cmd1 = "Start-Sleep -Seconds 30" cmd1 = "Start-Sleep -Seconds 30"
cmd2 = "Start-Sleep -Seconds 30" cmd2 = "Start-Sleep -Seconds 30"
} else { } else {
cmd1 = "sleep 30" cmd1 = "sleep 30"
cmd2 = "sleep 30" cmd2 = "sleep 30"
} }
r1 := tool.Execute(context.Background(), map[string]any{ r1 := tool.Execute(context.Background(), map[string]any{
"command": cmd1, "command": cmd1,
"background": true, "background": true,
}) })
if r1.IsError { if r1.IsError {
t.Fatalf("failed to start bg-1: %s", r1.ForLLM) t.Fatalf("failed to start bg-1: %s", r1.ForLLM)
} }
r2 := tool.Execute(context.Background(), map[string]any{ r2 := tool.Execute(context.Background(), map[string]any{
"command": cmd2, "command": cmd2,
"background": true, "background": true,
}) })
if r2.IsError { if r2.IsError {
t.Fatalf("failed to start bg-2: %s", r2.ForLLM) t.Fatalf("failed to start bg-2: %s", r2.ForLLM)
} }
// List should show both // List should show both
result = monitor.Execute(context.Background(), map[string]any{"action": "list"}) result = monitor.Execute(context.Background(), map[string]any{"action": "list"})
if result.IsError { if result.IsError {
t.Fatalf("unexpected error: %s", result.ForLLM) t.Fatalf("unexpected error: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "bg-1") { if !strings.Contains(result.ForLLM, "bg-1") {
t.Errorf("expected bg-1 in list, got: %s", result.ForLLM) t.Errorf("expected bg-1 in list, got: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "bg-2") { if !strings.Contains(result.ForLLM, "bg-2") {
t.Errorf("expected bg-2 in list, got: %s", result.ForLLM) t.Errorf("expected bg-2 in list, got: %s", result.ForLLM)
} }
// Cleanup // Cleanup
tool.Shutdown() tool.Shutdown()
} }
func TestBgMonitor_Watch_Match(t *testing.T) { func TestBgMonitor_Watch_Match(t *testing.T) {
tool, _ := NewExecTool("", false) tool, _ := NewExecTool("", false)
monitor := NewBgMonitorTool(tool) monitor := NewBgMonitorTool(tool)
var cmd string var cmd string
if runtime.GOOS == "windows" { if runtime.GOOS == "windows" {
cmd = "Write-Output 'Server ready on port 3000'; Start-Sleep -Seconds 30" cmd = "Write-Output 'Server ready on port 3000'; Start-Sleep -Seconds 30"
} else { } else {
@ -76,25 +95,34 @@ func TestBgMonitor_Watch_Match(t *testing.T) {
r := tool.Execute(context.Background(), map[string]any{ r := tool.Execute(context.Background(), map[string]any{
"command": cmd, "command": cmd,
"background": true, "background": true,
}) })
if r.IsError { if r.IsError {
t.Fatalf("failed to start bg: %s", r.ForLLM) t.Fatalf("failed to start bg: %s", r.ForLLM)
} }
// Watch for "ready" pattern — should match quickly // Watch for "ready" pattern — should match quickly
result := monitor.Execute(context.Background(), map[string]any{ result := monitor.Execute(context.Background(), map[string]any{
"action": "watch", "action": "watch",
"bg_id": "bg-1", "bg_id": "bg-1",
"pattern": "ready", "pattern": "ready",
"watch_timeout": float64(10), "watch_timeout": float64(10),
}) })
if result.IsError { if result.IsError {
t.Fatalf("expected watch to match, got error: %s", result.ForLLM) t.Fatalf("expected watch to match, got error: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "Match found") { if !strings.Contains(result.ForLLM, "Match found") {
t.Errorf("expected 'Match found' message, got: %s", result.ForLLM) t.Errorf("expected 'Match found' message, got: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "ready") { if !strings.Contains(result.ForLLM, "ready") {
t.Errorf("expected match to contain 'ready', got: %s", result.ForLLM) 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) { func TestBgMonitor_Watch_Timeout(t *testing.T) {
tool, _ := NewExecTool("", false) tool, _ := NewExecTool("", false)
monitor := NewBgMonitorTool(tool) monitor := NewBgMonitorTool(tool)
var cmd string var cmd string
if runtime.GOOS == "windows" { if runtime.GOOS == "windows" {
cmd = "Start-Sleep -Seconds 30" cmd = "Start-Sleep -Seconds 30"
} else { } else {
@ -115,22 +145,30 @@ func TestBgMonitor_Watch_Timeout(t *testing.T) {
r := tool.Execute(context.Background(), map[string]any{ r := tool.Execute(context.Background(), map[string]any{
"command": cmd, "command": cmd,
"background": true, "background": true,
}) })
if r.IsError { if r.IsError {
t.Fatalf("failed to start bg: %s", r.ForLLM) t.Fatalf("failed to start bg: %s", r.ForLLM)
} }
// Watch for a pattern that won't appear, with short timeout // Watch for a pattern that won't appear, with short timeout
result := monitor.Execute(context.Background(), map[string]any{ result := monitor.Execute(context.Background(), map[string]any{
"action": "watch", "action": "watch",
"bg_id": "bg-1", "bg_id": "bg-1",
"pattern": "never_going_to_match", "pattern": "never_going_to_match",
"watch_timeout": float64(1), "watch_timeout": float64(1),
}) })
if !result.IsError { if !result.IsError {
t.Fatalf("expected watch to timeout with error, got success: %s", result.ForLLM) t.Fatalf("expected watch to timeout with error, got success: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "timed out") { if !strings.Contains(result.ForLLM, "timed out") {
t.Errorf("expected 'timed out' message, got: %s", result.ForLLM) 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) { func TestBgMonitor_Watch_ProcessExit(t *testing.T) {
tool, _ := NewExecTool("", false) tool, _ := NewExecTool("", false)
monitor := NewBgMonitorTool(tool) monitor := NewBgMonitorTool(tool)
var cmd string var cmd string
if runtime.GOOS == "windows" { if runtime.GOOS == "windows" {
cmd = "Write-Output 'done quickly'" cmd = "Write-Output 'done quickly'"
} else { } else {
@ -151,25 +191,34 @@ func TestBgMonitor_Watch_ProcessExit(t *testing.T) {
r := tool.Execute(context.Background(), map[string]any{ r := tool.Execute(context.Background(), map[string]any{
"command": cmd, "command": cmd,
"background": true, "background": true,
}) })
if r.IsError { if r.IsError {
t.Fatalf("failed to start bg: %s", r.ForLLM) t.Fatalf("failed to start bg: %s", r.ForLLM)
} }
// Wait a bit for the process to exit // Wait a bit for the process to exit
time.Sleep(4 * time.Second) time.Sleep(4 * time.Second)
// Watch for a pattern that doesn't match — process should have exited // Watch for a pattern that doesn't match — process should have exited
result := monitor.Execute(context.Background(), map[string]any{ result := monitor.Execute(context.Background(), map[string]any{
"action": "watch", "action": "watch",
"bg_id": "bg-1", "bg_id": "bg-1",
"pattern": "never_match", "pattern": "never_match",
"watch_timeout": float64(5), "watch_timeout": float64(5),
}) })
if !result.IsError { if !result.IsError {
t.Fatalf("expected error when process exits, got: %s", result.ForLLM) t.Fatalf("expected error when process exits, got: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "exited") { if !strings.Contains(result.ForLLM, "exited") {
t.Errorf("expected 'exited' message, got: %s", result.ForLLM) 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) { func TestBgMonitor_Tail(t *testing.T) {
tool, _ := NewExecTool("", false) tool, _ := NewExecTool("", false)
monitor := NewBgMonitorTool(tool) monitor := NewBgMonitorTool(tool)
var cmd string var cmd string
if runtime.GOOS == "windows" { if runtime.GOOS == "windows" {
cmd = "1..5 | ForEach-Object { Write-Output \"line $_\" }; Start-Sleep -Seconds 30" cmd = "1..5 | ForEach-Object { Write-Output \"line $_\" }; Start-Sleep -Seconds 30"
} else { } else {
@ -190,24 +241,32 @@ func TestBgMonitor_Tail(t *testing.T) {
r := tool.Execute(context.Background(), map[string]any{ r := tool.Execute(context.Background(), map[string]any{
"command": cmd, "command": cmd,
"background": true, "background": true,
}) })
if r.IsError { if r.IsError {
t.Fatalf("failed to start bg: %s", r.ForLLM) t.Fatalf("failed to start bg: %s", r.ForLLM)
} }
// Wait for initial output to be captured // Wait for initial output to be captured
time.Sleep(4 * time.Second) time.Sleep(4 * time.Second)
// Tail last 3 lines // Tail last 3 lines
result := monitor.Execute(context.Background(), map[string]any{ result := monitor.Execute(context.Background(), map[string]any{
"action": "tail", "action": "tail",
"bg_id": "bg-1", "bg_id": "bg-1",
"lines": float64(3), "lines": float64(3),
}) })
if result.IsError { if result.IsError {
t.Fatalf("unexpected error: %s", result.ForLLM) t.Fatalf("unexpected error: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "line 5") { if !strings.Contains(result.ForLLM, "line 5") {
t.Errorf("expected tail to contain 'line 5', got: %s", result.ForLLM) 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) { func TestBgMonitor_InvalidAction(t *testing.T) {
tool, _ := NewExecTool("", false) tool, _ := NewExecTool("", false)
monitor := NewBgMonitorTool(tool) monitor := NewBgMonitorTool(tool)
result := monitor.Execute(context.Background(), map[string]any{"action": "invalid"}) result := monitor.Execute(context.Background(), map[string]any{"action": "invalid"})
if !result.IsError { if !result.IsError {
t.Fatalf("expected error for invalid action") t.Fatalf("expected error for invalid action")
} }
if !strings.Contains(result.ForLLM, "unknown action") { if !strings.Contains(result.ForLLM, "unknown action") {
t.Errorf("expected 'unknown action' message, got: %s", result.ForLLM) t.Errorf("expected 'unknown action' message, got: %s", result.ForLLM)
} }

View file

@ -10,27 +10,42 @@ import (
const ( const (
ciPollInterval = 30 * time.Second ciPollInterval = 30 * time.Second
ciPollTimeout = 15 * time.Minute ciPollTimeout = 15 * time.Minute
) )
// CreatePRTool creates a GitHub pull request from the current worktree branch. // CreatePRTool creates a GitHub pull request from the current worktree branch.
// //
// Safety invariants: // Safety invariants:
// - Only works inside a worktree (WorktreeInfo must be in context) // - Only works inside a worktree (WorktreeInfo must be in context)
// - Base branch is auto-detected from WorktreeInfo.BaseBranch // - Base branch is auto-detected from WorktreeInfo.BaseBranch
// - Requires the branch to be already pushed (use git_push first) // - Requires the branch to be already pushed (use git_push first)
// - Checks for merge conflicts with base before creating // - Checks for merge conflicts with base before creating
// - Uses `gh pr create` under the hood // - Uses `gh pr create` under the hood
// //
// Async behavior: // Async behavior:
// - PR creation itself is synchronous and returns immediately with the PR URL // - PR creation itself is synchronous and returns immediately with the PR URL
// - If CI runs are triggered, a background goroutine polls `gh pr checks` // - If CI runs are triggered, a background goroutine polls `gh pr checks`
// and calls the AsyncCallback when CI completes (pass or fail) // and calls the AsyncCallback when CI completes (pass or fail)
type CreatePRTool struct { type CreatePRTool struct {
callback AsyncCallback callback AsyncCallback
} }
// NewCreatePRTool creates a CreatePRTool. // NewCreatePRTool creates a CreatePRTool.
func NewCreatePRTool() *CreatePRTool { func NewCreatePRTool() *CreatePRTool {
return &CreatePRTool{} return &CreatePRTool{}
} }
@ -38,122 +53,187 @@ func NewCreatePRTool() *CreatePRTool {
func (t *CreatePRTool) Name() string { return "create_pr" } func (t *CreatePRTool) Name() string { return "create_pr" }
// SetCallback implements AsyncTool for CI completion notification. // SetCallback implements AsyncTool for CI completion notification.
func (t *CreatePRTool) SetCallback(cb AsyncCallback) { func (t *CreatePRTool) SetCallback(cb AsyncCallback) {
t.callback = cb t.callback = cb
} }
func (t *CreatePRTool) Description() string { func (t *CreatePRTool) Description() string {
return "Create a GitHub pull request from the current worktree branch. " + return "Create a GitHub pull request from the current worktree branch. " +
"The base branch is auto-detected from the worktree's parent branch. " + "The base branch is auto-detected from the worktree's parent branch. " +
"The branch must be pushed to origin first (use git_push). " + "The branch must be pushed to origin first (use git_push). " +
"Checks for merge conflicts with the base branch before creating. " + "Checks for merge conflicts with the base branch before creating. " +
"After PR creation, polls CI status in the background and notifies when complete. " + "After PR creation, polls CI status in the background and notifies when complete. " +
"Requires the `gh` CLI to be installed and authenticated." "Requires the `gh` CLI to be installed and authenticated."
} }
func (t *CreatePRTool) Parameters() map[string]any { func (t *CreatePRTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"title": map[string]any{ "title": map[string]any{
"type": "string", "type": "string",
"description": "Pull request title", "description": "Pull request title",
}, },
"body": map[string]any{ "body": map[string]any{
"type": "string", "type": "string",
"description": "Pull request body/description (supports markdown)", "description": "Pull request body/description (supports markdown)",
}, },
"draft": map[string]any{ "draft": map[string]any{
"type": "boolean", "type": "boolean",
"description": "Create as draft PR (default: false)", "description": "Create as draft PR (default: false)",
}, },
}, },
"required": []string{"title"}, "required": []string{"title"},
} }
} }
func (t *CreatePRTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *CreatePRTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
wt := WorktreeInfoFromCtx(ctx) wt := WorktreeInfoFromCtx(ctx)
if wt == nil { if wt == nil {
return ErrorResult( return ErrorResult(
"create_pr requires an active worktree.\n" + "create_pr requires an active worktree.\n" +
"This tool can only be used during worktree-based sessions " + "This tool can only be used during worktree-based sessions " +
"(e.g., heartbeat tasks or plan executing phase).\n" + "(e.g., heartbeat tasks or plan executing phase).\n" +
"The worktree provides the branch name and base branch for the PR.") "The worktree provides the branch name and base branch for the PR.")
} }
branch := wt.Branch branch := wt.Branch
if branch == "" { if branch == "" {
return ErrorResult( return ErrorResult(
"worktree has no branch name.\n" + "worktree has no branch name.\n" +
"The WorktreeInfo was set but Branch is empty. " + "The WorktreeInfo was set but Branch is empty. " +
"This is an internal error — the worktree may not have been created correctly.") "This is an internal error — the worktree may not have been created correctly.")
} }
baseBranch := wt.BaseBranch baseBranch := wt.BaseBranch
if baseBranch == "" { if baseBranch == "" {
baseBranch = "main" baseBranch = "main"
} }
title, ok := args["title"].(string) title, ok := args["title"].(string)
if !ok || strings.TrimSpace(title) == "" { if !ok || strings.TrimSpace(title) == "" {
return ErrorResult( return ErrorResult(
"title is required.\n" + "title is required.\n" +
"Provide a concise PR title describing the change (e.g., \"Add rate limiter to API endpoints\").") "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 // Verify the branch has been pushed by checking if the remote ref exists
checkCtx, checkCancel := context.WithTimeout(ctx, 15*time.Second) checkCtx, checkCancel := context.WithTimeout(ctx, 15*time.Second)
defer checkCancel() defer checkCancel()
checkCmd := exec.CommandContext(checkCtx, "git", "ls-remote", "--exit-code", "origin", branch) checkCmd := exec.CommandContext(checkCtx, "git", "ls-remote", "--exit-code", "origin", branch)
checkCmd.Dir = wt.Path checkCmd.Dir = wt.Path
if err := checkCmd.Run(); err != nil { if err := checkCmd.Run(); err != nil {
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"branch %q not found on origin.\n"+ "branch %q not found on origin.\n"+
"The branch must be pushed before creating a PR. Use the git_push tool first.\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.", "git_push will auto-commit uncommitted changes and push the worktree branch to origin.",
branch)) branch))
} }
// Fetch latest base branch and check for merge conflicts // Fetch latest base branch and check for merge conflicts
fetchCtx, fetchCancel := context.WithTimeout(ctx, 30*time.Second) fetchCtx, fetchCancel := context.WithTimeout(ctx, 30*time.Second)
defer fetchCancel() defer fetchCancel()
fetchCmd := exec.CommandContext(fetchCtx, "git", "fetch", "origin", baseBranch) fetchCmd := exec.CommandContext(fetchCtx, "git", "fetch", "origin", baseBranch)
fetchCmd.Dir = wt.Path fetchCmd.Dir = wt.Path
if out, err := fetchCmd.CombinedOutput(); err != nil { if out, err := fetchCmd.CombinedOutput(); err != nil {
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"failed to fetch origin/%s: %s\n%s\n"+ "failed to fetch origin/%s: %s\n%s\n"+
"Cannot verify merge compatibility without the latest base branch. "+ "Cannot verify merge compatibility without the latest base branch. "+
"Check network connectivity and that the base branch %q exists on origin.", "Check network connectivity and that the base branch %q exists on origin.",
baseBranch, err, strings.TrimSpace(string(out)), baseBranch)) baseBranch, err, strings.TrimSpace(string(out)), baseBranch))
} }
// Try a merge dry-run to detect conflicts. // Try a merge dry-run to detect conflicts.
// merge-tree --write-tree is a plumbing command (Git 2.38+) that performs a // 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. // three-way merge entirely in-memory without touching the working tree.
// Exit code 0 = clean merge, non-zero = conflicts detected. // Exit code 0 = clean merge, non-zero = conflicts detected.
mergeCtx, mergeCancel := context.WithTimeout(ctx, 30*time.Second) mergeCtx, mergeCancel := context.WithTimeout(ctx, 30*time.Second)
defer mergeCancel() defer mergeCancel()
mergeCmd := exec.CommandContext(mergeCtx, "git", "merge-tree", mergeCmd := exec.CommandContext(mergeCtx, "git", "merge-tree",
"--write-tree", "--no-messages", "--write-tree", "--no-messages",
branch, "origin/"+baseBranch) branch, "origin/"+baseBranch)
mergeCmd.Dir = wt.RepoRoot mergeCmd.Dir = wt.RepoRoot
mergeOut, mergeErr := mergeCmd.CombinedOutput() mergeOut, mergeErr := mergeCmd.CombinedOutput()
if mergeErr != nil { if mergeErr != nil {
conflictInfo := strings.TrimSpace(string(mergeOut)) conflictInfo := strings.TrimSpace(string(mergeOut))
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"merge conflict detected between %q and %s.\n"+ "merge conflict detected between %q and %s.\n"+
"The PR cannot be created cleanly. Resolve the conflicts in the worktree first, "+ "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"+ "then use git_push to push the resolution before retrying create_pr.\n"+
"Conflict details:\n%s", "Conflict details:\n%s",
branch, baseBranch, conflictInfo)) branch, baseBranch, conflictInfo))
} }
// Build gh pr create command // Build gh pr create command
ghArgs := []string{ ghArgs := []string{
"pr", "create", "pr", "create",
"--base", baseBranch, "--base", baseBranch,
"--head", branch, "--head", branch,
"--title", title, "--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) prCtx, prCancel := context.WithTimeout(ctx, 30*time.Second)
defer prCancel() defer prCancel()
cmd := exec.CommandContext(prCtx, "gh", ghArgs...) cmd := exec.CommandContext(prCtx, "gh", ghArgs...)
cmd.Dir = wt.RepoRoot cmd.Dir = wt.RepoRoot
out, err := cmd.CombinedOutput() out, err := cmd.CombinedOutput()
output := strings.TrimSpace(string(out)) output := strings.TrimSpace(string(out))
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"gh pr create failed: %s\n%s\n"+ "gh pr create failed: %s\n%s\n"+
"Possible causes:\n"+ "Possible causes:\n"+
"- gh CLI not installed or not authenticated (run `gh auth login`)\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"+ "- A PR already exists for branch %q (check with `gh pr list`)\n"+
"- Repository not configured as a GitHub remote", "- Repository not configured as a GitHub remote",
err, output, branch)) err, output, branch))
} }
prURL := output // gh pr create outputs the PR URL prURL := output // gh pr create outputs the PR URL
// Start background CI polling if callback is set // Start background CI polling if callback is set
if t.callback != nil && prURL != "" { if t.callback != nil && prURL != "" {
cb := t.callback cb := t.callback
repoRoot := wt.RepoRoot repoRoot := wt.RepoRoot
go pollCIStatus(repoRoot, prURL, cb) go pollCIStatus(repoRoot, prURL, cb)
} }
return AsyncResult(fmt.Sprintf( return AsyncResult(fmt.Sprintf(
"Pull request created: %s\n"+ "Pull request created: %s\n"+
"Branch: %s -> %s\n"+ "Branch: %s -> %s\n"+
"CI status will be reported asynchronously when checks complete.", "CI status will be reported asynchronously when checks complete.",
prURL, branch, baseBranch)) prURL, branch, baseBranch))
} }
// pollCIStatus polls `gh pr checks` in the background until all checks // pollCIStatus polls `gh pr checks` in the background until all checks
// pass, fail, or the timeout is reached. Reports back via AsyncCallback. // pass, fail, or the timeout is reached. Reports back via AsyncCallback.
func pollCIStatus(repoRoot, prURL string, callback AsyncCallback) { func pollCIStatus(repoRoot, prURL string, callback AsyncCallback) {
// Detached context with hard timeout — this goroutine outlives the tool call. // Detached context with hard timeout — this goroutine outlives the tool call.
ctx, cancel := context.WithTimeout(context.Background(), ciPollTimeout) ctx, cancel := context.WithTimeout(context.Background(), ciPollTimeout)
defer cancel() defer cancel()
// Initial wait: CI runs take a few seconds to register after PR creation // Initial wait: CI runs take a few seconds to register after PR creation
select { select {
case <-time.After(10 * time.Second): case <-time.After(10 * time.Second):
case <-ctx.Done(): case <-ctx.Done():
return return
} }
ticker := time.NewTicker(ciPollInterval) ticker := time.NewTicker(ciPollInterval)
defer ticker.Stop() defer ticker.Stop()
for { for {
status, detail := checkPRChecks(ctx, repoRoot, prURL) status, detail := checkPRChecks(ctx, repoRoot, prURL)
switch status { switch status {
case ciStatusPass: case ciStatusPass:
callback(ctx, NewToolResult(fmt.Sprintf( callback(ctx, NewToolResult(fmt.Sprintf(
"CI passed for %s\n%s", "CI passed for %s\n%s",
prURL, detail))) prURL, detail)))
return return
case ciStatusFail: case ciStatusFail:
callback(ctx, &ToolResult{ callback(ctx, &ToolResult{
ForLLM: fmt.Sprintf( ForLLM: fmt.Sprintf(
"CI failed for %s\n%s\n"+ "CI failed for %s\n%s\n"+
"Run `gh run view` for detailed logs.", "Run `gh run view` for detailed logs.",
prURL, detail), prURL, detail),
IsError: true, IsError: true,
}) })
return return
case ciStatusNone: case ciStatusNone:
callback(ctx, NewToolResult(fmt.Sprintf( callback(ctx, NewToolResult(fmt.Sprintf(
"No CI checks configured for %s. PR is ready for review.", "No CI checks configured for %s. PR is ready for review.",
prURL))) prURL)))
return return
case ciStatusPending: case ciStatusPending:
// Still running, continue polling // Still running, continue polling
} }
select { select {
case <-ticker.C: case <-ticker.C:
case <-ctx.Done(): case <-ctx.Done():
callback(ctx, &ToolResult{ callback(ctx, &ToolResult{
ForLLM: fmt.Sprintf( ForLLM: fmt.Sprintf(
"CI polling timed out after %s for %s.\n"+ "CI polling timed out after %s for %s.\n"+
"Checks may still be running. Run `gh pr checks %s` to check.", "Checks may still be running. Run `gh pr checks %s` to check.",
ciPollTimeout, prURL, prURL), ciPollTimeout, prURL, prURL),
IsError: true, IsError: true,
}) })
return return
} }
} }
@ -263,36 +394,51 @@ type ciStatus int
const ( const (
ciStatusPending ciStatus = iota ciStatusPending ciStatus = iota
ciStatusPass ciStatusPass
ciStatusFail ciStatusFail
ciStatusNone ciStatusNone
) )
// checkPRChecks runs `gh pr checks` and parses the result. // checkPRChecks runs `gh pr checks` and parses the result.
// Returns the aggregate status and raw output for the caller to include. // Returns the aggregate status and raw output for the caller to include.
func checkPRChecks(ctx context.Context, repoRoot, prURL string) (ciStatus, string) { func checkPRChecks(ctx context.Context, repoRoot, prURL string) (ciStatus, string) {
checkCtx, cancel := context.WithTimeout(ctx, 15*time.Second) checkCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
defer cancel() defer cancel()
cmd := exec.CommandContext(checkCtx, "gh", "pr", "checks", prURL) cmd := exec.CommandContext(checkCtx, "gh", "pr", "checks", prURL)
cmd.Dir = repoRoot cmd.Dir = repoRoot
out, err := cmd.CombinedOutput() out, err := cmd.CombinedOutput()
output := strings.TrimSpace(string(out)) output := strings.TrimSpace(string(out))
if err != nil { if err != nil {
// gh pr checks exits 1 when any check has failed // gh pr checks exits 1 when any check has failed
if strings.Contains(output, "fail") || strings.Contains(output, "X ") { if strings.Contains(output, "fail") || strings.Contains(output, "X ") {
return ciStatusFail, output return ciStatusFail, output
} }
// "no checks" case // "no checks" case
if strings.Contains(output, "no checks") || output == "" { if strings.Contains(output, "no checks") || output == "" {
return ciStatusNone, "" return ciStatusNone, ""
} }
// Transient error or still pending — keep polling // Transient error or still pending — keep polling
return ciStatusPending, output return ciStatusPending, output
} }
// Exit 0: all checks completed. Check for pending. // Exit 0: all checks completed. Check for pending.
if strings.Contains(output, "pending") || strings.Contains(output, "- ") { if strings.Contains(output, "pending") || strings.Contains(output, "- ") {
return ciStatusPending, output return ciStatusPending, output
} }

View file

@ -9,191 +9,259 @@ import (
) )
// TestCreatePRTool_NoWorktree verifies that create_pr fails without worktree context. // TestCreatePRTool_NoWorktree verifies that create_pr fails without worktree context.
func TestCreatePRTool_NoWorktree(t *testing.T) { func TestCreatePRTool_NoWorktree(t *testing.T) {
tool := NewCreatePRTool() tool := NewCreatePRTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"title": "Test PR", "title": "Test PR",
}) })
if !result.IsError { if !result.IsError {
t.Fatal("expected error when no worktree in context") t.Fatal("expected error when no worktree in context")
} }
assertContains(t, result.ForLLM, "worktree") assertContains(t, result.ForLLM, "worktree")
assertContains(t, result.ForLLM, "heartbeat") assertContains(t, result.ForLLM, "heartbeat")
} }
// TestCreatePRTool_EmptyBranch verifies that empty branch name is rejected. // TestCreatePRTool_EmptyBranch verifies that empty branch name is rejected.
func TestCreatePRTool_EmptyBranch(t *testing.T) { func TestCreatePRTool_EmptyBranch(t *testing.T) {
tool := NewCreatePRTool() tool := NewCreatePRTool()
ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{
Branch: "", Branch: "",
BaseBranch: "main", BaseBranch: "main",
Path: t.TempDir(), Path: t.TempDir(),
RepoRoot: t.TempDir(), RepoRoot: t.TempDir(),
}) })
result := tool.Execute(ctx, map[string]any{ result := tool.Execute(ctx, map[string]any{
"title": "Test PR", "title": "Test PR",
}) })
if !result.IsError { if !result.IsError {
t.Fatal("expected error for empty branch") t.Fatal("expected error for empty branch")
} }
assertContains(t, result.ForLLM, "no branch name") assertContains(t, result.ForLLM, "no branch name")
} }
// TestCreatePRTool_MissingTitle verifies that missing title is rejected. // TestCreatePRTool_MissingTitle verifies that missing title is rejected.
func TestCreatePRTool_MissingTitle(t *testing.T) { func TestCreatePRTool_MissingTitle(t *testing.T) {
tool := NewCreatePRTool() tool := NewCreatePRTool()
ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{
Branch: "plan/test", Branch: "plan/test",
BaseBranch: "main", BaseBranch: "main",
Path: t.TempDir(), Path: t.TempDir(),
RepoRoot: t.TempDir(), RepoRoot: t.TempDir(),
}) })
tests := []struct { tests := []struct {
name string name string
args map[string]any args map[string]any
}{ }{
{"no title key", map[string]any{}}, {"no title key", map[string]any{}},
{"empty title", map[string]any{"title": ""}}, {"empty title", map[string]any{"title": ""}},
{"whitespace title", map[string]any{"title": " "}}, {"whitespace title", map[string]any{"title": " "}},
} }
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
result := tool.Execute(ctx, tt.args) result := tool.Execute(ctx, tt.args)
if !result.IsError { if !result.IsError {
t.Fatal("expected error for missing/empty title") t.Fatal("expected error for missing/empty title")
} }
assertContains(t, result.ForLLM, "title is required") assertContains(t, result.ForLLM, "title is required")
}) })
} }
} }
// TestCreatePRTool_BranchNotPushed verifies the tool checks for remote branch existence. // TestCreatePRTool_BranchNotPushed verifies the tool checks for remote branch existence.
func TestCreatePRTool_BranchNotPushed(t *testing.T) { func TestCreatePRTool_BranchNotPushed(t *testing.T) {
tool := NewCreatePRTool() tool := NewCreatePRTool()
ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{
Branch: "plan/not-pushed", Branch: "plan/not-pushed",
BaseBranch: "main", BaseBranch: "main",
Path: t.TempDir(), Path: t.TempDir(),
RepoRoot: t.TempDir(), RepoRoot: t.TempDir(),
}) })
result := tool.Execute(ctx, map[string]any{ result := tool.Execute(ctx, map[string]any{
"title": "Test PR", "title": "Test PR",
}) })
if !result.IsError { if !result.IsError {
t.Fatal("expected error for unpushed branch") t.Fatal("expected error for unpushed branch")
} }
// Should mention git_push as the remedy // Should mention git_push as the remedy
assertContains(t, result.ForLLM, "git_push") assertContains(t, result.ForLLM, "git_push")
} }
// TestCreatePRTool_DefaultBaseBranch verifies fallback to "main" when BaseBranch is empty. // TestCreatePRTool_DefaultBaseBranch verifies fallback to "main" when BaseBranch is empty.
func TestCreatePRTool_DefaultBaseBranch(t *testing.T) { func TestCreatePRTool_DefaultBaseBranch(t *testing.T) {
tool := NewCreatePRTool() tool := NewCreatePRTool()
// With empty BaseBranch, tool should default to "main" // With empty BaseBranch, tool should default to "main"
ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{
Branch: "plan/test", Branch: "plan/test",
BaseBranch: "", BaseBranch: "",
Path: t.TempDir(), Path: t.TempDir(),
RepoRoot: t.TempDir(), RepoRoot: t.TempDir(),
}) })
result := tool.Execute(ctx, map[string]any{ result := tool.Execute(ctx, map[string]any{
"title": "Test PR", "title": "Test PR",
}) })
// Will fail at ls-remote (no real repo), but should not fail at baseBranch validation // Will fail at ls-remote (no real repo), but should not fail at baseBranch validation
if result.IsError && strings.Contains(result.ForLLM, "base branch") { if result.IsError && strings.Contains(result.ForLLM, "base branch") {
t.Fatal("should not fail on base branch when defaulting to main") t.Fatal("should not fail on base branch when defaulting to main")
} }
} }
// TestCreatePRTool_Interface verifies the tool satisfies both Tool and AsyncTool interfaces. // TestCreatePRTool_Interface verifies the tool satisfies both Tool and AsyncTool interfaces.
func TestCreatePRTool_Interface(t *testing.T) { func TestCreatePRTool_Interface(t *testing.T) {
var _ Tool = (*CreatePRTool)(nil) var _ Tool = (*CreatePRTool)(nil)
var _ AsyncTool = (*CreatePRTool)(nil) var _ AsyncTool = (*CreatePRTool)(nil)
tool := NewCreatePRTool() tool := NewCreatePRTool()
if tool.Name() != "create_pr" { if tool.Name() != "create_pr" {
t.Errorf("Name: got %q, want %q", tool.Name(), "create_pr") t.Errorf("Name: got %q, want %q", tool.Name(), "create_pr")
} }
if tool.Description() == "" { if tool.Description() == "" {
t.Error("Description should not be empty") t.Error("Description should not be empty")
} }
params := tool.Parameters() params := tool.Parameters()
if params == nil { if params == nil {
t.Fatal("Parameters should not be nil") t.Fatal("Parameters should not be nil")
} }
// Verify "title" is required // Verify "title" is required
required, ok := params["required"].([]string) required, ok := params["required"].([]string)
if !ok { if !ok {
t.Fatal("required should be []string") t.Fatal("required should be []string")
} }
foundTitle := false foundTitle := false
for _, r := range required { for _, r := range required {
if r == "title" { if r == "title" {
foundTitle = true foundTitle = true
} }
} }
if !foundTitle { if !foundTitle {
t.Error("title should be in required parameters") t.Error("title should be in required parameters")
} }
} }
// TestCreatePRTool_SetCallback verifies callback is stored. // TestCreatePRTool_SetCallback verifies callback is stored.
func TestCreatePRTool_SetCallback(t *testing.T) { func TestCreatePRTool_SetCallback(t *testing.T) {
tool := NewCreatePRTool() tool := NewCreatePRTool()
if tool.callback != nil { if tool.callback != nil {
t.Fatal("callback should be nil initially") t.Fatal("callback should be nil initially")
} }
called := false called := false
tool.SetCallback(func(ctx context.Context, result *ToolResult) { tool.SetCallback(func(ctx context.Context, result *ToolResult) {
called = true called = true
}) })
if tool.callback == nil { if tool.callback == nil {
t.Fatal("callback should be set after SetCallback") t.Fatal("callback should be set after SetCallback")
} }
// Verify it's callable (doesn't panic) // Verify it's callable (doesn't panic)
tool.callback(context.Background(), NewToolResult("test")) tool.callback(context.Background(), NewToolResult("test"))
if !called { if !called {
t.Fatal("callback was not invoked") t.Fatal("callback was not invoked")
} }
} }
// TestCheckPRChecks_ParseResults tests CI status parsing logic. // TestCheckPRChecks_ParseResults tests CI status parsing logic.
func TestCheckPRChecks_ParseResults(t *testing.T) { func TestCheckPRChecks_ParseResults(t *testing.T) {
// This tests the parsing logic conceptually — actual `gh` calls // This tests the parsing logic conceptually — actual `gh` calls
// would need integration tests. We verify the status constants exist // would need integration tests. We verify the status constants exist
// and the type is usable. // and the type is usable.
if ciStatusPending != 0 { if ciStatusPending != 0 {
t.Error("ciStatusPending should be 0 (default)") t.Error("ciStatusPending should be 0 (default)")
} }
if ciStatusPass == ciStatusFail { if ciStatusPass == ciStatusFail {
t.Error("ciStatusPass and ciStatusFail should differ") t.Error("ciStatusPass and ciStatusFail should differ")
} }
if ciStatusNone == ciStatusPending { if ciStatusNone == ciStatusPending {
t.Error("ciStatusNone and ciStatusPending should differ") t.Error("ciStatusNone and ciStatusPending should differ")
} }
} }
// TestAllowedToolsForPreset_GitTools checks git tools are correctly assigned to presets. // TestAllowedToolsForPreset_GitTools checks git tools are correctly assigned to presets.
func TestAllowedToolsForPreset_GitTools(t *testing.T) { func TestAllowedToolsForPreset_GitTools(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
preset Preset preset Preset
wantGitPush bool wantGitPush bool
wantCreatePR bool wantCreatePR bool
}{ }{
{"scout", PresetScout, false, false}, {"scout", PresetScout, false, false},
{"analyst", PresetAnalyst, false, false}, {"analyst", PresetAnalyst, false, false},
{"coder", PresetCoder, true, false}, {"coder", PresetCoder, true, false},
{"worker", PresetWorker, true, true}, {"worker", PresetWorker, true, true},
{"coordinator", PresetCoordinator, true, true}, {"coordinator", PresetCoordinator, true, true},
} }
@ -204,6 +272,7 @@ func TestAllowedToolsForPreset_GitTools(t *testing.T) {
if got := allowed["git_push"]; got != tt.wantGitPush { if got := allowed["git_push"]; got != tt.wantGitPush {
t.Errorf("git_push: got %v, want %v", got, tt.wantGitPush) t.Errorf("git_push: got %v, want %v", got, tt.wantGitPush)
} }
if got := allowed["create_pr"]; got != tt.wantCreatePR { if got := allowed["create_pr"]; got != tt.wantCreatePR {
t.Errorf("create_pr: got %v, want %v", 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 // JobExecutor is the interface for executing cron jobs through the agent
type JobExecutor interface { type JobExecutor interface {
ProcessDirectWithChannel(ctx context.Context, content, sessionKey, channel, chatID string) (string, error) ProcessDirectWithChannel(ctx context.Context, content, sessionKey, channel, chatID string) (string, error)
} }
// CronTool provides scheduling capabilities for the agent // CronTool provides scheduling capabilities for the agent
type CronTool struct { type CronTool struct {
cronService *cron.CronService cronService *cron.CronService
executor JobExecutor executor JobExecutor
msgBus *bus.MessageBus msgBus *bus.MessageBus
execTool *ExecTool execTool *ExecTool
channel string channel string
chatID string chatID string
mu sync.RWMutex mu sync.RWMutex
} }
// NewCronTool creates a new CronTool // NewCronTool creates a new CronTool
// execTimeout: 0 means no timeout, >0 sets the timeout duration // execTimeout: 0 means no timeout, >0 sets the timeout duration
func NewCronTool( func NewCronTool(
cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool, cronService *cron.CronService, executor JobExecutor, msgBus *bus.MessageBus, workspace string, restrict bool,
execTimeout time.Duration, config *config.Config, execTimeout time.Duration, config *config.Config,
) (*CronTool, error) { ) (*CronTool, error) {
execTool, err := NewExecToolWithConfig(workspace, restrict, config) execTool, err := NewExecToolWithConfig(workspace, restrict, config)
@ -41,102 +52,147 @@ func NewCronTool(
} }
execTool.SetTimeout(execTimeout) execTool.SetTimeout(execTimeout)
return &CronTool{ return &CronTool{
cronService: cronService, cronService: cronService,
executor: executor, executor: executor,
msgBus: msgBus, msgBus: msgBus,
execTool: execTool, execTool: execTool,
}, nil }, nil
} }
// Name returns the tool name // Name returns the tool name
func (t *CronTool) Name() string { func (t *CronTool) Name() string {
return "cron" return "cron"
} }
// Description returns the tool description // Description returns the tool description
func (t *CronTool) Description() string { 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." 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 // Parameters returns the tool parameters schema
func (t *CronTool) Parameters() map[string]any { func (t *CronTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"action": map[string]any{ "action": map[string]any{
"type": "string", "type": "string",
"enum": []string{"add", "list", "remove", "enable", "disable"}, "enum": []string{"add", "list", "remove", "enable", "disable"},
"description": "Action to perform. Use 'add' when user wants to schedule a reminder or task.", "description": "Action to perform. Use 'add' when user wants to schedule a reminder or task.",
}, },
"message": map[string]any{ "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.", "description": "The reminder/task message to display when triggered. If 'command' is used, this describes what the command does.",
}, },
"command": map[string]any{ "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.", "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{ "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'.", "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{ "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'.", "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{ "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.", "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{ "job_id": map[string]any{
"type": "string", "type": "string",
"description": "Job ID (for remove/enable/disable)", "description": "Job ID (for remove/enable/disable)",
}, },
"deliver": map[string]any{ "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", "description": "If true, send message directly to channel. If false, let agent process message (for complex tasks). Default: true",
}, },
}, },
"required": []string{"action"}, "required": []string{"action"},
} }
} }
// SetContext sets the current session context for job creation // SetContext sets the current session context for job creation
func (t *CronTool) SetContext(channel, chatID string) { func (t *CronTool) SetContext(channel, chatID string) {
t.mu.Lock() t.mu.Lock()
defer t.mu.Unlock() defer t.mu.Unlock()
t.channel = channel t.channel = channel
t.chatID = chatID t.chatID = chatID
} }
// Execute runs the tool with the given arguments // Execute runs the tool with the given arguments
func (t *CronTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *CronTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
action, ok := args["action"].(string) action, ok := args["action"].(string)
if !ok { if !ok {
return ErrorResult("action is required") return ErrorResult("action is required")
} }
switch action { switch action {
case "add": case "add":
return t.addJob(args) return t.addJob(args)
case "list": case "list":
return t.listJobs() return t.listJobs()
case "remove": case "remove":
return t.removeJob(args) return t.removeJob(args)
case "enable": case "enable":
return t.enableJob(args, true) return t.enableJob(args, true)
case "disable": case "disable":
return t.enableJob(args, false) return t.enableJob(args, false)
default: default:
return ErrorResult(fmt.Sprintf("unknown action: %s", action)) return ErrorResult(fmt.Sprintf("unknown action: %s", action))
} }
} }
func (t *CronTool) addJob(args map[string]any) *ToolResult { func (t *CronTool) addJob(args map[string]any) *ToolResult {
t.mu.RLock() t.mu.RLock()
channel := t.channel channel := t.channel
chatID := t.chatID chatID := t.chatID
t.mu.RUnlock() t.mu.RUnlock()
if channel == "" || chatID == "" { if channel == "" || chatID == "" {
@ -144,6 +200,7 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult {
} }
message, ok := args["message"].(string) message, ok := args["message"].(string)
if !ok || message == "" { if !ok || message == "" {
return ErrorResult("message is required for add") return ErrorResult("message is required for add")
} }
@ -151,26 +208,35 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult {
var schedule cron.CronSchedule var schedule cron.CronSchedule
// Check for at_seconds (one-time), every_seconds (recurring), or cron_expr // Check for at_seconds (one-time), every_seconds (recurring), or cron_expr
atSeconds, hasAt := args["at_seconds"].(float64) atSeconds, hasAt := args["at_seconds"].(float64)
everySeconds, hasEvery := args["every_seconds"].(float64) everySeconds, hasEvery := args["every_seconds"].(float64)
cronExpr, hasCron := args["cron_expr"].(string) cronExpr, hasCron := args["cron_expr"].(string)
// Priority: at_seconds > every_seconds > cron_expr // Priority: at_seconds > every_seconds > cron_expr
if hasAt { if hasAt {
atMS := time.Now().UnixMilli() + int64(atSeconds)*1000 atMS := time.Now().UnixMilli() + int64(atSeconds)*1000
schedule = cron.CronSchedule{ schedule = cron.CronSchedule{
Kind: "at", Kind: "at",
AtMS: &atMS, AtMS: &atMS,
} }
} else if hasEvery { } else if hasEvery {
everyMS := int64(everySeconds) * 1000 everyMS := int64(everySeconds) * 1000
schedule = cron.CronSchedule{ schedule = cron.CronSchedule{
Kind: "every", Kind: "every",
EveryMS: &everyMS, EveryMS: &everyMS,
} }
} else if hasCron { } else if hasCron {
schedule = cron.CronSchedule{ schedule = cron.CronSchedule{
Kind: "cron", Kind: "cron",
Expr: cronExpr, Expr: cronExpr,
} }
} else { } else {
@ -178,29 +244,43 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult {
} }
// Read deliver parameter, default to true // Read deliver parameter, default to true
deliver := true deliver := true
if d, ok := args["deliver"].(bool); ok { if d, ok := args["deliver"].(bool); ok {
deliver = d deliver = d
} }
command, _ := args["command"].(string) command, _ := args["command"].(string)
if command != "" { if command != "" {
// Commands must be processed by agent/exec tool, so deliver must be false (or handled specifically) // 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 // 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. // 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. // However, logically, it's not "delivered" to chat directly as is.
deliver = false deliver = false
} }
// Truncate message for job name (max 30 chars) // Truncate message for job name (max 30 chars)
messagePreview := utils.Truncate(message, 30) messagePreview := utils.Truncate(message, 30)
job, err := t.cronService.AddJob( job, err := t.cronService.AddJob(
messagePreview, messagePreview,
schedule, schedule,
message, message,
deliver, deliver,
channel, channel,
chatID, chatID,
) )
if err != nil { if err != nil {
@ -209,7 +289,9 @@ func (t *CronTool) addJob(args map[string]any) *ToolResult {
if command != "" { if command != "" {
job.Payload.Command = command job.Payload.Command = command
// Need to save the updated payload // Need to save the updated payload
t.cronService.UpdateJob(job) t.cronService.UpdateJob(job)
} }
@ -224,9 +306,12 @@ func (t *CronTool) listJobs() *ToolResult {
} }
var sb strings.Builder var sb strings.Builder
sb.WriteString("Scheduled jobs:\n") sb.WriteString("Scheduled jobs:\n")
for _, j := range jobs { for _, j := range jobs {
var scheduleInfo string var scheduleInfo string
if j.Schedule.Kind == "every" && j.Schedule.EveryMS != nil { if j.Schedule.Kind == "every" && j.Schedule.EveryMS != nil {
scheduleInfo = fmt.Sprintf("every %ds", *j.Schedule.EveryMS/1000) scheduleInfo = fmt.Sprintf("every %ds", *j.Schedule.EveryMS/1000)
} else if j.Schedule.Kind == "cron" { } else if j.Schedule.Kind == "cron" {
@ -236,6 +321,7 @@ func (t *CronTool) listJobs() *ToolResult {
} else { } else {
scheduleInfo = "unknown" scheduleInfo = "unknown"
} }
fmt.Fprintf(&sb, "- %s (id: %s, %s)\n", j.Name, j.ID, scheduleInfo) 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 { func (t *CronTool) removeJob(args map[string]any) *ToolResult {
jobID, ok := args["job_id"].(string) jobID, ok := args["job_id"].(string)
if !ok || jobID == "" { if !ok || jobID == "" {
return ErrorResult("job_id is required for remove") 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) { if t.cronService.RemoveJob(jobID) {
return SilentResult(fmt.Sprintf("Cron job removed: %s", jobID)) return SilentResult(fmt.Sprintf("Cron job removed: %s", jobID))
} }
return ErrorResult(fmt.Sprintf("Job %s not found", jobID)) return ErrorResult(fmt.Sprintf("Job %s not found", jobID))
} }
func (t *CronTool) enableJob(args map[string]any, enable bool) *ToolResult { func (t *CronTool) enableJob(args map[string]any, enable bool) *ToolResult {
jobID, ok := args["job_id"].(string) jobID, ok := args["job_id"].(string)
if !ok || jobID == "" { if !ok || jobID == "" {
return ErrorResult("job_id is required for enable/disable") return ErrorResult("job_id is required for enable/disable")
} }
job := t.cronService.EnableJob(jobID, enable) job := t.cronService.EnableJob(jobID, enable)
if job == nil { if job == nil {
return ErrorResult(fmt.Sprintf("Job %s not found", jobID)) return ErrorResult(fmt.Sprintf("Job %s not found", jobID))
} }
status := "enabled" status := "enabled"
if !enable { if !enable {
status = "disabled" status = "disabled"
} }
return SilentResult(fmt.Sprintf("Cron job '%s' %s", job.Name, status)) return SilentResult(fmt.Sprintf("Cron job '%s' %s", job.Name, status))
} }
// ExecuteJob executes a cron job through the agent // ExecuteJob executes a cron job through the agent
func (t *CronTool) ExecuteJob(ctx context.Context, job *cron.CronJob) string { func (t *CronTool) ExecuteJob(ctx context.Context, job *cron.CronJob) string {
// Get channel/chatID from job payload // Get channel/chatID from job payload
channel := job.Payload.Channel channel := job.Payload.Channel
chatID := job.Payload.To chatID := job.Payload.To
// Default values if not set // Default values if not set
if channel == "" { if channel == "" {
channel = "cli" channel = "cli"
} }
if chatID == "" { if chatID == "" {
chatID = "direct" chatID = "direct"
} }
// Execute command if present // Execute command if present
if job.Payload.Command != "" { if job.Payload.Command != "" {
args := map[string]any{ args := map[string]any{
"command": job.Payload.Command, "command": job.Payload.Command,
} }
result := t.execTool.Execute(ctx, args) result := t.execTool.Execute(ctx, args)
var output string var output string
if result.IsError { if result.IsError {
output = fmt.Sprintf("Error executing scheduled command: %s", result.ForLLM) output = fmt.Sprintf("Error executing scheduled command: %s", result.ForLLM)
} else { } 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) pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer pubCancel() defer pubCancel()
t.msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{ t.msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{
Channel: channel, Channel: channel,
ChatID: chatID, ChatID: chatID,
Content: output, Content: output,
}) })
return "ok" return "ok"
} }
// If deliver=true, send message directly without agent processing // If deliver=true, send message directly without agent processing
if job.Payload.Deliver { if job.Payload.Deliver {
pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second) pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer pubCancel() defer pubCancel()
t.msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{ t.msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{
Channel: channel, Channel: channel,
ChatID: chatID, ChatID: chatID,
Content: job.Payload.Message, Content: job.Payload.Message,
}) })
return "ok" return "ok"
} }
// For deliver=false, process through agent (for complex tasks) // For deliver=false, process through agent (for complex tasks)
sessionKey := fmt.Sprintf("cron-%s", job.ID) sessionKey := fmt.Sprintf("cron-%s", job.ID)
// Call agent with job's message // Call agent with job's message
response, err := t.executor.ProcessDirectWithChannel( response, err := t.executor.ProcessDirectWithChannel(
ctx, ctx,
job.Payload.Message, job.Payload.Message,
sessionKey, sessionKey,
channel, channel,
chatID, chatID,
) )
if err != nil { 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 is automatically sent via MessageBus by AgentLoop
_ = response // Will be sent by AgentLoop _ = response // Will be sent by AgentLoop
return "ok" return "ok"
} }

View file

@ -10,11 +10,13 @@ import (
) )
// DevPreviewTool allows the agent to control the Mini App dev reverse proxy. // DevPreviewTool allows the agent to control the Mini App dev reverse proxy.
type DevPreviewTool struct { type DevPreviewTool struct {
manager miniapp.DevTargetManager manager miniapp.DevTargetManager
} }
// NewDevPreviewTool creates a new DevPreviewTool. // NewDevPreviewTool creates a new DevPreviewTool.
func NewDevPreviewTool(manager miniapp.DevTargetManager) *DevPreviewTool { func NewDevPreviewTool(manager miniapp.DevTargetManager) *DevPreviewTool {
return &DevPreviewTool{manager: manager} return &DevPreviewTool{manager: manager}
} }
@ -28,113 +30,157 @@ func (t *DevPreviewTool) Description() string {
func (t *DevPreviewTool) Parameters() map[string]any { func (t *DevPreviewTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"action": map[string]any{ "action": map[string]any{
"type": "string", "type": "string",
"enum": []string{"start", "stop", "unregister", "status"}, "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).", "description": "Action to perform: start (register + activate target), stop (deactivate proxy), unregister (remove a registered target), status (list all targets).",
}, },
"target": map[string]any{ "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.", "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{ "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.", "description": "Display name for the target (e.g. 'frontend'). Optional for 'start' action; auto-generated from host:port if omitted.",
}, },
"id": map[string]any{ "id": map[string]any{
"type": "string", "type": "string",
"description": "Target ID. Required for 'unregister' action.", "description": "Target ID. Required for 'unregister' action.",
}, },
}, },
"required": []string{"action"}, "required": []string{"action"},
} }
} }
func (t *DevPreviewTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *DevPreviewTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
action, ok := args["action"].(string) action, ok := args["action"].(string)
if !ok { if !ok {
return ErrorResult("action is required") return ErrorResult("action is required")
} }
switch action { switch action {
case "start": case "start":
target, _ := args["target"].(string) target, _ := args["target"].(string)
if target == "" { if target == "" {
return ErrorResult("target is required for start action") return ErrorResult("target is required for start action")
} }
name, _ := args["name"].(string) name, _ := args["name"].(string)
if name == "" { if name == "" {
name = inferName(target) name = inferName(target)
} }
id, err := t.manager.RegisterDevTarget(name, target) id, err := t.manager.RegisterDevTarget(name, target)
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("failed to register dev target: %v", err)) return ErrorResult(fmt.Sprintf("failed to register dev target: %v", err))
} }
if err := t.manager.ActivateDevTarget(id); err != nil { if err := t.manager.ActivateDevTarget(id); err != nil {
return ErrorResult(fmt.Sprintf("failed to activate dev target: %v", err)) return ErrorResult(fmt.Sprintf("failed to activate dev target: %v", err))
} }
return SilentResult( return SilentResult(
fmt.Sprintf( fmt.Sprintf(
"Dev preview started (id=%s, name=%s). Target: %s\nUsers can view it in the Mini App Dev tab.", "Dev preview started (id=%s, name=%s). Target: %s\nUsers can view it in the Mini App Dev tab.",
id, id,
name, name,
target, target,
), ),
) )
case "stop": case "stop":
if err := t.manager.DeactivateDevTarget(); err != nil { if err := t.manager.DeactivateDevTarget(); err != nil {
return ErrorResult(fmt.Sprintf("failed to stop dev preview: %v", err)) return ErrorResult(fmt.Sprintf("failed to stop dev preview: %v", err))
} }
return SilentResult("Dev preview stopped.") return SilentResult("Dev preview stopped.")
case "unregister": case "unregister":
id, _ := args["id"].(string) id, _ := args["id"].(string)
if id == "" { if id == "" {
return ErrorResult("id is required for unregister action") return ErrorResult("id is required for unregister action")
} }
if err := t.manager.UnregisterDevTarget(id); err != nil { if err := t.manager.UnregisterDevTarget(id); err != nil {
return ErrorResult(fmt.Sprintf("failed to unregister target: %v", err)) return ErrorResult(fmt.Sprintf("failed to unregister target: %v", err))
} }
return SilentResult(fmt.Sprintf("Dev target %s unregistered.", id)) return SilentResult(fmt.Sprintf("Dev target %s unregistered.", id))
case "status": case "status":
targets := t.manager.ListDevTargets() targets := t.manager.ListDevTargets()
active := t.manager.GetDevTarget() active := t.manager.GetDevTarget()
if len(targets) == 0 { if len(targets) == 0 {
if active == "" { if active == "" {
return SilentResult("Dev preview is not active. No targets registered.") return SilentResult("Dev preview is not active. No targets registered.")
} }
return SilentResult(fmt.Sprintf("Dev preview is active. Target: %s\nNo registered targets.", active)) return SilentResult(fmt.Sprintf("Dev preview is active. Target: %s\nNo registered targets.", active))
} }
var sb strings.Builder var sb strings.Builder
if active != "" { if active != "" {
sb.WriteString(fmt.Sprintf("Dev preview is active. Target: %s\n", active)) sb.WriteString(fmt.Sprintf("Dev preview is active. Target: %s\n", active))
} else { } else {
sb.WriteString("Dev preview is not active.\n") sb.WriteString("Dev preview is not active.\n")
} }
sb.WriteString("Registered targets:\n") sb.WriteString("Registered targets:\n")
for _, dt := range targets { for _, dt := range targets {
sb.WriteString(fmt.Sprintf(" [%s] %s → %s\n", dt.ID, dt.Name, dt.Target)) sb.WriteString(fmt.Sprintf(" [%s] %s → %s\n", dt.ID, dt.Name, dt.Target))
} }
return SilentResult(sb.String()) return SilentResult(sb.String())
default: default:
return ErrorResult(fmt.Sprintf("unknown action: %s", action)) return ErrorResult(fmt.Sprintf("unknown action: %s", action))
} }
} }
// inferName generates a display name from a target URL (e.g. "localhost:3000"). // inferName generates a display name from a target URL (e.g. "localhost:3000").
func inferName(target string) string { func inferName(target string) string {
u, err := url.Parse(target) u, err := url.Parse(target)
if err != nil { if err != nil {
return target return target
} }
host := u.Hostname() host := u.Hostname()
port := u.Port() port := u.Port()
if port != "" { if port != "" {
return host + ":" + port return host + ":" + port
} }
return host return host
} }

View file

@ -10,11 +10,16 @@ import (
) )
// mockDevTargetManager implements miniapp.DevTargetManager for testing. // mockDevTargetManager implements miniapp.DevTargetManager for testing.
type mockDevTargetManager struct { type mockDevTargetManager struct {
targets map[string]*miniapp.DevTarget targets map[string]*miniapp.DevTarget
nextID int nextID int
activeID string activeID string
active string // active target URL active string // active target URL
regErr error regErr error
} }
@ -26,9 +31,13 @@ func (m *mockDevTargetManager) RegisterDevTarget(name, target string) (string, e
if m.regErr != nil { if m.regErr != nil {
return "", m.regErr return "", m.regErr
} }
m.nextID++ m.nextID++
id := fmt.Sprintf("%d", m.nextID) id := fmt.Sprintf("%d", m.nextID)
m.targets[id] = &miniapp.DevTarget{ID: id, Name: name, Target: target} m.targets[id] = &miniapp.DevTarget{ID: id, Name: name, Target: target}
return id, nil return id, nil
} }
@ -36,27 +45,37 @@ func (m *mockDevTargetManager) UnregisterDevTarget(id string) error {
if _, ok := m.targets[id]; !ok { if _, ok := m.targets[id]; !ok {
return fmt.Errorf("target %q not found", id) return fmt.Errorf("target %q not found", id)
} }
delete(m.targets, id) delete(m.targets, id)
if m.activeID == id { if m.activeID == id {
m.activeID = "" m.activeID = ""
m.active = "" m.active = ""
} }
return nil return nil
} }
func (m *mockDevTargetManager) ActivateDevTarget(id string) error { func (m *mockDevTargetManager) ActivateDevTarget(id string) error {
dt, ok := m.targets[id] dt, ok := m.targets[id]
if !ok { if !ok {
return fmt.Errorf("target %q not found", id) return fmt.Errorf("target %q not found", id)
} }
m.activeID = id m.activeID = id
m.active = dt.Target m.active = dt.Target
return nil return nil
} }
func (m *mockDevTargetManager) DeactivateDevTarget() error { func (m *mockDevTargetManager) DeactivateDevTarget() error {
m.activeID = "" m.activeID = ""
m.active = "" m.active = ""
return nil return nil
} }
@ -66,34 +85,43 @@ func (m *mockDevTargetManager) GetDevTarget() string {
func (m *mockDevTargetManager) ListDevTargets() []miniapp.DevTarget { func (m *mockDevTargetManager) ListDevTargets() []miniapp.DevTarget {
out := make([]miniapp.DevTarget, 0, len(m.targets)) out := make([]miniapp.DevTarget, 0, len(m.targets))
for _, dt := range m.targets { for _, dt := range m.targets {
out = append(out, *dt) out = append(out, *dt)
} }
return out return out
} }
func TestDevPreviewTool_Start(t *testing.T) { func TestDevPreviewTool_Start(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://localhost:3000", "target": "http://localhost:3000",
"name": "frontend", "name": "frontend",
}) })
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
if len(mgr.targets) != 1 { if len(mgr.targets) != 1 {
t.Errorf("expected 1 registered target, got %d", len(mgr.targets)) t.Errorf("expected 1 registered target, got %d", len(mgr.targets))
} }
if mgr.active != "http://localhost:3000" { if mgr.active != "http://localhost:3000" {
t.Errorf("expected active target http://localhost:3000, got %q", mgr.active) t.Errorf("expected active target http://localhost:3000, got %q", mgr.active)
} }
if !strings.Contains(result.ForLLM, "started") { if !strings.Contains(result.ForLLM, "started") {
t.Errorf("expected result to contain 'started', got %q", result.ForLLM) t.Errorf("expected result to contain 'started', got %q", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "frontend") { if !strings.Contains(result.ForLLM, "frontend") {
t.Errorf("expected result to contain 'frontend', got %q", result.ForLLM) 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) { func TestDevPreviewTool_StartAutoName(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://localhost:3000", "target": "http://localhost:3000",
}) })
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
// Auto-generated name should be "localhost:3000" // Auto-generated name should be "localhost:3000"
for _, dt := range mgr.targets { for _, dt := range mgr.targets {
if dt.Name != "localhost:3000" { if dt.Name != "localhost:3000" {
t.Errorf("expected auto-name 'localhost:3000', got %q", dt.Name) 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) { func TestDevPreviewTool_StartMissingTarget(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -134,11 +167,14 @@ func TestDevPreviewTool_StartMissingTarget(t *testing.T) {
func TestDevPreviewTool_StartError(t *testing.T) { func TestDevPreviewTool_StartError(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
mgr.regErr = fmt.Errorf("only localhost") mgr.regErr = fmt.Errorf("only localhost")
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://example.com:3000", "target": "http://example.com:3000",
}) })
@ -149,7 +185,9 @@ func TestDevPreviewTool_StartError(t *testing.T) {
func TestDevPreviewTool_Stop(t *testing.T) { func TestDevPreviewTool_Stop(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
mgr.active = "http://localhost:3000" mgr.active = "http://localhost:3000"
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -159,6 +197,7 @@ func TestDevPreviewTool_Stop(t *testing.T) {
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
if mgr.active != "" { if mgr.active != "" {
t.Errorf("expected empty active target after stop, got %q", 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) { func TestDevPreviewTool_Unregister(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
// Register a target first // Register a target first
id, _ := mgr.RegisterDevTarget("frontend", "http://localhost:3000") id, _ := mgr.RegisterDevTarget("frontend", "http://localhost:3000")
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "unregister", "action": "unregister",
"id": id, "id": id,
}) })
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
if len(mgr.targets) != 0 { if len(mgr.targets) != 0 {
t.Errorf("expected 0 targets after unregister, got %d", len(mgr.targets)) 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) { func TestDevPreviewTool_UnregisterMissingID(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -199,10 +243,12 @@ func TestDevPreviewTool_UnregisterMissingID(t *testing.T) {
func TestDevPreviewTool_UnregisterNotFound(t *testing.T) { func TestDevPreviewTool_UnregisterNotFound(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "unregister", "action": "unregister",
"id": "999", "id": "999",
}) })
@ -213,10 +259,13 @@ func TestDevPreviewTool_UnregisterNotFound(t *testing.T) {
func TestDevPreviewTool_Status(t *testing.T) { func TestDevPreviewTool_Status(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
mgr.RegisterDevTarget("api", "http://localhost:8080") mgr.RegisterDevTarget("api", "http://localhost:8080")
mgr.RegisterDevTarget("frontend", "http://localhost:3000") mgr.RegisterDevTarget("frontend", "http://localhost:3000")
mgr.active = "http://localhost:8080" mgr.active = "http://localhost:8080"
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -226,15 +275,19 @@ func TestDevPreviewTool_Status(t *testing.T) {
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "active") { if !strings.Contains(result.ForLLM, "active") {
t.Errorf("expected 'active' in result, got %q", result.ForLLM) t.Errorf("expected 'active' in result, got %q", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "http://localhost:8080") { if !strings.Contains(result.ForLLM, "http://localhost:8080") {
t.Errorf("expected target URL in result, got %q", result.ForLLM) t.Errorf("expected target URL in result, got %q", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "api") { if !strings.Contains(result.ForLLM, "api") {
t.Errorf("expected 'api' in result, got %q", result.ForLLM) t.Errorf("expected 'api' in result, got %q", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "frontend") { if !strings.Contains(result.ForLLM, "frontend") {
t.Errorf("expected 'frontend' in result, got %q", result.ForLLM) 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) { func TestDevPreviewTool_StatusInactive(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -251,6 +305,7 @@ func TestDevPreviewTool_StatusInactive(t *testing.T) {
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "not active") { if !strings.Contains(result.ForLLM, "not active") {
t.Errorf("expected 'not active' in result, got %q", result.ForLLM) 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) { func TestDevPreviewTool_UnknownAction(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -271,6 +327,7 @@ func TestDevPreviewTool_UnknownAction(t *testing.T) {
func TestDevPreviewTool_MissingAction(t *testing.T) { func TestDevPreviewTool_MissingAction(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{}) result := tool.Execute(context.Background(), map[string]any{})
@ -282,15 +339,19 @@ func TestDevPreviewTool_MissingAction(t *testing.T) {
func TestDevPreviewTool_NameAndSchema(t *testing.T) { func TestDevPreviewTool_NameAndSchema(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
if tool.Name() != "dev_preview" { if tool.Name() != "dev_preview" {
t.Errorf("expected name dev_preview, got %q", tool.Name()) t.Errorf("expected name dev_preview, got %q", tool.Name())
} }
if tool.Description() == "" { if tool.Description() == "" {
t.Error("expected non-empty description") t.Error("expected non-empty description")
} }
params := tool.Parameters() params := tool.Parameters()
if params == nil { if params == nil {
t.Fatal("expected non-nil parameters") t.Fatal("expected non-nil parameters")
} }
@ -300,26 +361,35 @@ func TestDevPreviewTool_NameAndSchema(t *testing.T) {
func TestDevPreviewTool_StartMultipleTargets(t *testing.T) { func TestDevPreviewTool_StartMultipleTargets(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
r1 := tool.Execute(context.Background(), map[string]any{ r1 := tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://localhost:8080", "target": "http://localhost:8080",
"name": "api", "name": "api",
}) })
r2 := tool.Execute(context.Background(), map[string]any{ r2 := tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://localhost:3000", "target": "http://localhost:3000",
"name": "frontend", "name": "frontend",
}) })
if r1.IsError || r2.IsError { if r1.IsError || r2.IsError {
t.Fatalf("expected both starts to succeed, got err1=%v err2=%v", 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 { if len(mgr.targets) != 2 {
t.Errorf("expected 2 registered targets, got %d", len(mgr.targets)) t.Errorf("expected 2 registered targets, got %d", len(mgr.targets))
} }
// The second start should make the frontend active // The second start should make the frontend active
if mgr.active != "http://localhost:3000" { if mgr.active != "http://localhost:3000" {
t.Errorf("expected last started target to be active, got %q", mgr.active) t.Errorf("expected last started target to be active, got %q", mgr.active)
} }
@ -327,11 +397,14 @@ func TestDevPreviewTool_StartMultipleTargets(t *testing.T) {
func TestDevPreviewTool_StopPreservesRegistrations(t *testing.T) { func TestDevPreviewTool_StopPreservesRegistrations(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
tool.Execute(context.Background(), map[string]any{ tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://localhost:3000", "target": "http://localhost:3000",
"name": "frontend", "name": "frontend",
}) })
@ -342,11 +415,15 @@ func TestDevPreviewTool_StopPreservesRegistrations(t *testing.T) {
if result.IsError { if result.IsError {
t.Fatalf("stop failed: %s", result.ForLLM) t.Fatalf("stop failed: %s", result.ForLLM)
} }
// Registration should still be there // Registration should still be there
if len(mgr.targets) != 1 { if len(mgr.targets) != 1 {
t.Errorf("expected 1 registered target after stop, got %d", len(mgr.targets)) t.Errorf("expected 1 registered target after stop, got %d", len(mgr.targets))
} }
// But active should be cleared // But active should be cleared
if mgr.active != "" { if mgr.active != "" {
t.Errorf("expected inactive after stop, got %q", 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) { func TestDevPreviewTool_StatusWithTargetsButInactive(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
mgr.RegisterDevTarget("api", "http://localhost:8080") mgr.RegisterDevTarget("api", "http://localhost:8080")
// active remains empty // active remains empty
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -366,9 +445,11 @@ func TestDevPreviewTool_StatusWithTargetsButInactive(t *testing.T) {
if result.IsError { if result.IsError {
t.Fatalf("status failed: %s", result.ForLLM) t.Fatalf("status failed: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "not active") { if !strings.Contains(result.ForLLM, "not active") {
t.Errorf("expected 'not active' in status, got %q", result.ForLLM) t.Errorf("expected 'not active' in status, got %q", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "api") { if !strings.Contains(result.ForLLM, "api") {
t.Errorf("expected 'api' listed in status, got %q", result.ForLLM) 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) { func TestDevPreviewTool_ResultIsSilent(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
cases := []struct { cases := []struct {
name string name string
args map[string]any args map[string]any
}{ }{
{"start", map[string]any{"action": "start", "target": "http://localhost:3000"}}, {"start", map[string]any{"action": "start", "target": "http://localhost:3000"}},
{"stop", map[string]any{"action": "stop"}}, {"stop", map[string]any{"action": "stop"}},
{"status", map[string]any{"action": "status"}}, {"status", map[string]any{"action": "status"}},
} }
for _, tc := range cases { for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) { t.Run(tc.name, func(t *testing.T) {
result := tool.Execute(context.Background(), tc.args) result := tool.Execute(context.Background(), tc.args)
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
if result.Silent != true { if result.Silent != true {
t.Errorf("expected SilentResult (IsSilent=true), got IsSilent=%v", result.Silent) 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) { func TestDevPreviewTool_ActionTypeNotString(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": 123, "action": 123,
}) })
if !result.IsError { if !result.IsError {
t.Error("expected error for non-string action") t.Error("expected error for non-string action")
} }
@ -414,17 +504,26 @@ func TestDevPreviewTool_ActionTypeNotString(t *testing.T) {
func TestDevPreviewTool_InferName(t *testing.T) { func TestDevPreviewTool_InferName(t *testing.T) {
cases := []struct { cases := []struct {
target string target string
want string want string
}{ }{
{"http://localhost:3000", "localhost:3000"}, {"http://localhost:3000", "localhost:3000"},
{"http://localhost:8080", "localhost:8080"}, {"http://localhost:8080", "localhost:8080"},
{"http://127.0.0.1:9000", "127.0.0.1:9000"}, {"http://127.0.0.1:9000", "127.0.0.1:9000"},
{"http://localhost", "localhost"}, {"http://localhost", "localhost"},
{"http://[::1]:5000", "::1:5000"}, {"http://[::1]:5000", "::1:5000"},
{"not-a-url", ""}, // url.Parse succeeds but Hostname() is empty {"not-a-url", ""}, // url.Parse succeeds but Hostname() is empty
} }
for _, tc := range cases { for _, tc := range cases {
got := inferName(tc.target) got := inferName(tc.target)
if got != tc.want { if got != tc.want {
t.Errorf("inferName(%q) = %q, want %q", tc.target, 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) { func TestDevPreviewTool_StartEmptyName(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
// Explicitly pass empty name — should auto-infer // Explicitly pass empty name — should auto-infer
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://localhost:5000", "target": "http://localhost:5000",
"name": "", "name": "",
}) })
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
for _, dt := range mgr.targets { for _, dt := range mgr.targets {
if dt.Name != "localhost:5000" { if dt.Name != "localhost:5000" {
t.Errorf("expected auto-name 'localhost:5000', got %q", dt.Name) 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) { func TestDevPreviewTool_UnregisterActiveTarget(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
// Register and activate // Register and activate
tool.Execute(context.Background(), map[string]any{ tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "http://localhost:3000", "target": "http://localhost:3000",
"name": "frontend", "name": "frontend",
}) })
// Find the registered ID // Find the registered ID
var id string var id string
for k := range mgr.targets { for k := range mgr.targets {
id = k id = k
} }
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "unregister", "action": "unregister",
"id": id, "id": id,
}) })
if result.IsError { if result.IsError {
t.Fatalf("unregister failed: %s", result.ForLLM) t.Fatalf("unregister failed: %s", result.ForLLM)
} }
if len(mgr.targets) != 0 { if len(mgr.targets) != 0 {
t.Errorf("expected 0 targets, got %d", len(mgr.targets)) t.Errorf("expected 0 targets, got %d", len(mgr.targets))
} }
if mgr.active != "" { if mgr.active != "" {
t.Errorf("expected no active target, got %q", 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) { func TestDevPreviewTool_StartTargetEmptyString(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"action": "start", "action": "start",
"target": "", "target": "",
}) })
@ -501,8 +616,11 @@ func TestDevPreviewTool_StartTargetEmptyString(t *testing.T) {
func TestDevPreviewTool_StatusActiveNoTargets(t *testing.T) { func TestDevPreviewTool_StatusActiveNoTargets(t *testing.T) {
// Edge case: active proxy but no registered targets (shouldn't normally happen) // Edge case: active proxy but no registered targets (shouldn't normally happen)
mgr := newMockManager() mgr := newMockManager()
mgr.active = "http://localhost:9999" // active but targets map is empty mgr.active = "http://localhost:9999" // active but targets map is empty
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -512,12 +630,15 @@ func TestDevPreviewTool_StatusActiveNoTargets(t *testing.T) {
if result.IsError { if result.IsError {
t.Fatalf("expected success, got error: %s", result.ForLLM) t.Fatalf("expected success, got error: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "active") { if !strings.Contains(result.ForLLM, "active") {
t.Errorf("expected 'active' in result, got %q", result.ForLLM) t.Errorf("expected 'active' in result, got %q", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "http://localhost:9999") { if !strings.Contains(result.ForLLM, "http://localhost:9999") {
t.Errorf("expected target URL in result, got %q", result.ForLLM) t.Errorf("expected target URL in result, got %q", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "No registered targets") { if !strings.Contains(result.ForLLM, "No registered targets") {
t.Errorf("expected 'No registered targets' in result, got %q", result.ForLLM) 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) { func TestDevPreviewTool_StatusOutputFormat(t *testing.T) {
mgr := newMockManager() mgr := newMockManager()
tool := NewDevPreviewTool(mgr) tool := NewDevPreviewTool(mgr)
id1, _ := mgr.RegisterDevTarget("api", "http://localhost:8080") id1, _ := mgr.RegisterDevTarget("api", "http://localhost:8080")
mgr.RegisterDevTarget("frontend", "http://localhost:3000") mgr.RegisterDevTarget("frontend", "http://localhost:3000")
mgr.ActivateDevTarget(id1) mgr.ActivateDevTarget(id1)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -538,15 +662,21 @@ func TestDevPreviewTool_StatusOutputFormat(t *testing.T) {
if result.IsError { if result.IsError {
t.Fatalf("status failed: %s", result.ForLLM) t.Fatalf("status failed: %s", result.ForLLM)
} }
// Should contain IDs in bracket format // Should contain IDs in bracket format
if !strings.Contains(result.ForLLM, "["+id1+"]") { if !strings.Contains(result.ForLLM, "["+id1+"]") {
t.Errorf("expected [%s] in output, got %q", id1, result.ForLLM) t.Errorf("expected [%s] in output, got %q", id1, result.ForLLM)
} }
// Should contain the arrow // Should contain the arrow
if !strings.Contains(result.ForLLM, "→") { if !strings.Contains(result.ForLLM, "→") {
t.Errorf("expected arrow in output, got %q", result.ForLLM) t.Errorf("expected arrow in output, got %q", result.ForLLM)
} }
// Should contain "Registered targets:" // Should contain "Registered targets:"
if !strings.Contains(result.ForLLM, "Registered targets:") { if !strings.Contains(result.ForLLM, "Registered targets:") {
t.Errorf("expected 'Registered targets:' header, got %q", result.ForLLM) 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. // EditFileTool edits a file by replacing old_text with new_text.
// The old_text must exist exactly in the file. // The old_text must exist exactly in the file.
type EditFileTool struct { type EditFileTool struct {
fs fileSystem fs fileSystem
} }
// NewEditFileTool creates a new EditFileTool with optional directory restriction. // NewEditFileTool creates a new EditFileTool with optional directory restriction.
func NewEditFileTool(workspace string, restrict bool) *EditFileTool { func NewEditFileTool(workspace string, restrict bool) *EditFileTool {
var fs fileSystem var fs fileSystem
if restrict { if restrict {
fs = &sandboxFs{workspace: workspace} fs = &sandboxFs{workspace: workspace}
} else { } else {
fs = &hostFs{} fs = &hostFs{}
} }
return &EditFileTool{fs: fs} return &EditFileTool{fs: fs}
} }
@ -36,36 +41,46 @@ func (t *EditFileTool) Description() string {
func (t *EditFileTool) Parameters() map[string]any { func (t *EditFileTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"path": map[string]any{ "path": map[string]any{
"type": "string", "type": "string",
"description": "The file path to edit", "description": "The file path to edit",
}, },
"old_text": map[string]any{ "old_text": map[string]any{
"type": "string", "type": "string",
"description": "The exact text to find and replace", "description": "The exact text to find and replace",
}, },
"new_text": map[string]any{ "new_text": map[string]any{
"type": "string", "type": "string",
"description": "The text to replace with", "description": "The text to replace with",
}, },
}, },
"required": []string{"path", "old_text", "new_text"}, "required": []string{"path", "old_text", "new_text"},
} }
} }
func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
path, ok := args["path"].(string) path, ok := args["path"].(string)
if !ok { if !ok {
return ErrorResult("path is required") return ErrorResult("path is required")
} }
oldText, ok := args["old_text"].(string) oldText, ok := args["old_text"].(string)
if !ok { if !ok {
return ErrorResult("old_text is required") return ErrorResult("old_text is required")
} }
newText, ok := args["new_text"].(string) newText, ok := args["new_text"].(string)
if !ok { if !ok {
return ErrorResult("new_text is required") 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 { if err := editFile(resolveFS(ctx, t.fs, path), path, oldText, newText); err != nil {
return ErrorResult(err.Error()) return ErrorResult(err.Error())
} }
return SilentResult(fmt.Sprintf("File edited: %s", path)) return SilentResult(fmt.Sprintf("File edited: %s", path))
} }
@ -82,11 +98,13 @@ type AppendFileTool struct {
func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool { func NewAppendFileTool(workspace string, restrict bool) *AppendFileTool {
var fs fileSystem var fs fileSystem
if restrict { if restrict {
fs = &sandboxFs{workspace: workspace} fs = &sandboxFs{workspace: workspace}
} else { } else {
fs = &hostFs{} fs = &hostFs{}
} }
return &AppendFileTool{fs: fs} return &AppendFileTool{fs: fs}
} }
@ -101,27 +119,34 @@ func (t *AppendFileTool) Description() string {
func (t *AppendFileTool) Parameters() map[string]any { func (t *AppendFileTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"path": map[string]any{ "path": map[string]any{
"type": "string", "type": "string",
"description": "The file path to append to", "description": "The file path to append to",
}, },
"content": map[string]any{ "content": map[string]any{
"type": "string", "type": "string",
"description": "The content to append", "description": "The content to append",
}, },
}, },
"required": []string{"path", "content"}, "required": []string{"path", "content"},
} }
} }
func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
path, ok := args["path"].(string) path, ok := args["path"].(string)
if !ok { if !ok {
return ErrorResult("path is required") return ErrorResult("path is required")
} }
content, ok := args["content"].(string) content, ok := args["content"].(string)
if !ok { if !ok {
return ErrorResult("content is required") 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 { if err := appendFile(resolveFS(ctx, t.fs, path), path, content); err != nil {
return ErrorResult(err.Error()) return ErrorResult(err.Error())
} }
return SilentResult(fmt.Sprintf("Appended to %s", path)) return SilentResult(fmt.Sprintf("Appended to %s", path))
} }
// editFile reads the file via sysFs, performs the replacement, and writes back. // 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. // It uses a fileSystem interface, allowing the same logic for both restricted and unrestricted modes.
func editFile(sysFs fileSystem, path, oldText, newText string) error { func editFile(sysFs fileSystem, path, oldText, newText string) error {
content, err := sysFs.ReadFile(path) content, err := sysFs.ReadFile(path)
if err != nil { 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. // appendFile reads the existing content (if any) via sysFs, appends new content, and writes back.
func appendFile(sysFs fileSystem, path, appendContent string) error { func appendFile(sysFs fileSystem, path, appendContent string) error {
content, err := sysFs.ReadFile(path) content, err := sysFs.ReadFile(path)
if err != nil && !errors.Is(err, fs.ErrNotExist) { if err != nil && !errors.Is(err, fs.ErrNotExist) {
return err return err
} }
newContent := append(content, []byte(appendContent)...) newContent := append(content, []byte(appendContent)...)
return sysFs.WriteFile(path, newContent) return sysFs.WriteFile(path, newContent)
} }
// replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText. // replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText.
func replaceEditContent(content []byte, oldText, newText string) ([]byte, error) { func replaceEditContent(content []byte, oldText, newText string) ([]byte, error) {
contentStr := string(content) contentStr := string(content)
@ -168,10 +200,12 @@ func replaceEditContent(content []byte, oldText, newText string) ([]byte, error)
} }
count := strings.Count(contentStr, oldText) count := strings.Count(contentStr, oldText)
if count > 1 { if count > 1 {
return nil, fmt.Errorf("old_text appears %d times. Please provide more context to make it unique", count) 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) newContent := strings.Replace(contentStr, oldText, newText, 1)
return []byte(newContent), nil return []byte(newContent), nil
} }

View file

@ -11,261 +11,349 @@ import (
) )
// TestEditTool_EditFile_Success verifies successful file editing // TestEditTool_EditFile_Success verifies successful file editing
func TestEditTool_EditFile_Success(t *testing.T) { func TestEditTool_EditFile_Success(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.txt") testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("Hello World\nThis is a test"), 0o644) os.WriteFile(testFile, []byte("Hello World\nThis is a test"), 0o644)
tool := NewEditFileTool(tmpDir, true) tool := NewEditFileTool(tmpDir, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"old_text": "World", "old_text": "World",
"new_text": "Universe", "new_text": "Universe",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// Should return SilentResult // Should return SilentResult
if !result.Silent { if !result.Silent {
t.Errorf("Expected Silent=true for EditFile, got false") t.Errorf("Expected Silent=true for EditFile, got false")
} }
// ForUser should be empty (silent result) // ForUser should be empty (silent result)
if result.ForUser != "" { if result.ForUser != "" {
t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser)
} }
// Verify file was actually edited // Verify file was actually edited
content, err := os.ReadFile(testFile) content, err := os.ReadFile(testFile)
if err != nil { if err != nil {
t.Fatalf("Failed to read edited file: %v", err) t.Fatalf("Failed to read edited file: %v", err)
} }
contentStr := string(content) contentStr := string(content)
if !strings.Contains(contentStr, "Hello Universe") { if !strings.Contains(contentStr, "Hello Universe") {
t.Errorf("Expected file to contain 'Hello Universe', got: %s", contentStr) t.Errorf("Expected file to contain 'Hello Universe', got: %s", contentStr)
} }
if strings.Contains(contentStr, "Hello World") { if strings.Contains(contentStr, "Hello World") {
t.Errorf("Expected 'Hello World' to be replaced, got: %s", contentStr) t.Errorf("Expected 'Hello World' to be replaced, got: %s", contentStr)
} }
} }
// TestEditTool_EditFile_NotFound verifies error handling for non-existent file // TestEditTool_EditFile_NotFound verifies error handling for non-existent file
func TestEditTool_EditFile_NotFound(t *testing.T) { func TestEditTool_EditFile_NotFound(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "nonexistent.txt") testFile := filepath.Join(tmpDir, "nonexistent.txt")
tool := NewEditFileTool(tmpDir, true) tool := NewEditFileTool(tmpDir, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"old_text": "old", "old_text": "old",
"new_text": "new", "new_text": "new",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error for non-existent file") t.Errorf("Expected error for non-existent file")
} }
// Should mention file not found // Should mention file not found
if !strings.Contains(result.ForLLM, "not found") && !strings.Contains(result.ForUser, "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) t.Errorf("Expected 'file not found' message, got ForLLM: %s", result.ForLLM)
} }
} }
// TestEditTool_EditFile_OldTextNotFound verifies error when old_text doesn't exist // TestEditTool_EditFile_OldTextNotFound verifies error when old_text doesn't exist
func TestEditTool_EditFile_OldTextNotFound(t *testing.T) { func TestEditTool_EditFile_OldTextNotFound(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.txt") testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("Hello World"), 0o644) os.WriteFile(testFile, []byte("Hello World"), 0o644)
tool := NewEditFileTool(tmpDir, true) tool := NewEditFileTool(tmpDir, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"old_text": "Goodbye", "old_text": "Goodbye",
"new_text": "Hello", "new_text": "Hello",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when old_text not found") t.Errorf("Expected error when old_text not found")
} }
// Should mention old_text not found // Should mention old_text not found
if !strings.Contains(result.ForLLM, "not found") && !strings.Contains(result.ForUser, "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) t.Errorf("Expected 'not found' message, got ForLLM: %s", result.ForLLM)
} }
} }
// TestEditTool_EditFile_MultipleMatches verifies error when old_text appears multiple times // TestEditTool_EditFile_MultipleMatches verifies error when old_text appears multiple times
func TestEditTool_EditFile_MultipleMatches(t *testing.T) { func TestEditTool_EditFile_MultipleMatches(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.txt") testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("test test test"), 0o644) os.WriteFile(testFile, []byte("test test test"), 0o644)
tool := NewEditFileTool(tmpDir, true) tool := NewEditFileTool(tmpDir, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"old_text": "test", "old_text": "test",
"new_text": "done", "new_text": "done",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when old_text appears multiple times") t.Errorf("Expected error when old_text appears multiple times")
} }
// Should mention multiple occurrences // Should mention multiple occurrences
if !strings.Contains(result.ForLLM, "times") && !strings.Contains(result.ForUser, "times") { if !strings.Contains(result.ForLLM, "times") && !strings.Contains(result.ForUser, "times") {
t.Errorf("Expected 'multiple times' message, got ForLLM: %s", result.ForLLM) t.Errorf("Expected 'multiple times' message, got ForLLM: %s", result.ForLLM)
} }
} }
// TestEditTool_EditFile_OutsideAllowedDir verifies error when path is outside allowed directory // TestEditTool_EditFile_OutsideAllowedDir verifies error when path is outside allowed directory
func TestEditTool_EditFile_OutsideAllowedDir(t *testing.T) { func TestEditTool_EditFile_OutsideAllowedDir(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
otherDir := t.TempDir() otherDir := t.TempDir()
testFile := filepath.Join(otherDir, "test.txt") testFile := filepath.Join(otherDir, "test.txt")
os.WriteFile(testFile, []byte("content"), 0o644) os.WriteFile(testFile, []byte("content"), 0o644)
tool := NewEditFileTool(tmpDir, true) // Restrict to tmpDir tool := NewEditFileTool(tmpDir, true) // Restrict to tmpDir
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"old_text": "content", "old_text": "content",
"new_text": "new", "new_text": "new",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
assert.True(t, result.IsError, "Expected error when path is outside allowed directory") assert.True(t, result.IsError, "Expected error when path is outside allowed directory")
// Should mention outside allowed directory // Should mention outside allowed directory
// Note: ErrorResult only sets ForLLM by default, so ForUser might be empty. // Note: ErrorResult only sets ForLLM by default, so ForUser might be empty.
// We check ForLLM as it's the primary error channel. // We check ForLLM as it's the primary error channel.
assert.True( assert.True(
t, t,
strings.Contains(result.ForLLM, "outside") || strings.Contains(result.ForLLM, "access denied") || strings.Contains(result.ForLLM, "outside") || strings.Contains(result.ForLLM, "access denied") ||
strings.Contains(result.ForLLM, "escapes"), strings.Contains(result.ForLLM, "escapes"),
"Expected 'outside allowed' or 'access denied' message, got ForLLM: %s", "Expected 'outside allowed' or 'access denied' message, got ForLLM: %s",
result.ForLLM, result.ForLLM,
) )
} }
// TestEditTool_EditFile_MissingPath verifies error handling for missing path // TestEditTool_EditFile_MissingPath verifies error handling for missing path
func TestEditTool_EditFile_MissingPath(t *testing.T) { func TestEditTool_EditFile_MissingPath(t *testing.T) {
tool := NewEditFileTool("", false) tool := NewEditFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"old_text": "old", "old_text": "old",
"new_text": "new", "new_text": "new",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when path is missing") t.Errorf("Expected error when path is missing")
} }
} }
// TestEditTool_EditFile_MissingOldText verifies error handling for missing old_text // TestEditTool_EditFile_MissingOldText verifies error handling for missing old_text
func TestEditTool_EditFile_MissingOldText(t *testing.T) { func TestEditTool_EditFile_MissingOldText(t *testing.T) {
tool := NewEditFileTool("", false) tool := NewEditFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "/tmp/test.txt", "path": "/tmp/test.txt",
"new_text": "new", "new_text": "new",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when old_text is missing") t.Errorf("Expected error when old_text is missing")
} }
} }
// TestEditTool_EditFile_MissingNewText verifies error handling for missing new_text // TestEditTool_EditFile_MissingNewText verifies error handling for missing new_text
func TestEditTool_EditFile_MissingNewText(t *testing.T) { func TestEditTool_EditFile_MissingNewText(t *testing.T) {
tool := NewEditFileTool("", false) tool := NewEditFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "/tmp/test.txt", "path": "/tmp/test.txt",
"old_text": "old", "old_text": "old",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when new_text is missing") t.Errorf("Expected error when new_text is missing")
} }
} }
// TestEditTool_AppendFile_Success verifies successful file appending // TestEditTool_AppendFile_Success verifies successful file appending
func TestEditTool_AppendFile_Success(t *testing.T) { func TestEditTool_AppendFile_Success(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.txt") testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("Initial content"), 0o644) os.WriteFile(testFile, []byte("Initial content"), 0o644)
tool := NewAppendFileTool("", false) tool := NewAppendFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"content": "\nAppended content", "content": "\nAppended content",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// Should return SilentResult // Should return SilentResult
if !result.Silent { if !result.Silent {
t.Errorf("Expected Silent=true for AppendFile, got false") t.Errorf("Expected Silent=true for AppendFile, got false")
} }
// ForUser should be empty (silent result) // ForUser should be empty (silent result)
if result.ForUser != "" { if result.ForUser != "" {
t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser)
} }
// Verify content was actually appended // Verify content was actually appended
content, err := os.ReadFile(testFile) content, err := os.ReadFile(testFile)
if err != nil { if err != nil {
t.Fatalf("Failed to read file: %v", err) t.Fatalf("Failed to read file: %v", err)
} }
contentStr := string(content) contentStr := string(content)
if !strings.Contains(contentStr, "Initial content") { if !strings.Contains(contentStr, "Initial content") {
t.Errorf("Expected original content to remain, got: %s", contentStr) t.Errorf("Expected original content to remain, got: %s", contentStr)
} }
if !strings.Contains(contentStr, "Appended content") { if !strings.Contains(contentStr, "Appended content") {
t.Errorf("Expected appended content, got: %s", contentStr) t.Errorf("Expected appended content, got: %s", contentStr)
} }
} }
// TestEditTool_AppendFile_MissingPath verifies error handling for missing path // TestEditTool_AppendFile_MissingPath verifies error handling for missing path
func TestEditTool_AppendFile_MissingPath(t *testing.T) { func TestEditTool_AppendFile_MissingPath(t *testing.T) {
tool := NewAppendFileTool("", false) tool := NewAppendFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"content": "test", "content": "test",
} }
@ -273,15 +361,19 @@ func TestEditTool_AppendFile_MissingPath(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when path is missing") t.Errorf("Expected error when path is missing")
} }
} }
// TestEditTool_AppendFile_MissingContent verifies error handling for missing content // TestEditTool_AppendFile_MissingContent verifies error handling for missing content
func TestEditTool_AppendFile_MissingContent(t *testing.T) { func TestEditTool_AppendFile_MissingContent(t *testing.T) {
tool := NewAppendFileTool("", false) tool := NewAppendFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "/tmp/test.txt", "path": "/tmp/test.txt",
} }
@ -289,43 +381,67 @@ func TestEditTool_AppendFile_MissingContent(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when content is missing") t.Errorf("Expected error when content is missing")
} }
} }
// TestReplaceEditContent verifies the helper function replaceEditContent // TestReplaceEditContent verifies the helper function replaceEditContent
func TestReplaceEditContent(t *testing.T) { func TestReplaceEditContent(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
content []byte content []byte
oldText string oldText string
newText string newText string
expected []byte expected []byte
expectError bool expectError bool
}{ }{
{ {
name: "successful replacement", name: "successful replacement",
content: []byte("hello world"), content: []byte("hello world"),
oldText: "world", oldText: "world",
newText: "universe", newText: "universe",
expected: []byte("hello universe"), expected: []byte("hello universe"),
expectError: false, expectError: false,
}, },
{ {
name: "old text not found", name: "old text not found",
content: []byte("hello world"), content: []byte("hello world"),
oldText: "golang", oldText: "golang",
newText: "rust", newText: "rust",
expected: nil, expected: nil,
expectError: true, expectError: true,
}, },
{ {
name: "multiple matches found", name: "multiple matches found",
content: []byte("test text test"), content: []byte("test text test"),
oldText: "test", oldText: "test",
newText: "done", newText: "done",
expected: nil, expected: nil,
expectError: true, expectError: true,
}, },
} }
@ -333,10 +449,12 @@ func TestReplaceEditContent(t *testing.T) {
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
result, err := replaceEditContent(tt.content, tt.oldText, tt.newText) result, err := replaceEditContent(tt.content, tt.oldText, tt.newText)
if tt.expectError { if tt.expectError {
assert.Error(t, err) assert.Error(t, err)
} else { } else {
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, tt.expected, result) assert.Equal(t, tt.expected, result)
} }
}) })
@ -344,94 +462,142 @@ func TestReplaceEditContent(t *testing.T) {
} }
// TestAppendFileTool_AppendToNonExistent_Restricted verifies that AppendFileTool in restricted mode // 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. // 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. // This exercises the errors.Is(err, fs.ErrNotExist) path in appendFile + sandboxFs.
func TestAppendFileTool_AppendToNonExistent_Restricted(t *testing.T) { func TestAppendFileTool_AppendToNonExistent_Restricted(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
tool := NewAppendFileTool(workspace, true) tool := NewAppendFileTool(workspace, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "brand_new_file.txt", "path": "brand_new_file.txt",
"content": "first content", "content": "first content",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
assert.False( assert.False(
t, t,
result.IsError, result.IsError,
"Expected success when appending to non-existent file in restricted mode, got: %s", "Expected success when appending to non-existent file in restricted mode, got: %s",
result.ForLLM, result.ForLLM,
) )
// Verify the file was created with correct content // Verify the file was created with correct content
data, err := os.ReadFile(filepath.Join(workspace, "brand_new_file.txt")) data, err := os.ReadFile(filepath.Join(workspace, "brand_new_file.txt"))
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, "first content", string(data)) assert.Equal(t, "first content", string(data))
} }
// TestAppendFileTool_Restricted_Success verifies that AppendFileTool in restricted mode // TestAppendFileTool_Restricted_Success verifies that AppendFileTool in restricted mode
// correctly appends to an existing file within the sandbox. // correctly appends to an existing file within the sandbox.
func TestAppendFileTool_Restricted_Success(t *testing.T) { func TestAppendFileTool_Restricted_Success(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
testFile := "existing.txt" testFile := "existing.txt"
err := os.WriteFile(filepath.Join(workspace, testFile), []byte("initial"), 0o644) err := os.WriteFile(filepath.Join(workspace, testFile), []byte("initial"), 0o644)
assert.NoError(t, err) assert.NoError(t, err)
tool := NewAppendFileTool(workspace, true) tool := NewAppendFileTool(workspace, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"content": " appended", "content": " appended",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM)
assert.True(t, result.Silent) assert.True(t, result.Silent)
data, err := os.ReadFile(filepath.Join(workspace, testFile)) data, err := os.ReadFile(filepath.Join(workspace, testFile))
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, "initial appended", string(data)) assert.Equal(t, "initial appended", string(data))
} }
// TestEditFileTool_Restricted_InPlaceEdit verifies that EditFileTool in restricted mode // TestEditFileTool_Restricted_InPlaceEdit verifies that EditFileTool in restricted mode
// correctly edits a file using the sandboxFs path. // correctly edits a file using the sandboxFs path.
func TestEditFileTool_Restricted_InPlaceEdit(t *testing.T) { func TestEditFileTool_Restricted_InPlaceEdit(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
testFile := "edit_target.txt" testFile := "edit_target.txt"
err := os.WriteFile(filepath.Join(workspace, testFile), []byte("Hello World"), 0o644) err := os.WriteFile(filepath.Join(workspace, testFile), []byte("Hello World"), 0o644)
assert.NoError(t, err) assert.NoError(t, err)
tool := NewEditFileTool(workspace, true) tool := NewEditFileTool(workspace, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"old_text": "World", "old_text": "World",
"new_text": "Go", "new_text": "Go",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM)
assert.True(t, result.Silent) assert.True(t, result.Silent)
data, err := os.ReadFile(filepath.Join(workspace, testFile)) data, err := os.ReadFile(filepath.Join(workspace, testFile))
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, "Hello Go", string(data)) assert.Equal(t, "Hello Go", string(data))
} }
// TestEditFileTool_Restricted_FileNotFound verifies that editFile returns a proper // TestEditFileTool_Restricted_FileNotFound verifies that editFile returns a proper
// error message when the target file does not exist. // error message when the target file does not exist.
func TestEditFileTool_Restricted_FileNotFound(t *testing.T) { func TestEditFileTool_Restricted_FileNotFound(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
tool := NewEditFileTool(workspace, true) tool := NewEditFileTool(workspace, true)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "no_such_file.txt", "path": "no_such_file.txt",
"old_text": "old", "old_text": "old",
"new_text": "new", "new_text": "new",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
assert.True(t, result.IsError) assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "not found") 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. // validatePath ensures the given path is within the workspace if restrict is true.
// Used by shell.go for working directory validation. // Used by shell.go for working directory validation.
func validatePath(path, workspace string, restrict bool) (string, error) { func validatePath(path, workspace string, restrict bool) (string, error) {
if workspace == "" { if workspace == "" {
return path, fmt.Errorf("workspace is not defined") return path, fmt.Errorf("workspace is not defined")
@ -25,6 +27,7 @@ func validatePath(path, workspace string, restrict bool) (string, error) {
} }
var absPath string var absPath string
if filepath.IsAbs(path) { if filepath.IsAbs(path) {
absPath = filepath.Clean(path) absPath = filepath.Clean(path)
} else { } else {
@ -40,7 +43,9 @@ func validatePath(path, workspace string, restrict bool) (string, error) {
} }
var resolved string var resolved string
workspaceReal := absWorkspace workspaceReal := absWorkspace
if resolved, err = filepath.EvalSymlinks(absWorkspace); err == nil { if resolved, err = filepath.EvalSymlinks(absWorkspace); err == nil {
workspaceReal = resolved workspaceReal = resolved
} }
@ -51,6 +56,7 @@ func validatePath(path, workspace string, restrict bool) (string, error) {
} }
} else if os.IsNotExist(err) { } else if os.IsNotExist(err) {
var parentResolved string var parentResolved string
if parentResolved, err = resolveExistingAncestor(filepath.Dir(absPath)); err == nil { if parentResolved, err = resolveExistingAncestor(filepath.Dir(absPath)); err == nil {
if !isWithinWorkspace(parentResolved, workspaceReal) { if !isWithinWorkspace(parentResolved, workspaceReal) {
return "", fmt.Errorf("access denied: symlink resolves outside workspace") return "", fmt.Errorf("access denied: symlink resolves outside workspace")
@ -73,6 +79,7 @@ func resolveExistingAncestor(path string) (string, error) {
} else if !os.IsNotExist(err) { } else if !os.IsNotExist(err) {
return "", err return "", err
} }
if filepath.Dir(current) == current { if filepath.Dir(current) == current {
return "", os.ErrNotExist return "", os.ErrNotExist
} }
@ -81,6 +88,7 @@ func resolveExistingAncestor(path string) (string, error) {
func isWithinWorkspace(candidate, workspace string) bool { func isWithinWorkspace(candidate, workspace string) bool {
rel, err := filepath.Rel(filepath.Clean(workspace), filepath.Clean(candidate)) rel, err := filepath.Rel(filepath.Clean(workspace), filepath.Clean(candidate))
return err == nil && filepath.IsLocal(rel) return err == nil && filepath.IsLocal(rel)
} }
@ -90,11 +98,13 @@ type ReadFileTool struct {
func NewReadFileTool(workspace string, restrict bool) *ReadFileTool { func NewReadFileTool(workspace string, restrict bool) *ReadFileTool {
var fs fileSystem var fs fileSystem
if restrict { if restrict {
fs = &sandboxFs{workspace: workspace} fs = &sandboxFs{workspace: workspace}
} else { } else {
fs = &hostFs{} fs = &hostFs{}
} }
return &ReadFileTool{fs: fs} return &ReadFileTool{fs: fs}
} }
@ -109,18 +119,22 @@ func (t *ReadFileTool) Description() string {
func (t *ReadFileTool) Parameters() map[string]any { func (t *ReadFileTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"path": map[string]any{ "path": map[string]any{
"type": "string", "type": "string",
"description": "Path to the file to read", "description": "Path to the file to read",
}, },
}, },
"required": []string{"path"}, "required": []string{"path"},
} }
} }
func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
path, ok := args["path"].(string) path, ok := args["path"].(string)
if !ok { if !ok {
return ErrorResult("path is required") return ErrorResult("path is required")
} }
@ -129,6 +143,7 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe
if err != nil { if err != nil {
return ErrorResult(err.Error()) return ErrorResult(err.Error())
} }
return NewToolResult(string(content)) return NewToolResult(string(content))
} }
@ -138,11 +153,13 @@ type WriteFileTool struct {
func NewWriteFileTool(workspace string, restrict bool) *WriteFileTool { func NewWriteFileTool(workspace string, restrict bool) *WriteFileTool {
var fs fileSystem var fs fileSystem
if restrict { if restrict {
fs = &sandboxFs{workspace: workspace} fs = &sandboxFs{workspace: workspace}
} else { } else {
fs = &hostFs{} fs = &hostFs{}
} }
return &WriteFileTool{fs: fs} return &WriteFileTool{fs: fs}
} }
@ -157,27 +174,34 @@ func (t *WriteFileTool) Description() string {
func (t *WriteFileTool) Parameters() map[string]any { func (t *WriteFileTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"path": map[string]any{ "path": map[string]any{
"type": "string", "type": "string",
"description": "Path to the file to write", "description": "Path to the file to write",
}, },
"content": map[string]any{ "content": map[string]any{
"type": "string", "type": "string",
"description": "Content to write to the file", "description": "Content to write to the file",
}, },
}, },
"required": []string{"path", "content"}, "required": []string{"path", "content"},
} }
} }
func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
path, ok := args["path"].(string) path, ok := args["path"].(string)
if !ok { if !ok {
return ErrorResult("path is required") return ErrorResult("path is required")
} }
content, ok := args["content"].(string) content, ok := args["content"].(string)
if !ok { if !ok {
return ErrorResult("content is required") return ErrorResult("content is required")
} }
@ -195,11 +219,13 @@ type ListDirTool struct {
func NewListDirTool(workspace string, restrict bool) *ListDirTool { func NewListDirTool(workspace string, restrict bool) *ListDirTool {
var fs fileSystem var fs fileSystem
if restrict { if restrict {
fs = &sandboxFs{workspace: workspace} fs = &sandboxFs{workspace: workspace}
} else { } else {
fs = &hostFs{} fs = &hostFs{}
} }
return &ListDirTool{fs: fs} return &ListDirTool{fs: fs}
} }
@ -214,18 +240,22 @@ func (t *ListDirTool) Description() string {
func (t *ListDirTool) Parameters() map[string]any { func (t *ListDirTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"path": map[string]any{ "path": map[string]any{
"type": "string", "type": "string",
"description": "Path to list", "description": "Path to list",
}, },
}, },
"required": []string{"path"}, "required": []string{"path"},
} }
} }
func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
path, ok := args["path"].(string) path, ok := args["path"].(string)
if !ok { if !ok {
path = "." path = "."
} }
@ -234,32 +264,42 @@ func (t *ListDirTool) Execute(ctx context.Context, args map[string]any) *ToolRes
if err != nil { if err != nil {
return ErrorResult(err.Error()) return ErrorResult(err.Error())
} }
return formatDirEntries(entries) return formatDirEntries(entries)
} }
func formatDirEntries(entries []os.DirEntry) *ToolResult { func formatDirEntries(entries []os.DirEntry) *ToolResult {
var result strings.Builder var result strings.Builder
for _, entry := range entries { for _, entry := range entries {
if entry.IsDir() { if entry.IsDir() {
result.WriteString("DIR: ") result.WriteString("DIR: ")
} else { } else {
result.WriteString("FILE: ") result.WriteString("FILE: ")
} }
result.WriteString(entry.Name()) result.WriteString(entry.Name())
result.WriteByte('\n') result.WriteByte('\n')
} }
return NewToolResult(result.String()) return NewToolResult(result.String())
} }
// fileSystem abstracts reading, writing, and listing files, allowing both // fileSystem abstracts reading, writing, and listing files, allowing both
// unrestricted (host filesystem) and sandbox (os.Root) implementations to share the same polymorphic interface. // unrestricted (host filesystem) and sandbox (os.Root) implementations to share the same polymorphic interface.
type fileSystem interface { type fileSystem interface {
ReadFile(path string) ([]byte, error) ReadFile(path string) ([]byte, error)
WriteFile(path string, data []byte) error WriteFile(path string, data []byte) error
ReadDir(path string) ([]os.DirEntry, error) ReadDir(path string) ([]os.DirEntry, error)
} }
// hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem. // hostFs is an unrestricted fileReadWriter that operates directly on the host filesystem.
type hostFs struct{} type hostFs struct{}
func (h *hostFs) ReadFile(path string) ([]byte, error) { func (h *hostFs) ReadFile(path string) ([]byte, error) {
@ -268,11 +308,14 @@ func (h *hostFs) ReadFile(path string) ([]byte, error) {
if os.IsNotExist(err) { if os.IsNotExist(err) {
return nil, fmt.Errorf("failed to read file: file not found: %w", err) return nil, fmt.Errorf("failed to read file: file not found: %w", err)
} }
if os.IsPermission(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: access denied: %w", err)
} }
return nil, fmt.Errorf("failed to read file: %w", err) return nil, fmt.Errorf("failed to read file: %w", err)
} }
return content, nil return content, nil
} }
@ -281,16 +324,20 @@ func (h *hostFs) ReadDir(path string) ([]os.DirEntry, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to read directory: %w", err) return nil, fmt.Errorf("failed to read directory: %w", err)
} }
return entries, nil return entries, nil
} }
func (h *hostFs) WriteFile(path string, data []byte) error { func (h *hostFs) WriteFile(path string, data []byte) error {
// Use unified atomic write utility with explicit sync for flash storage reliability. // Use unified atomic write utility with explicit sync for flash storage reliability.
// Using 0o600 (owner read/write only) for secure default permissions. // Using 0o600 (owner read/write only) for secure default permissions.
return fileutil.WriteFileAtomic(path, data, 0o600) return fileutil.WriteFileAtomic(path, data, 0o600)
} }
// sandboxFs is a sandboxed fileSystem that operates within a strictly defined workspace using os.Root. // sandboxFs is a sandboxed fileSystem that operates within a strictly defined workspace using os.Root.
type sandboxFs struct { type sandboxFs struct {
workspace string workspace string
} }
@ -304,6 +351,7 @@ func (r *sandboxFs) execute(path string, fn func(root *os.Root, relPath string)
if err != nil { if err != nil {
return fmt.Errorf("failed to open workspace: %w", err) return fmt.Errorf("failed to open workspace: %w", err)
} }
defer root.Close() defer root.Close()
relPath, err := getSafeRelPath(r.workspace, path) 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) { func (r *sandboxFs) ReadFile(path string) ([]byte, error) {
var content []byte var content []byte
err := r.execute(path, func(root *os.Root, relPath string) error { err := r.execute(path, func(root *os.Root, relPath string) error {
fileContent, err := root.ReadFile(relPath) fileContent, err := root.ReadFile(relPath)
if err != nil { if err != nil {
if os.IsNotExist(err) { if os.IsNotExist(err) {
return fmt.Errorf("failed to read file: file not found: %w", err) return fmt.Errorf("failed to read file: file not found: %w", err)
} }
// os.Root returns "escapes from parent" for paths outside the root // os.Root returns "escapes from parent" for paths outside the root
if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") || if os.IsPermission(err) || strings.Contains(err.Error(), "escapes from parent") ||
strings.Contains(err.Error(), "permission denied") { strings.Contains(err.Error(), "permission denied") {
return fmt.Errorf("failed to read file: access denied: %w", err) return fmt.Errorf("failed to read file: access denied: %w", err)
} }
return fmt.Errorf("failed to read file: %w", err) return fmt.Errorf("failed to read file: %w", err)
} }
content = fileContent content = fileContent
return nil return nil
}) })
return content, err return content, err
} }
func (r *sandboxFs) WriteFile(path string, data []byte) error { func (r *sandboxFs) WriteFile(path string, data []byte) error {
return r.execute(path, func(root *os.Root, relPath string) error { return r.execute(path, func(root *os.Root, relPath string) error {
dir := filepath.Dir(relPath) dir := filepath.Dir(relPath)
if dir != "." && dir != "/" { if dir != "." && dir != "/" {
if err := root.MkdirAll(dir, 0o755); err != nil { if err := root.MkdirAll(dir, 0o755); err != nil {
return fmt.Errorf("failed to create parent directories: %w", err) 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. // Use atomic write pattern with explicit sync for flash storage reliability.
// Using 0o600 (owner read/write only) for secure default permissions. // Using 0o600 (owner read/write only) for secure default permissions.
tmpRelPath := fmt.Sprintf(".tmp-%d-%d", os.Getpid(), time.Now().UnixNano()) 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) tmpFile, err := root.OpenFile(tmpRelPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
if err != nil { if err != nil {
root.Remove(tmpRelPath) root.Remove(tmpRelPath)
return fmt.Errorf("failed to open temp file: %w", err) return fmt.Errorf("failed to open temp file: %w", err)
} }
if _, err := tmpFile.Write(data); err != nil { if _, err := tmpFile.Write(data); err != nil {
tmpFile.Close() tmpFile.Close()
root.Remove(tmpRelPath) root.Remove(tmpRelPath)
return fmt.Errorf("failed to write temp file: %w", err) return fmt.Errorf("failed to write temp file: %w", err)
} }
// CRITICAL: Force sync to storage medium before rename. // CRITICAL: Force sync to storage medium before rename.
// This ensures data is physically written to disk, not just cached. // This ensures data is physically written to disk, not just cached.
if err := tmpFile.Sync(); err != nil { if err := tmpFile.Sync(); err != nil {
tmpFile.Close() tmpFile.Close()
root.Remove(tmpRelPath) root.Remove(tmpRelPath)
return fmt.Errorf("failed to sync temp file: %w", err) return fmt.Errorf("failed to sync temp file: %w", err)
} }
if err := tmpFile.Close(); err != nil { if err := tmpFile.Close(); err != nil {
root.Remove(tmpRelPath) root.Remove(tmpRelPath)
return fmt.Errorf("failed to close temp file: %w", err) return fmt.Errorf("failed to close temp file: %w", err)
} }
if err := root.Rename(tmpRelPath, relPath); err != nil { if err := root.Rename(tmpRelPath, relPath); err != nil {
root.Remove(tmpRelPath) root.Remove(tmpRelPath)
return fmt.Errorf("failed to rename temp file over target: %w", err) return fmt.Errorf("failed to rename temp file over target: %w", err)
} }
// Sync directory to ensure rename is durable // Sync directory to ensure rename is durable
if dirFile, err := root.Open("."); err == nil { if dirFile, err := root.Open("."); err == nil {
_ = dirFile.Sync() _ = dirFile.Sync()
dirFile.Close() dirFile.Close()
} }
@ -390,26 +460,33 @@ func (r *sandboxFs) WriteFile(path string, data []byte) error {
func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) { func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) {
var entries []os.DirEntry var entries []os.DirEntry
err := r.execute(path, func(root *os.Root, relPath string) error { err := r.execute(path, func(root *os.Root, relPath string) error {
dirEntries, err := fs.ReadDir(root.FS(), relPath) dirEntries, err := fs.ReadDir(root.FS(), relPath)
if err != nil { if err != nil {
return err return err
} }
entries = dirEntries entries = dirEntries
return nil return nil
}) })
return entries, err return entries, err
} }
// Helper to get a safe relative path for os.Root usage // Helper to get a safe relative path for os.Root usage
func getSafeRelPath(workspace, path string) (string, error) { func getSafeRelPath(workspace, path string) (string, error) {
if workspace == "" { if workspace == "" {
return "", fmt.Errorf("workspace is not defined") return "", fmt.Errorf("workspace is not defined")
} }
rel := filepath.Clean(path) rel := filepath.Clean(path)
if filepath.IsAbs(rel) { if filepath.IsAbs(rel) {
var err error var err error
rel, err = filepath.Rel(workspace, rel) rel, err = filepath.Rel(workspace, rel)
if err != nil { if err != nil {
return "", fmt.Errorf("failed to calculate relative path: %w", err) return "", fmt.Errorf("failed to calculate relative path: %w", err)

View file

@ -12,13 +12,18 @@ import (
) )
// TestFilesystemTool_ReadFile_Success verifies successful file reading // TestFilesystemTool_ReadFile_Success verifies successful file reading
func TestFilesystemTool_ReadFile_Success(t *testing.T) { func TestFilesystemTool_ReadFile_Success(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "test.txt") testFile := filepath.Join(tmpDir, "test.txt")
os.WriteFile(testFile, []byte("test content"), 0o644) os.WriteFile(testFile, []byte("test content"), 0o644)
tool := NewReadFileTool("", false) tool := NewReadFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
} }
@ -26,26 +31,33 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// ForLLM should contain file content // ForLLM should contain file content
if !strings.Contains(result.ForLLM, "test content") { if !strings.Contains(result.ForLLM, "test content") {
t.Errorf("Expected ForLLM to contain 'test content', got: %s", result.ForLLM) t.Errorf("Expected ForLLM to contain 'test content', got: %s", result.ForLLM)
} }
// ReadFile returns NewToolResult which only sets ForLLM, not ForUser // ReadFile returns NewToolResult which only sets ForLLM, not ForUser
// This is the expected behavior - file content goes to LLM, not directly to user // This is the expected behavior - file content goes to LLM, not directly to user
if result.ForUser != "" { if result.ForUser != "" {
t.Errorf("Expected ForUser to be empty for NewToolResult, got: %s", result.ForUser) t.Errorf("Expected ForUser to be empty for NewToolResult, got: %s", result.ForUser)
} }
} }
// TestFilesystemTool_ReadFile_NotFound verifies error handling for missing file // TestFilesystemTool_ReadFile_NotFound verifies error handling for missing file
func TestFilesystemTool_ReadFile_NotFound(t *testing.T) { func TestFilesystemTool_ReadFile_NotFound(t *testing.T) {
tool := NewReadFileTool("", false) tool := NewReadFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "/nonexistent_file_12345.txt", "path": "/nonexistent_file_12345.txt",
} }
@ -53,107 +65,135 @@ func TestFilesystemTool_ReadFile_NotFound(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Failure should be marked as error // Failure should be marked as error
if !result.IsError { if !result.IsError {
t.Errorf("Expected error for missing file, got IsError=false") t.Errorf("Expected error for missing file, got IsError=false")
} }
// Should contain error message // Should contain error message
if !strings.Contains(result.ForLLM, "failed to read") && !strings.Contains(result.ForUser, "failed to read") { 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) t.Errorf("Expected error message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
} }
} }
// TestFilesystemTool_ReadFile_MissingPath verifies error handling for missing path // TestFilesystemTool_ReadFile_MissingPath verifies error handling for missing path
func TestFilesystemTool_ReadFile_MissingPath(t *testing.T) { func TestFilesystemTool_ReadFile_MissingPath(t *testing.T) {
tool := &ReadFileTool{} tool := &ReadFileTool{}
ctx := context.Background() ctx := context.Background()
args := map[string]any{} args := map[string]any{}
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when path is missing") t.Errorf("Expected error when path is missing")
} }
// Should mention required parameter // Should mention required parameter
if !strings.Contains(result.ForLLM, "path is required") && !strings.Contains(result.ForUser, "path is required") { 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) t.Errorf("Expected 'path is required' message, got ForLLM: %s", result.ForLLM)
} }
} }
// TestFilesystemTool_WriteFile_Success verifies successful file writing // TestFilesystemTool_WriteFile_Success verifies successful file writing
func TestFilesystemTool_WriteFile_Success(t *testing.T) { func TestFilesystemTool_WriteFile_Success(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "newfile.txt") testFile := filepath.Join(tmpDir, "newfile.txt")
tool := NewWriteFileTool("", false) tool := NewWriteFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"content": "hello world", "content": "hello world",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// WriteFile returns SilentResult // WriteFile returns SilentResult
if !result.Silent { if !result.Silent {
t.Errorf("Expected Silent=true for WriteFile, got false") t.Errorf("Expected Silent=true for WriteFile, got false")
} }
// ForUser should be empty (silent result) // ForUser should be empty (silent result)
if result.ForUser != "" { if result.ForUser != "" {
t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser)
} }
// Verify file was actually written // Verify file was actually written
content, err := os.ReadFile(testFile) content, err := os.ReadFile(testFile)
if err != nil { if err != nil {
t.Fatalf("Failed to read written file: %v", err) t.Fatalf("Failed to read written file: %v", err)
} }
if string(content) != "hello world" { if string(content) != "hello world" {
t.Errorf("Expected file content 'hello world', got: %s", string(content)) t.Errorf("Expected file content 'hello world', got: %s", string(content))
} }
} }
// TestFilesystemTool_WriteFile_CreateDir verifies directory creation // TestFilesystemTool_WriteFile_CreateDir verifies directory creation
func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) { func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "subdir", "newfile.txt") testFile := filepath.Join(tmpDir, "subdir", "newfile.txt")
tool := NewWriteFileTool("", false) tool := NewWriteFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"content": "test", "content": "test",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success with directory creation, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success with directory creation, got IsError=true: %s", result.ForLLM)
} }
// Verify directory was created and file written // Verify directory was created and file written
content, err := os.ReadFile(testFile) content, err := os.ReadFile(testFile)
if err != nil { if err != nil {
t.Fatalf("Failed to read written file: %v", err) t.Fatalf("Failed to read written file: %v", err)
} }
if string(content) != "test" { if string(content) != "test" {
t.Errorf("Expected file content 'test', got: %s", string(content)) t.Errorf("Expected file content 'test', got: %s", string(content))
} }
} }
// TestFilesystemTool_WriteFile_MissingPath verifies error handling for missing path // TestFilesystemTool_WriteFile_MissingPath verifies error handling for missing path
func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) { func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) {
tool := NewWriteFileTool("", false) tool := NewWriteFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"content": "test", "content": "test",
} }
@ -161,15 +201,19 @@ func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when path is missing") t.Errorf("Expected error when path is missing")
} }
} }
// TestFilesystemTool_WriteFile_MissingContent verifies error handling for missing content // TestFilesystemTool_WriteFile_MissingContent verifies error handling for missing content
func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) { func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) {
tool := NewWriteFileTool("", false) tool := NewWriteFileTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "/tmp/test.txt", "path": "/tmp/test.txt",
} }
@ -177,26 +221,35 @@ func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when content is missing") t.Errorf("Expected error when content is missing")
} }
// Should mention required parameter // Should mention required parameter
if !strings.Contains(result.ForLLM, "content is required") && if !strings.Contains(result.ForLLM, "content is required") &&
!strings.Contains(result.ForUser, "content is required") { !strings.Contains(result.ForUser, "content is required") {
t.Errorf("Expected 'content is required' message, got ForLLM: %s", result.ForLLM) t.Errorf("Expected 'content is required' message, got ForLLM: %s", result.ForLLM)
} }
} }
// TestFilesystemTool_ListDir_Success verifies successful directory listing // TestFilesystemTool_ListDir_Success verifies successful directory listing
func TestFilesystemTool_ListDir_Success(t *testing.T) { func TestFilesystemTool_ListDir_Success(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
os.WriteFile(filepath.Join(tmpDir, "file1.txt"), []byte("content"), 0o644) os.WriteFile(filepath.Join(tmpDir, "file1.txt"), []byte("content"), 0o644)
os.WriteFile(filepath.Join(tmpDir, "file2.txt"), []byte("content"), 0o644) os.WriteFile(filepath.Join(tmpDir, "file2.txt"), []byte("content"), 0o644)
os.Mkdir(filepath.Join(tmpDir, "subdir"), 0o755) os.Mkdir(filepath.Join(tmpDir, "subdir"), 0o755)
tool := NewListDirTool("", false) tool := NewListDirTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": tmpDir, "path": tmpDir,
} }
@ -204,23 +257,29 @@ func TestFilesystemTool_ListDir_Success(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// Should list files and directories // Should list files and directories
if !strings.Contains(result.ForLLM, "file1.txt") || !strings.Contains(result.ForLLM, "file2.txt") { if !strings.Contains(result.ForLLM, "file1.txt") || !strings.Contains(result.ForLLM, "file2.txt") {
t.Errorf("Expected files in listing, got: %s", result.ForLLM) t.Errorf("Expected files in listing, got: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "subdir") { if !strings.Contains(result.ForLLM, "subdir") {
t.Errorf("Expected subdir in listing, got: %s", result.ForLLM) t.Errorf("Expected subdir in listing, got: %s", result.ForLLM)
} }
} }
// TestFilesystemTool_ListDir_NotFound verifies error handling for non-existent directory // TestFilesystemTool_ListDir_NotFound verifies error handling for non-existent directory
func TestFilesystemTool_ListDir_NotFound(t *testing.T) { func TestFilesystemTool_ListDir_NotFound(t *testing.T) {
tool := NewListDirTool("", false) tool := NewListDirTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"path": "/nonexistent_directory_12345", "path": "/nonexistent_directory_12345",
} }
@ -228,49 +287,61 @@ func TestFilesystemTool_ListDir_NotFound(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Failure should be marked as error // Failure should be marked as error
if !result.IsError { if !result.IsError {
t.Errorf("Expected error for non-existent directory, got IsError=false") t.Errorf("Expected error for non-existent directory, got IsError=false")
} }
// Should contain error message // Should contain error message
if !strings.Contains(result.ForLLM, "failed to read") && !strings.Contains(result.ForUser, "failed to read") { 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) t.Errorf("Expected error message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser)
} }
} }
// TestFilesystemTool_ListDir_DefaultPath verifies default to current directory // TestFilesystemTool_ListDir_DefaultPath verifies default to current directory
func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) { func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) {
tool := NewListDirTool("", false) tool := NewListDirTool("", false)
ctx := context.Background() ctx := context.Background()
args := map[string]any{} args := map[string]any{}
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should use "." as default path // Should use "." as default path
if result.IsError { if result.IsError {
t.Errorf("Expected success with default path '.', got IsError=true: %s", result.ForLLM) t.Errorf("Expected success with default path '.', got IsError=true: %s", result.ForLLM)
} }
} }
// Block paths that look inside workspace but point outside via symlink. // Block paths that look inside workspace but point outside via symlink.
func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) { func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) {
root := t.TempDir() root := t.TempDir()
workspace := filepath.Join(root, "workspace") workspace := filepath.Join(root, "workspace")
if err := os.MkdirAll(workspace, 0o755); err != nil { if err := os.MkdirAll(workspace, 0o755); err != nil {
t.Fatalf("failed to create workspace: %v", err) t.Fatalf("failed to create workspace: %v", err)
} }
secret := filepath.Join(root, "secret.txt") secret := filepath.Join(root, "secret.txt")
if err := os.WriteFile(secret, []byte("top secret"), 0o644); err != nil { if err := os.WriteFile(secret, []byte("top secret"), 0o644); err != nil {
t.Fatalf("failed to write secret file: %v", err) t.Fatalf("failed to write secret file: %v", err)
} }
link := filepath.Join(workspace, "leak.txt") link := filepath.Join(workspace, "leak.txt")
if err := os.Symlink(secret, link); err != nil { if err := os.Symlink(secret, link); err != nil {
t.Skipf("symlink not supported in this environment: %v", err) t.Skipf("symlink not supported in this environment: %v", err)
} }
tool := NewReadFileTool(workspace, true) tool := NewReadFileTool(workspace, true)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"path": link, "path": link,
}) })
@ -278,10 +349,15 @@ func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) {
if !result.IsError { if !result.IsError {
t.Fatalf("expected symlink escape to be blocked") t.Fatalf("expected symlink escape to be blocked")
} }
// os.Root might return different errors depending on platform/implementation // os.Root might return different errors depending on platform/implementation
// but it definitely should error. // but it definitely should error.
// Our wrapper returns "access denied or file not found" // Our wrapper returns "access denied or file not found"
if !strings.Contains(result.ForLLM, "access denied") && !strings.Contains(result.ForLLM, "file not found") && if !strings.Contains(result.ForLLM, "access denied") && !strings.Contains(result.ForLLM, "file not found") &&
!strings.Contains(result.ForLLM, "no such file") { !strings.Contains(result.ForLLM, "no such file") {
t.Fatalf("expected symlink escape error, got: %s", result.ForLLM) 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="" tool := NewReadFileTool("", true) // restrict=true but workspace=""
// Try to read a sensitive file (simulated by a temp file outside workspace) // Try to read a sensitive file (simulated by a temp file outside workspace)
tmpDir := t.TempDir() tmpDir := t.TempDir()
secretFile := filepath.Join(tmpDir, "shadow") secretFile := filepath.Join(tmpDir, "shadow")
os.WriteFile(secretFile, []byte("secret data"), 0o600) os.WriteFile(secretFile, []byte("secret data"), 0o600)
result := tool.Execute(context.Background(), map[string]any{ 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) // 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) assert.True(t, result.IsError, "Security Regression: Empty workspace allowed access! content: %s", result.ForLLM)
// Verify it failed for the right reason // Verify it failed for the right reason
assert.Contains(t, result.ForLLM, "workspace is not defined", "Expected 'workspace is not defined' error") 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: // 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. // single dir, deeply nested dirs, already-existing dirs, and a file blocking a directory path.
func TestRootMkdirAll(t *testing.T) { func TestRootMkdirAll(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
root, err := os.OpenRoot(workspace) root, err := os.OpenRoot(workspace)
if err != nil { if err != nil {
t.Fatalf("failed to open root: %v", err) t.Fatalf("failed to open root: %v", err)
} }
defer root.Close() defer root.Close()
// Case 1: Single directory // Case 1: Single directory
err = root.MkdirAll("dir1", 0o755) err = root.MkdirAll("dir1", 0o755)
assert.NoError(t, err) assert.NoError(t, err)
_, err = os.Stat(filepath.Join(workspace, "dir1")) _, err = os.Stat(filepath.Join(workspace, "dir1"))
assert.NoError(t, err) assert.NoError(t, err)
// Case 2: Deeply nested directory // Case 2: Deeply nested directory
err = root.MkdirAll("a/b/c/d", 0o755) err = root.MkdirAll("a/b/c/d", 0o755)
assert.NoError(t, err) assert.NoError(t, err)
_, err = os.Stat(filepath.Join(workspace, "a/b/c/d")) _, err = os.Stat(filepath.Join(workspace, "a/b/c/d"))
assert.NoError(t, err) assert.NoError(t, err)
// Case 3: Already exists — must be idempotent // Case 3: Already exists — must be idempotent
err = root.MkdirAll("a/b/c/d", 0o755) err = root.MkdirAll("a/b/c/d", 0o755)
assert.NoError(t, err) assert.NoError(t, err)
// Case 4: A regular file blocks directory creation — must error // Case 4: A regular file blocks directory creation — must error
err = os.WriteFile(filepath.Join(workspace, "file_exists"), []byte("data"), 0o644) err = os.WriteFile(filepath.Join(workspace, "file_exists"), []byte("data"), 0o644)
assert.NoError(t, err) assert.NoError(t, err)
err = root.MkdirAll("file_exists", 0o755) err = root.MkdirAll("file_exists", 0o755)
assert.Error(t, err, "expected error when a file exists at the directory path") assert.Error(t, err, "expected error when a file exists at the directory path")
} }
func TestFilesystemTool_WriteFile_Restricted_CreateDir(t *testing.T) { func TestFilesystemTool_WriteFile_Restricted_CreateDir(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
tool := NewWriteFileTool(workspace, true) tool := NewWriteFileTool(workspace, true)
ctx := context.Background() ctx := context.Background()
testFile := "deep/nested/path/to/file.txt" testFile := "deep/nested/path/to/file.txt"
content := "deep content" content := "deep content"
args := map[string]any{ args := map[string]any{
"path": testFile, "path": testFile,
"content": content, "content": content,
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM)
// Verify file content // Verify file content
actualPath := filepath.Join(workspace, testFile) actualPath := filepath.Join(workspace, testFile)
data, err := os.ReadFile(actualPath) data, err := os.ReadFile(actualPath)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, content, string(data)) assert.Equal(t, content, string(data))
} }
// TestHostFs_Read_PermissionDenied verifies that hostFs.ReadFile surfaces access denied errors. // TestHostFs_Read_PermissionDenied verifies that hostFs.ReadFile surfaces access denied errors.
func TestHostFs_Read_PermissionDenied(t *testing.T) { func TestHostFs_Read_PermissionDenied(t *testing.T) {
if os.Getuid() == 0 { if os.Getuid() == 0 {
t.Skip("skipping permission test: running as root") t.Skip("skipping permission test: running as root")
} }
tmpDir := t.TempDir() tmpDir := t.TempDir()
protected := filepath.Join(tmpDir, "protected.txt") protected := filepath.Join(tmpDir, "protected.txt")
err := os.WriteFile(protected, []byte("secret"), 0o000) err := os.WriteFile(protected, []byte("secret"), 0o000)
assert.NoError(t, err) assert.NoError(t, err)
defer os.Chmod(protected, 0o644) // ensure cleanup defer os.Chmod(protected, 0o644) // ensure cleanup
_, err = (&hostFs{}).ReadFile(protected) _, err = (&hostFs{}).ReadFile(protected)
assert.Error(t, err) assert.Error(t, err)
assert.Contains(t, err.Error(), "access denied") assert.Contains(t, err.Error(), "access denied")
} }
// TestHostFs_Read_Directory verifies that hostFs.ReadFile returns an error when given a directory path. // TestHostFs_Read_Directory verifies that hostFs.ReadFile returns an error when given a directory path.
func TestHostFs_Read_Directory(t *testing.T) { func TestHostFs_Read_Directory(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
_, err := (&hostFs{}).ReadFile(tmpDir) _, err := (&hostFs{}).ReadFile(tmpDir)
assert.Error(t, err, "expected error when reading a directory as a file") 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. // TestSandboxFs_Read_Directory verifies that sandboxFs.ReadFile returns an error when given a directory.
func TestSandboxFs_Read_Directory(t *testing.T) { func TestSandboxFs_Read_Directory(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
root, err := os.OpenRoot(workspace) root, err := os.OpenRoot(workspace)
assert.NoError(t, err) assert.NoError(t, err)
defer root.Close() defer root.Close()
// Create a subdirectory // Create a subdirectory
err = root.Mkdir("subdir", 0o755) err = root.Mkdir("subdir", 0o755)
assert.NoError(t, err) assert.NoError(t, err)
_, err = (&sandboxFs{workspace: workspace}).ReadFile("subdir") _, err = (&sandboxFs{workspace: workspace}).ReadFile("subdir")
assert.Error(t, err, "expected error when reading a directory as a file") assert.Error(t, err, "expected error when reading a directory as a file")
} }
// TestHostFs_Write_ParentDirMissing verifies that hostFs.WriteFile creates parent dirs automatically. // TestHostFs_Write_ParentDirMissing verifies that hostFs.WriteFile creates parent dirs automatically.
func TestHostFs_Write_ParentDirMissing(t *testing.T) { func TestHostFs_Write_ParentDirMissing(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
target := filepath.Join(tmpDir, "a", "b", "c", "file.txt") target := filepath.Join(tmpDir, "a", "b", "c", "file.txt")
err := (&hostFs{}).WriteFile(target, []byte("hello")) err := (&hostFs{}).WriteFile(target, []byte("hello"))
assert.NoError(t, err) assert.NoError(t, err)
data, err := os.ReadFile(target) data, err := os.ReadFile(target)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, "hello", string(data)) assert.Equal(t, "hello", string(data))
} }
// TestSandboxFs_Write_ParentDirMissing verifies that sandboxFs.WriteFile creates // TestSandboxFs_Write_ParentDirMissing verifies that sandboxFs.WriteFile creates
// nested parent directories automatically within the sandbox. // nested parent directories automatically within the sandbox.
func TestSandboxFs_Write_ParentDirMissing(t *testing.T) { func TestSandboxFs_Write_ParentDirMissing(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
relPath := "x/y/z/file.txt" relPath := "x/y/z/file.txt"
err := (&sandboxFs{workspace: workspace}).WriteFile(relPath, []byte("nested")) err := (&sandboxFs{workspace: workspace}).WriteFile(relPath, []byte("nested"))
assert.NoError(t, err) assert.NoError(t, err)
data, err := os.ReadFile(filepath.Join(workspace, relPath)) data, err := os.ReadFile(filepath.Join(workspace, relPath))
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, "nested", string(data)) assert.Equal(t, "nested", string(data))
} }
// TestHostFs_Write verifies the hostFs.WriteFile helper function // TestHostFs_Write verifies the hostFs.WriteFile helper function
func TestHostFs_Write(t *testing.T) { func TestHostFs_Write(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
testFile := filepath.Join(tmpDir, "atomic_test.txt") testFile := filepath.Join(tmpDir, "atomic_test.txt")
testData := []byte("atomic test content") testData := []byte("atomic test content")
err := (&hostFs{}).WriteFile(testFile, testData) err := (&hostFs{}).WriteFile(testFile, testData)
assert.NoError(t, err) assert.NoError(t, err)
content, err := os.ReadFile(testFile) content, err := os.ReadFile(testFile)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, testData, content) assert.Equal(t, testData, content)
// Verify it overwrites correctly // Verify it overwrites correctly
newData := []byte("new atomic content") newData := []byte("new atomic content")
err = (&hostFs{}).WriteFile(testFile, newData) err = (&hostFs{}).WriteFile(testFile, newData)
assert.NoError(t, err) assert.NoError(t, err)
content, err = os.ReadFile(testFile) content, err = os.ReadFile(testFile)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, newData, content) assert.Equal(t, newData, content)
} }
// TestSandboxFs_Write verifies the sandboxFs.WriteFile helper function // TestSandboxFs_Write verifies the sandboxFs.WriteFile helper function
func TestSandboxFs_Write(t *testing.T) { func TestSandboxFs_Write(t *testing.T) {
tmpDir := t.TempDir() tmpDir := t.TempDir()
relPath := "atomic_root_test.txt" relPath := "atomic_root_test.txt"
testData := []byte("atomic root test content") testData := []byte("atomic root test content")
erw := &sandboxFs{workspace: tmpDir} erw := &sandboxFs{workspace: tmpDir}
err := erw.WriteFile(relPath, testData) err := erw.WriteFile(relPath, testData)
assert.NoError(t, err) assert.NoError(t, err)
root, err := os.OpenRoot(tmpDir) root, err := os.OpenRoot(tmpDir)
assert.NoError(t, err) assert.NoError(t, err)
defer root.Close() defer root.Close()
f, err := root.Open(relPath) f, err := root.Open(relPath)
assert.NoError(t, err) assert.NoError(t, err)
defer f.Close() defer f.Close()
content, err := io.ReadAll(f) content, err := io.ReadAll(f)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, testData, content) assert.Equal(t, testData, content)
// Verify it overwrites correctly // Verify it overwrites correctly
newData := []byte("new root atomic content") newData := []byte("new root atomic content")
err = erw.WriteFile(relPath, newData) err = erw.WriteFile(relPath, newData)
assert.NoError(t, err) assert.NoError(t, err)
f2, err := root.Open(relPath) f2, err := root.Open(relPath)
assert.NoError(t, err) assert.NoError(t, err)
defer f2.Close() defer f2.Close()
content, err = io.ReadAll(f2) content, err = io.ReadAll(f2)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, newData, content) assert.Equal(t, newData, content)
} }
// TestValidatePath_OutsideWorkspace_IncludesPath verifies that the access // TestValidatePath_OutsideWorkspace_IncludesPath verifies that the access
// denied error includes the workspace path so the caller knows the boundary. // denied error includes the workspace path so the caller knows the boundary.
func TestValidatePath_OutsideWorkspace_IncludesPath(t *testing.T) { func TestValidatePath_OutsideWorkspace_IncludesPath(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
outsidePath := filepath.Join(t.TempDir(), "secret.txt") outsidePath := filepath.Join(t.TempDir(), "secret.txt")
_, err := validatePath(outsidePath, workspace, true) _, err := validatePath(outsidePath, workspace, true)
assert.Error(t, err) assert.Error(t, err)
assert.Contains(t, err.Error(), "access denied") assert.Contains(t, err.Error(), "access denied")
assert.Contains(t, err.Error(), workspace) assert.Contains(t, err.Error(), workspace)
} }

View file

@ -12,35 +12,49 @@ import (
) )
// worktreeInfoKey is the context key for passing WorktreeInfo to tools. // worktreeInfoKey is the context key for passing WorktreeInfo to tools.
type worktreeInfoKey struct{} type worktreeInfoKey struct{}
// WithWorktreeInfo returns a context carrying the active WorktreeInfo. // WithWorktreeInfo returns a context carrying the active WorktreeInfo.
func WithWorktreeInfo(ctx context.Context, wt *git.WorktreeInfo) context.Context { func WithWorktreeInfo(ctx context.Context, wt *git.WorktreeInfo) context.Context {
return context.WithValue(ctx, worktreeInfoKey{}, wt) return context.WithValue(ctx, worktreeInfoKey{}, wt)
} }
// WorktreeInfoFromCtx extracts the WorktreeInfo from context, or nil. // WorktreeInfoFromCtx extracts the WorktreeInfo from context, or nil.
func WorktreeInfoFromCtx(ctx context.Context) *git.WorktreeInfo { func WorktreeInfoFromCtx(ctx context.Context) *git.WorktreeInfo {
if v, ok := ctx.Value(worktreeInfoKey{}).(*git.WorktreeInfo); ok { if v, ok := ctx.Value(worktreeInfoKey{}).(*git.WorktreeInfo); ok {
return v return v
} }
return nil return nil
} }
// protectedBranches are branch names that can never be pushed to. // protectedBranches are branch names that can never be pushed to.
var protectedBranches = regexp.MustCompile(`^(main|master|develop|release/.*)$`) var protectedBranches = regexp.MustCompile(`^(main|master|develop|release/.*)$`)
// GitPushTool implements safe git push restricted to worktree branches. // GitPushTool implements safe git push restricted to worktree branches.
// //
// Safety invariants: // Safety invariants:
// - Only works inside a worktree (WorktreeInfo must be in context) // - Only works inside a worktree (WorktreeInfo must be in context)
// - Pushes only the worktree's branch — no arbitrary branch targets // - Pushes only the worktree's branch — no arbitrary branch targets
// - Protected branches (main, master, develop, release/*) are blocked // - Protected branches (main, master, develop, release/*) are blocked
// - Force push is never allowed // - Force push is never allowed
// - Auto-commits uncommitted changes before pushing // - Auto-commits uncommitted changes before pushing
type GitPushTool struct{} type GitPushTool struct{}
// NewGitPushTool creates a GitPushTool. // NewGitPushTool creates a GitPushTool.
func NewGitPushTool() *GitPushTool { func NewGitPushTool() *GitPushTool {
return &GitPushTool{} return &GitPushTool{}
} }
@ -49,94 +63,138 @@ func (t *GitPushTool) Name() string { return "git_push" }
func (t *GitPushTool) Description() string { func (t *GitPushTool) Description() string {
return "Push the current worktree branch to origin. Only works inside a git worktree. " + return "Push the current worktree branch to origin. Only works inside a git worktree. " +
"Auto-commits uncommitted changes before pushing. " + "Auto-commits uncommitted changes before pushing. " +
"Protected branches (main, master, develop) cannot be pushed to. Force push is not allowed." "Protected branches (main, master, develop) cannot be pushed to. Force push is not allowed."
} }
func (t *GitPushTool) Parameters() map[string]any { func (t *GitPushTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"commit_message": 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.", "description": "Commit message for uncommitted changes. If omitted, uncommitted changes are auto-committed with a default message.",
}, },
}, },
"required": []string{}, "required": []string{},
} }
} }
func (t *GitPushTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *GitPushTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
wt := WorktreeInfoFromCtx(ctx) wt := WorktreeInfoFromCtx(ctx)
if wt == nil { if wt == nil {
return ErrorResult( return ErrorResult(
"git_push requires an active worktree.\n" + "git_push requires an active worktree.\n" +
"This tool can only be used during worktree-based sessions " + "This tool can only be used during worktree-based sessions " +
"(e.g., heartbeat tasks or plan executing phase).\n" + "(e.g., heartbeat tasks or plan executing phase).\n" +
"The worktree provides the branch name and isolation boundary — " + "The worktree provides the branch name and isolation boundary — " +
"without it, git_push cannot determine which branch to push.") "without it, git_push cannot determine which branch to push.")
} }
branch := wt.Branch branch := wt.Branch
if branch == "" { if branch == "" {
return ErrorResult( return ErrorResult(
"worktree has no branch name.\n" + "worktree has no branch name.\n" +
"The WorktreeInfo was set but Branch is empty. " + "The WorktreeInfo was set but Branch is empty. " +
"This is an internal error — the worktree may not have been created correctly.") "This is an internal error — the worktree may not have been created correctly.")
} }
// Block protected branches // Block protected branches
if protectedBranches.MatchString(branch) { if protectedBranches.MatchString(branch) {
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"cannot push to protected branch %q.\n"+ "cannot push to protected branch %q.\n"+
"Protected branches (main, master, develop, release/*) are blocked to prevent "+ "Protected branches (main, master, develop, release/*) are blocked to prevent "+
"accidental overwrites. Work should be done on feature branches created by worktrees.", "accidental overwrites. Work should be done on feature branches created by worktrees.",
branch)) branch))
} }
// Auto-commit uncommitted changes // Auto-commit uncommitted changes
if git.HasUncommittedChanges(wt.Path) { if git.HasUncommittedChanges(wt.Path) {
commitMsg := "auto: save before push" commitMsg := "auto: save before push"
if msg, ok := args["commit_message"].(string); ok && msg != "" { if msg, ok := args["commit_message"].(string); ok && msg != "" {
commitMsg = msg commitMsg = msg
} }
if err := git.AutoCommit(wt.Path, commitMsg); err != nil { if err := git.AutoCommit(wt.Path, commitMsg); err != nil {
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"auto-commit failed before push: %v\n"+ "auto-commit failed before push: %v\n"+
"git_push auto-commits uncommitted changes before pushing. "+ "git_push auto-commits uncommitted changes before pushing. "+
"The commit failed, so no push was attempted. "+ "The commit failed, so no push was attempted. "+
"Check if the worktree at %q is in a valid state (e.g., no merge conflicts).", "Check if the worktree at %q is in a valid state (e.g., no merge conflicts).",
err, wt.Path)) err, wt.Path))
} }
} }
// Check there are commits to push // Check there are commits to push
ahead := git.CommitsAhead(wt.RepoRoot, wt.BaseBranch, branch) ahead := git.CommitsAhead(wt.RepoRoot, wt.BaseBranch, branch)
if ahead == 0 { if ahead == 0 {
return NewToolResult(fmt.Sprintf( return NewToolResult(fmt.Sprintf(
"Nothing to push: branch %q has no commits ahead of %s.\n"+ "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.", "The branch is identical to the base. Make changes and commit before pushing.",
branch, wt.BaseBranch)) branch, wt.BaseBranch))
} }
// Push with -u (set upstream tracking) // Push with -u (set upstream tracking)
pushCtx, cancel := context.WithTimeout(ctx, 2*time.Minute) pushCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
defer cancel() defer cancel()
cmd := exec.CommandContext(pushCtx, "git", "push", "-u", "origin", branch) cmd := exec.CommandContext(pushCtx, "git", "push", "-u", "origin", branch)
cmd.Dir = wt.Path cmd.Dir = wt.Path
out, err := cmd.CombinedOutput() out, err := cmd.CombinedOutput()
output := strings.TrimSpace(string(out)) output := strings.TrimSpace(string(out))
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"git push failed for branch %q: %s\n%s\n"+ "git push failed for branch %q: %s\n%s\n"+
"Possible causes: network error, authentication failure, or remote rejected the push. "+ "Possible causes: network error, authentication failure, or remote rejected the push. "+
"If the remote branch has diverged, resolve the divergence in the worktree first — "+ "If the remote branch has diverged, resolve the divergence in the worktree first — "+
"force push is not available.", "force push is not available.",
branch, err, output)) branch, err, output))
} }
return NewToolResult(fmt.Sprintf("Pushed branch %q to origin (%d commit(s) ahead of %s)\n%s", return NewToolResult(fmt.Sprintf("Pushed branch %q to origin (%d commit(s) ahead of %s)\n%s",
branch, ahead, wt.BaseBranch, output)) branch, ahead, wt.BaseBranch, output))
} }

View file

@ -9,22 +9,29 @@ import (
) )
// TestGitPushTool_NoWorktree verifies that git_push fails without worktree context. // TestGitPushTool_NoWorktree verifies that git_push fails without worktree context.
func TestGitPushTool_NoWorktree(t *testing.T) { func TestGitPushTool_NoWorktree(t *testing.T) {
tool := NewGitPushTool() tool := NewGitPushTool()
result := tool.Execute(context.Background(), map[string]any{}) result := tool.Execute(context.Background(), map[string]any{})
if !result.IsError { if !result.IsError {
t.Fatal("expected error when no worktree in context") t.Fatal("expected error when no worktree in context")
} }
if result.ForLLM == "" { if result.ForLLM == "" {
t.Fatal("error message should not be empty") t.Fatal("error message should not be empty")
} }
// Verify helpful guidance is included // Verify helpful guidance is included
assertContains(t, result.ForLLM, "worktree") assertContains(t, result.ForLLM, "worktree")
assertContains(t, result.ForLLM, "heartbeat") assertContains(t, result.ForLLM, "heartbeat")
} }
// TestGitPushTool_ProtectedBranch verifies that protected branches are blocked. // TestGitPushTool_ProtectedBranch verifies that protected branches are blocked.
func TestGitPushTool_ProtectedBranch(t *testing.T) { func TestGitPushTool_ProtectedBranch(t *testing.T) {
tool := NewGitPushTool() tool := NewGitPushTool()
@ -34,39 +41,55 @@ func TestGitPushTool_ProtectedBranch(t *testing.T) {
t.Run(branch, func(t *testing.T) { t.Run(branch, func(t *testing.T) {
ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{
Branch: branch, Branch: branch,
BaseBranch: "main", BaseBranch: "main",
Path: t.TempDir(), Path: t.TempDir(),
RepoRoot: t.TempDir(), RepoRoot: t.TempDir(),
}) })
result := tool.Execute(ctx, map[string]any{}) result := tool.Execute(ctx, map[string]any{})
if !result.IsError { if !result.IsError {
t.Fatalf("expected error for protected branch %q", branch) t.Fatalf("expected error for protected branch %q", branch)
} }
assertContains(t, result.ForLLM, "protected") assertContains(t, result.ForLLM, "protected")
assertContains(t, result.ForLLM, branch) assertContains(t, result.ForLLM, branch)
}) })
} }
} }
// TestGitPushTool_EmptyBranch verifies that empty branch name is rejected. // TestGitPushTool_EmptyBranch verifies that empty branch name is rejected.
func TestGitPushTool_EmptyBranch(t *testing.T) { func TestGitPushTool_EmptyBranch(t *testing.T) {
tool := NewGitPushTool() tool := NewGitPushTool()
ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{
Branch: "", Branch: "",
BaseBranch: "main", BaseBranch: "main",
Path: t.TempDir(), Path: t.TempDir(),
RepoRoot: t.TempDir(), RepoRoot: t.TempDir(),
}) })
result := tool.Execute(ctx, map[string]any{}) result := tool.Execute(ctx, map[string]any{})
if !result.IsError { if !result.IsError {
t.Fatal("expected error for empty branch") t.Fatal("expected error for empty branch")
} }
assertContains(t, result.ForLLM, "no branch name") assertContains(t, result.ForLLM, "no branch name")
} }
// TestGitPushTool_AllowedBranch verifies that non-protected branches pass the branch check. // 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.) // (Push itself will fail because there's no real git repo, but it should get past validation.)
func TestGitPushTool_AllowedBranch(t *testing.T) { func TestGitPushTool_AllowedBranch(t *testing.T) {
tool := NewGitPushTool() tool := NewGitPushTool()
@ -76,12 +99,18 @@ func TestGitPushTool_AllowedBranch(t *testing.T) {
t.Run(branch, func(t *testing.T) { t.Run(branch, func(t *testing.T) {
ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{ ctx := WithWorktreeInfo(context.Background(), &git.WorktreeInfo{
Branch: branch, Branch: branch,
BaseBranch: "main", BaseBranch: "main",
Path: t.TempDir(), Path: t.TempDir(),
RepoRoot: t.TempDir(), RepoRoot: t.TempDir(),
}) })
result := tool.Execute(ctx, map[string]any{}) result := tool.Execute(ctx, map[string]any{})
// Should NOT fail with "protected branch" error // Should NOT fail with "protected branch" error
if result.IsError && strings.Contains(result.ForLLM, "protected") { if result.IsError && strings.Contains(result.ForLLM, "protected") {
t.Fatalf("branch %q should not be blocked as protected", branch) 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. // TestProtectedBranchesRegex tests the regex directly.
func TestProtectedBranchesRegex(t *testing.T) { func TestProtectedBranchesRegex(t *testing.T) {
tests := []struct { tests := []struct {
branch string branch string
protected bool protected bool
}{ }{
{"main", true}, {"main", true},
{"master", true}, {"master", true},
{"develop", true}, {"develop", true},
{"release/v1.0", true}, {"release/v1.0", true},
{"release/2026-03", true}, {"release/2026-03", true},
{"plan/add-feature", false}, {"plan/add-feature", false},
{"feature/main", false}, // "main" not at start {"feature/main", false}, // "main" not at start
{"main-backup", false}, // "main" followed by suffix {"main-backup", false}, // "main" followed by suffix
{"hotfix/urgent", false}, {"hotfix/urgent", false},
} }
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.branch, func(t *testing.T) { t.Run(tt.branch, func(t *testing.T) {
got := protectedBranches.MatchString(tt.branch) got := protectedBranches.MatchString(tt.branch)
if got != tt.protected { if got != tt.protected {
t.Errorf("branch %q: got protected=%v, want %v", tt.branch, 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. // TestWorktreeInfoContext verifies context round-trip.
func TestWorktreeInfoContext(t *testing.T) { func TestWorktreeInfoContext(t *testing.T) {
wt := &git.WorktreeInfo{ wt := &git.WorktreeInfo{
Branch: "plan/test", Branch: "plan/test",
BaseBranch: "main", BaseBranch: "main",
Path: "/tmp/wt", Path: "/tmp/wt",
RepoRoot: "/tmp/repo", RepoRoot: "/tmp/repo",
} }
ctx := WithWorktreeInfo(context.Background(), wt) ctx := WithWorktreeInfo(context.Background(), wt)
got := WorktreeInfoFromCtx(ctx) got := WorktreeInfoFromCtx(ctx)
if got == nil { if got == nil {
t.Fatal("expected non-nil WorktreeInfo from context") t.Fatal("expected non-nil WorktreeInfo from context")
} }
if got.Branch != wt.Branch { if got.Branch != wt.Branch {
t.Errorf("Branch: got %q, want %q", got.Branch, wt.Branch) t.Errorf("Branch: got %q, want %q", got.Branch, wt.Branch)
} }
if got.BaseBranch != wt.BaseBranch { if got.BaseBranch != wt.BaseBranch {
t.Errorf("BaseBranch: got %q, want %q", got.BaseBranch, wt.BaseBranch) t.Errorf("BaseBranch: got %q, want %q", got.BaseBranch, wt.BaseBranch)
} }
// Nil case // Nil case
got2 := WorktreeInfoFromCtx(context.Background()) got2 := WorktreeInfoFromCtx(context.Background())
if got2 != nil { if got2 != nil {
t.Errorf("expected nil WorktreeInfo from bare context, got %+v", got2) t.Errorf("expected nil WorktreeInfo from bare context, got %+v", got2)
} }
} }
// TestGitPushTool_Interface verifies the tool satisfies the Tool interface. // TestGitPushTool_Interface verifies the tool satisfies the Tool interface.
func TestGitPushTool_Interface(t *testing.T) { func TestGitPushTool_Interface(t *testing.T) {
var _ Tool = (*GitPushTool)(nil) var _ Tool = (*GitPushTool)(nil)
tool := NewGitPushTool() tool := NewGitPushTool()
if tool.Name() != "git_push" { if tool.Name() != "git_push" {
t.Errorf("Name: got %q, want %q", tool.Name(), "git_push") t.Errorf("Name: got %q, want %q", tool.Name(), "git_push")
} }
if tool.Description() == "" { if tool.Description() == "" {
t.Error("Description should not be empty") t.Error("Description should not be empty")
} }
params := tool.Parameters() params := tool.Parameters()
if params == nil { if params == nil {
t.Fatal("Parameters should not be nil") t.Fatal("Parameters should not be nil")
} }
if params["type"] != "object" { if params["type"] != "object" {
t.Errorf("Parameters type: got %v, want object", params["type"]) 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) { func assertContains(t *testing.T, s, substr string) {
t.Helper() t.Helper()
if !strings.Contains(s, substr) { if !strings.Contains(s, substr) {
t.Errorf("expected %q to contain %q", 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. // I2CTool provides I2C bus interaction for reading sensors and controlling peripherals.
type I2CTool struct{} type I2CTool struct{}
func NewI2CTool() *I2CTool { func NewI2CTool() *I2CTool {
@ -27,38 +28,55 @@ func (t *I2CTool) Description() string {
func (t *I2CTool) Parameters() map[string]any { func (t *I2CTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"action": map[string]any{ "action": map[string]any{
"type": "string", "type": "string",
"enum": []string{"detect", "scan", "read", "write"}, "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)", "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{ "bus": map[string]any{
"type": "string", "type": "string",
"description": "I2C bus number (e.g. \"1\" for /dev/i2c-1). Required for scan/read/write.", "description": "I2C bus number (e.g. \"1\" for /dev/i2c-1). Required for scan/read/write.",
}, },
"address": map[string]any{ "address": map[string]any{
"type": "integer", "type": "integer",
"description": "7-bit I2C device address (0x03-0x77). Required for read/write.", "description": "7-bit I2C device address (0x03-0x77). Required for read/write.",
}, },
"register": map[string]any{ "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.", "description": "Register address to read from or write to. If set, sends register byte before read/write.",
}, },
"data": map[string]any{ "data": map[string]any{
"type": "array", "type": "array",
"items": map[string]any{"type": "integer"}, "items": map[string]any{"type": "integer"},
"description": "Bytes to write (0-255 each). Required for write action.", "description": "Bytes to write (0-255 each). Required for write action.",
}, },
"length": map[string]any{ "length": map[string]any{
"type": "integer", "type": "integer",
"description": "Number of bytes to read (1-256). Default: 1. Used with read action.", "description": "Number of bytes to read (1-256). Default: 1. Used with read action.",
}, },
"confirm": map[string]any{ "confirm": map[string]any{
"type": "boolean", "type": "boolean",
"description": "Must be true for write operations. Safety guard to prevent accidental writes.", "description": "Must be true for write operations. Safety guard to prevent accidental writes.",
}, },
}, },
"required": []string{"action"}, "required": []string{"action"},
} }
} }
@ -69,25 +87,36 @@ func (t *I2CTool) Execute(ctx context.Context, args map[string]any) *ToolResult
} }
action, ok := args["action"].(string) action, ok := args["action"].(string)
if !ok { if !ok {
return ErrorResult("action is required") return ErrorResult("action is required")
} }
switch action { switch action {
case "detect": case "detect":
return t.detect() return t.detect()
case "scan": case "scan":
return t.scan(args) return t.scan(args)
case "read": case "read":
return t.readDevice(args) return t.readDevice(args)
case "write": case "write":
return t.writeDevice(args) return t.writeDevice(args)
default: default:
return ErrorResult(fmt.Sprintf("unknown action: %s (valid: detect, scan, read, write)", action)) return ErrorResult(fmt.Sprintf("unknown action: %s (valid: detect, scan, read, write)", action))
} }
} }
// detect lists available I2C buses by globbing /dev/i2c-* // detect lists available I2C buses by globbing /dev/i2c-*
func (t *I2CTool) detect() *ToolResult { func (t *I2CTool) detect() *ToolResult {
matches, err := filepath.Glob("/dev/i2c-*") matches, err := filepath.Glob("/dev/i2c-*")
if err != nil { if err != nil {
@ -102,11 +131,14 @@ func (t *I2CTool) detect() *ToolResult {
type busInfo struct { type busInfo struct {
Path string `json:"path"` Path string `json:"path"`
Bus string `json:"bus"` Bus string `json:"bus"`
} }
buses := make([]busInfo, 0, len(matches)) buses := make([]busInfo, 0, len(matches))
re := regexp.MustCompile(`/dev/i2c-(\d+)`) re := regexp.MustCompile(`/dev/i2c-(\d+)`)
for _, m := range matches { for _, m := range matches {
if sub := re.FindStringSubmatch(m); sub != nil { if sub := re.FindStringSubmatch(m); sub != nil {
buses = append(buses, busInfo{Path: m, Bus: sub[1]}) buses = append(buses, busInfo{Path: m, Bus: sub[1]})
@ -114,44 +146,62 @@ func (t *I2CTool) detect() *ToolResult {
} }
result, _ := json.MarshalIndent(buses, "", " ") result, _ := json.MarshalIndent(buses, "", " ")
return SilentResult(fmt.Sprintf("Found %d I2C bus(es):\n%s", len(buses), string(result))) 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) // Helper functions for I2C operations (used by platform-specific implementations)
// isValidBusID checks that a bus identifier is a simple number (prevents path injection) // isValidBusID checks that a bus identifier is a simple number (prevents path injection)
// //
//nolint:unused // Used by i2c_linux.go //nolint:unused // Used by i2c_linux.go
func isValidBusID(id string) bool { func isValidBusID(id string) bool {
matched, _ := regexp.MatchString(`^\d+$`, id) matched, _ := regexp.MatchString(`^\d+$`, id)
return matched return matched
} }
// parseI2CAddress extracts and validates an I2C address from args // parseI2CAddress extracts and validates an I2C address from args
// //
//nolint:unused // Used by i2c_linux.go //nolint:unused // Used by i2c_linux.go
func parseI2CAddress(args map[string]any) (int, *ToolResult) { func parseI2CAddress(args map[string]any) (int, *ToolResult) {
addrFloat, ok := args["address"].(float64) addrFloat, ok := args["address"].(float64)
if !ok { if !ok {
return 0, ErrorResult("address is required (e.g. 0x38 for AHT20)") return 0, ErrorResult("address is required (e.g. 0x38 for AHT20)")
} }
addr := int(addrFloat) addr := int(addrFloat)
if addr < 0x03 || addr > 0x77 { if addr < 0x03 || addr > 0x77 {
return 0, ErrorResult("address must be in valid 7-bit range (0x03-0x77)") return 0, ErrorResult("address must be in valid 7-bit range (0x03-0x77)")
} }
return addr, nil return addr, nil
} }
// parseI2CBus extracts and validates an I2C bus from args // parseI2CBus extracts and validates an I2C bus from args
// //
//nolint:unused // Used by i2c_linux.go //nolint:unused // Used by i2c_linux.go
func parseI2CBus(args map[string]any) (string, *ToolResult) { func parseI2CBus(args map[string]any) (string, *ToolResult) {
bus, ok := args["bus"].(string) bus, ok := args["bus"].(string)
if !ok || bus == "" { if !ok || bus == "" {
return "", ErrorResult("bus is required (e.g. \"1\" for /dev/i2c-1)") return "", ErrorResult("bus is required (e.g. \"1\" for /dev/i2c-1)")
} }
if !isValidBusID(bus) { if !isValidBusID(bus) {
return "", ErrorResult("invalid bus identifier: must be a number (e.g. \"1\")") return "", ErrorResult("invalid bus identifier: must be a number (e.g. \"1\")")
} }
return bus, nil return bus, nil
} }

View file

@ -8,279 +8,465 @@ import (
) )
// I2C ioctl constants from Linux kernel headers (<linux/i2c-dev.h>, <linux/i2c.h>) // I2C ioctl constants from Linux kernel headers (<linux/i2c-dev.h>, <linux/i2c.h>)
const ( const (
i2cSlave = 0x0703 // Set slave address (fails if in use by driver) i2cSlave = 0x0703 // Set slave address (fails if in use by driver)
i2cFuncs = 0x0705 // Query adapter functionality bitmask i2cFuncs = 0x0705 // Query adapter functionality bitmask
i2cSmbus = 0x0720 // Perform SMBus transaction i2cSmbus = 0x0720 // Perform SMBus transaction
// I2C_FUNC capability bits // I2C_FUNC capability bits
i2cFuncSmbusQuick = 0x00010000 i2cFuncSmbusQuick = 0x00010000
i2cFuncSmbusReadByte = 0x00020000 i2cFuncSmbusReadByte = 0x00020000
// SMBus transaction types // SMBus transaction types
i2cSmbusRead = 0 i2cSmbusRead = 0
i2cSmbusWrite = 1 i2cSmbusWrite = 1
// SMBus protocol sizes // SMBus protocol sizes
i2cSmbusQuick = 0 i2cSmbusQuick = 0
i2cSmbusByte = 1 i2cSmbusByte = 1
) )
// i2cSmbusData matches the kernel union i2c_smbus_data (34 bytes max). // 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). // For quick and byte transactions only the first byte is used (if at all).
type i2cSmbusData [34]byte type i2cSmbusData [34]byte
// i2cSmbusArgs matches the kernel struct i2c_smbus_ioctl_data. // i2cSmbusArgs matches the kernel struct i2c_smbus_ioctl_data.
type i2cSmbusArgs struct { type i2cSmbusArgs struct {
readWrite uint8 readWrite uint8
command uint8 command uint8
size uint32 size uint32
data *i2cSmbusData data *i2cSmbusData
} }
// smbusProbe performs a single SMBus probe at the given address. // smbusProbe performs a single SMBus probe at the given address.
// Uses SMBus Quick Write (safest) or falls back to SMBus Read Byte for // Uses SMBus Quick Write (safest) or falls back to SMBus Read Byte for
// EEPROM address ranges where quick write can corrupt AT24RF08 chips. // EEPROM address ranges where quick write can corrupt AT24RF08 chips.
// This matches i2cdetect's MODE_AUTO behavior. // This matches i2cdetect's MODE_AUTO behavior.
func smbusProbe(fd int, addr int, hasQuick bool) bool { func smbusProbe(fd int, addr int, hasQuick bool) bool {
// EEPROM ranges: use read byte (quick write can corrupt AT24RF08) // EEPROM ranges: use read byte (quick write can corrupt AT24RF08)
useReadByte := (addr >= 0x30 && addr <= 0x37) || (addr >= 0x50 && addr <= 0x5F) useReadByte := (addr >= 0x30 && addr <= 0x37) || (addr >= 0x50 && addr <= 0x5F)
if !useReadByte && hasQuick { if !useReadByte && hasQuick {
// SMBus Quick Write: [START] [ADDR|W] [ACK/NACK] [STOP] // SMBus Quick Write: [START] [ADDR|W] [ACK/NACK] [STOP]
// Safest probe — no data transferred // Safest probe — no data transferred
args := i2cSmbusArgs{ args := i2cSmbusArgs{
readWrite: i2cSmbusWrite, readWrite: i2cSmbusWrite,
command: 0, command: 0,
size: i2cSmbusQuick, size: i2cSmbusQuick,
data: nil, data: nil,
} }
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSmbus, uintptr(unsafe.Pointer(&args))) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSmbus, uintptr(unsafe.Pointer(&args)))
return errno == 0 return errno == 0
} }
// SMBus Read Byte: [START] [ADDR|R] [ACK/NACK] [DATA] [STOP] // SMBus Read Byte: [START] [ADDR|R] [ACK/NACK] [DATA] [STOP]
var data i2cSmbusData var data i2cSmbusData
args := i2cSmbusArgs{ args := i2cSmbusArgs{
readWrite: i2cSmbusRead, readWrite: i2cSmbusRead,
command: 0, command: 0,
size: i2cSmbusByte, size: i2cSmbusByte,
data: &data, data: &data,
} }
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSmbus, uintptr(unsafe.Pointer(&args))) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSmbus, uintptr(unsafe.Pointer(&args)))
return errno == 0 return errno == 0
} }
// scan probes valid 7-bit addresses on a bus for connected devices. // scan probes valid 7-bit addresses on a bus for connected devices.
// Uses the same hybrid probe strategy as i2cdetect's MODE_AUTO: // Uses the same hybrid probe strategy as i2cdetect's MODE_AUTO:
// SMBus Quick Write for most addresses, SMBus Read Byte for EEPROM ranges. // SMBus Quick Write for most addresses, SMBus Read Byte for EEPROM ranges.
func (t *I2CTool) scan(args map[string]any) *ToolResult { func (t *I2CTool) scan(args map[string]any) *ToolResult {
bus, errResult := parseI2CBus(args) bus, errResult := parseI2CBus(args)
if errResult != nil { if errResult != nil {
return errResult return errResult
} }
devPath := fmt.Sprintf("/dev/i2c-%s", bus) devPath := fmt.Sprintf("/dev/i2c-%s", bus)
fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) fd, err := syscall.Open(devPath, syscall.O_RDWR, 0)
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("failed to open %s: %v (check permissions and i2c-dev module)", devPath, err)) return ErrorResult(fmt.Sprintf("failed to open %s: %v (check permissions and i2c-dev module)", devPath, err))
} }
defer syscall.Close(fd) defer syscall.Close(fd)
// Query adapter capabilities to determine available probe methods. // Query adapter capabilities to determine available probe methods.
// I2C_FUNCS writes an unsigned long, which is word-sized on Linux. // I2C_FUNCS writes an unsigned long, which is word-sized on Linux.
var funcs uintptr var funcs uintptr
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cFuncs, uintptr(unsafe.Pointer(&funcs))) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cFuncs, uintptr(unsafe.Pointer(&funcs)))
if errno != 0 { if errno != 0 {
return ErrorResult(fmt.Sprintf("failed to query I2C adapter capabilities on %s: %v", devPath, errno)) return ErrorResult(fmt.Sprintf("failed to query I2C adapter capabilities on %s: %v", devPath, errno))
} }
hasQuick := funcs&i2cFuncSmbusQuick != 0 hasQuick := funcs&i2cFuncSmbusQuick != 0
hasReadByte := funcs&i2cFuncSmbusReadByte != 0 hasReadByte := funcs&i2cFuncSmbusReadByte != 0
if !hasQuick && !hasReadByte { if !hasQuick && !hasReadByte {
return ErrorResult( return ErrorResult(
fmt.Sprintf("I2C adapter %s supports neither SMBus Quick nor Read Byte — cannot probe safely", devPath), fmt.Sprintf("I2C adapter %s supports neither SMBus Quick nor Read Byte — cannot probe safely", devPath),
) )
} }
type deviceEntry struct { type deviceEntry struct {
Address string `json:"address"` Address string `json:"address"`
Status string `json:"status,omitempty"` Status string `json:"status,omitempty"`
} }
var found []deviceEntry var found []deviceEntry
// Scan 0x08-0x77, skipping I2C reserved addresses 0x00-0x07 // Scan 0x08-0x77, skipping I2C reserved addresses 0x00-0x07
for addr := 0x08; addr <= 0x77; addr++ { for addr := 0x08; addr <= 0x77; addr++ {
// Set slave address — EBUSY means a kernel driver owns this address // Set slave address — EBUSY means a kernel driver owns this address
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr)) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr))
if errno != 0 { if errno != 0 {
if errno == syscall.EBUSY { if errno == syscall.EBUSY {
found = append(found, deviceEntry{ found = append(found, deviceEntry{
Address: fmt.Sprintf("0x%02x", addr), Address: fmt.Sprintf("0x%02x", addr),
Status: "busy (in use by kernel driver)", Status: "busy (in use by kernel driver)",
}) })
} }
continue continue
} }
if smbusProbe(fd, addr, hasQuick) { if smbusProbe(fd, addr, hasQuick) {
found = append(found, deviceEntry{ found = append(found, deviceEntry{
Address: fmt.Sprintf("0x%02x", addr), Address: fmt.Sprintf("0x%02x", addr),
}) })
} }
} }
if len(found) == 0 { if len(found) == 0 {
return SilentResult(fmt.Sprintf("No devices found on %s. Check wiring and pull-up resistors.", devPath)) return SilentResult(fmt.Sprintf("No devices found on %s. Check wiring and pull-up resistors.", devPath))
} }
result, _ := json.MarshalIndent(map[string]any{ result, _ := json.MarshalIndent(map[string]any{
"bus": devPath, "bus": devPath,
"devices": found, "devices": found,
"count": len(found), "count": len(found),
}, "", " ") }, "", " ")
return SilentResult(fmt.Sprintf("Scan of %s:\n%s", devPath, string(result))) return SilentResult(fmt.Sprintf("Scan of %s:\n%s", devPath, string(result)))
} }
// readDevice reads bytes from an I2C device, optionally at a specific register // readDevice reads bytes from an I2C device, optionally at a specific register
func (t *I2CTool) readDevice(args map[string]any) *ToolResult { func (t *I2CTool) readDevice(args map[string]any) *ToolResult {
bus, errResult := parseI2CBus(args) bus, errResult := parseI2CBus(args)
if errResult != nil { if errResult != nil {
return errResult return errResult
} }
addr, errResult := parseI2CAddress(args) addr, errResult := parseI2CAddress(args)
if errResult != nil { if errResult != nil {
return errResult return errResult
} }
length := 1 length := 1
if l, ok := args["length"].(float64); ok { if l, ok := args["length"].(float64); ok {
length = int(l) length = int(l)
} }
if length < 1 || length > 256 { if length < 1 || length > 256 {
return ErrorResult("length must be between 1 and 256") return ErrorResult("length must be between 1 and 256")
} }
devPath := fmt.Sprintf("/dev/i2c-%s", bus) devPath := fmt.Sprintf("/dev/i2c-%s", bus)
fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) fd, err := syscall.Open(devPath, syscall.O_RDWR, 0)
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("failed to open %s: %v", devPath, err)) return ErrorResult(fmt.Sprintf("failed to open %s: %v", devPath, err))
} }
defer syscall.Close(fd) defer syscall.Close(fd)
// Set slave address // Set slave address
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr)) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr))
if errno != 0 { if errno != 0 {
return ErrorResult(fmt.Sprintf("failed to set I2C address 0x%02x: %v", addr, errno)) return ErrorResult(fmt.Sprintf("failed to set I2C address 0x%02x: %v", addr, errno))
} }
// If register is specified, write it first // If register is specified, write it first
if regFloat, ok := args["register"].(float64); ok { if regFloat, ok := args["register"].(float64); ok {
reg := int(regFloat) reg := int(regFloat)
if reg < 0 || reg > 255 { if reg < 0 || reg > 255 {
return ErrorResult("register must be between 0x00 and 0xFF") return ErrorResult("register must be between 0x00 and 0xFF")
} }
_, err = syscall.Write(fd, []byte{byte(reg)}) _, err = syscall.Write(fd, []byte{byte(reg)})
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("failed to write register 0x%02x: %v", reg, err)) return ErrorResult(fmt.Sprintf("failed to write register 0x%02x: %v", reg, err))
} }
} }
// Read data // Read data
buf := make([]byte, length) buf := make([]byte, length)
n, err := syscall.Read(fd, buf) n, err := syscall.Read(fd, buf)
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("failed to read from device 0x%02x: %v", addr, err)) return ErrorResult(fmt.Sprintf("failed to read from device 0x%02x: %v", addr, err))
} }
// Format as hex bytes // Format as hex bytes
hexBytes := make([]string, n) hexBytes := make([]string, n)
intBytes := make([]int, n) intBytes := make([]int, n)
for i := 0; i < n; i++ { for i := 0; i < n; i++ {
hexBytes[i] = fmt.Sprintf("0x%02x", buf[i]) hexBytes[i] = fmt.Sprintf("0x%02x", buf[i])
intBytes[i] = int(buf[i]) intBytes[i] = int(buf[i])
} }
result, _ := json.MarshalIndent(map[string]any{ result, _ := json.MarshalIndent(map[string]any{
"bus": devPath, "bus": devPath,
"address": fmt.Sprintf("0x%02x", addr), "address": fmt.Sprintf("0x%02x", addr),
"bytes": intBytes, "bytes": intBytes,
"hex": hexBytes, "hex": hexBytes,
"length": n, "length": n,
}, "", " ") }, "", " ")
return SilentResult(string(result)) return SilentResult(string(result))
} }
// writeDevice writes bytes to an I2C device, optionally at a specific register // writeDevice writes bytes to an I2C device, optionally at a specific register
func (t *I2CTool) writeDevice(args map[string]any) *ToolResult { func (t *I2CTool) writeDevice(args map[string]any) *ToolResult {
confirm, _ := args["confirm"].(bool) confirm, _ := args["confirm"].(bool)
if !confirm { if !confirm {
return ErrorResult( return ErrorResult(
"write operations require confirm: true. Please confirm with the user before writing to I2C devices, as incorrect writes can misconfigure hardware.", "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) bus, errResult := parseI2CBus(args)
if errResult != nil { if errResult != nil {
return errResult return errResult
} }
addr, errResult := parseI2CAddress(args) addr, errResult := parseI2CAddress(args)
if errResult != nil { if errResult != nil {
return errResult return errResult
} }
dataRaw, ok := args["data"].([]any) dataRaw, ok := args["data"].([]any)
if !ok || len(dataRaw) == 0 { if !ok || len(dataRaw) == 0 {
return ErrorResult("data is required for write (array of byte values 0-255)") return ErrorResult("data is required for write (array of byte values 0-255)")
} }
if len(dataRaw) > 256 { if len(dataRaw) > 256 {
return ErrorResult("data too long: maximum 256 bytes per I2C transaction") return ErrorResult("data too long: maximum 256 bytes per I2C transaction")
} }
data := make([]byte, 0, len(dataRaw)+1) data := make([]byte, 0, len(dataRaw)+1)
// If register is specified, prepend it to the data // If register is specified, prepend it to the data
if regFloat, ok := args["register"].(float64); ok { if regFloat, ok := args["register"].(float64); ok {
reg := int(regFloat) reg := int(regFloat)
if reg < 0 || reg > 255 { if reg < 0 || reg > 255 {
return ErrorResult("register must be between 0x00 and 0xFF") return ErrorResult("register must be between 0x00 and 0xFF")
} }
data = append(data, byte(reg)) data = append(data, byte(reg))
} }
for i, v := range dataRaw { for i, v := range dataRaw {
f, ok := v.(float64) f, ok := v.(float64)
if !ok { if !ok {
return ErrorResult(fmt.Sprintf("data[%d] is not a valid byte value", i)) return ErrorResult(fmt.Sprintf("data[%d] is not a valid byte value", i))
} }
b := int(f) b := int(f)
if b < 0 || b > 255 { if b < 0 || b > 255 {
return ErrorResult(fmt.Sprintf("data[%d] = %d is out of byte range (0-255)", i, b)) return ErrorResult(fmt.Sprintf("data[%d] = %d is out of byte range (0-255)", i, b))
} }
data = append(data, byte(b)) data = append(data, byte(b))
} }
devPath := fmt.Sprintf("/dev/i2c-%s", bus) devPath := fmt.Sprintf("/dev/i2c-%s", bus)
fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) fd, err := syscall.Open(devPath, syscall.O_RDWR, 0)
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("failed to open %s: %v", devPath, err)) return ErrorResult(fmt.Sprintf("failed to open %s: %v", devPath, err))
} }
defer syscall.Close(fd) defer syscall.Close(fd)
// Set slave address // Set slave address
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr)) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), i2cSlave, uintptr(addr))
if errno != 0 { if errno != 0 {
return ErrorResult(fmt.Sprintf("failed to set I2C address 0x%02x: %v", addr, errno)) return ErrorResult(fmt.Sprintf("failed to set I2C address 0x%02x: %v", addr, errno))
} }
// Write data // Write data
n, err := syscall.Write(fd, data) n, err := syscall.Write(fd, data)
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("failed to write to device 0x%02x: %v", addr, err)) 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)) 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. // 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. // Designed for token-efficient log analysis: defaults to WARN level to exclude noise.
type LogsTool struct{} type LogsTool struct{}
func NewLogsTool() *LogsTool { func NewLogsTool() *LogsTool {
@ -20,29 +22,40 @@ func (t *LogsTool) Name() string { return "logs" }
func (t *LogsTool) Description() string { func (t *LogsTool) Description() string {
return "Retrieve recent application logs from the in-memory ring buffer. " + return "Retrieve recent application logs from the in-memory ring buffer. " +
"Use level filter to minimize token usage (default: WARN). " + "Use level filter to minimize token usage (default: WARN). " +
"Call this when the user asks about errors, issues, or system health." "Call this when the user asks about errors, issues, or system health."
} }
func (t *LogsTool) Parameters() map[string]any { func (t *LogsTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"level": map[string]any{ "level": map[string]any{
"type": "string", "type": "string",
"description": "Minimum log level: DEBUG, INFO, WARN, ERROR. Default: WARN", "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{ "component": map[string]any{
"type": "string", "type": "string",
"description": "Filter by component name (e.g. telegram, discord, slack, agent)", "description": "Filter by component name (e.g. telegram, discord, slack, agent)",
}, },
"limit": map[string]any{ "limit": map[string]any{
"type": "integer", "type": "integer",
"description": "Maximum number of log entries to return. Default: 50", "description": "Maximum number of log entries to return. Default: 50",
}, },
"query": map[string]any{ "query": map[string]any{
"type": "string", "type": "string",
"description": "Filter by substring match in log message", "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 { func (t *LogsTool) Execute(_ context.Context, args map[string]any) *ToolResult {
// Parse level (default: WARN) // Parse level (default: WARN)
level := logger.WARN level := logger.WARN
if lvlStr, ok := args["level"].(string); ok && lvlStr != "" { if lvlStr, ok := args["level"].(string); ok && lvlStr != "" {
level = logger.ParseLevel(lvlStr) level = logger.ParseLevel(lvlStr)
} }
// Parse component // Parse component
component, _ := args["component"].(string) component, _ := args["component"].(string)
// Parse limit (default: 50, max: 300) // Parse limit (default: 50, max: 300)
limit := 50 limit := 50
if l, ok := args["limit"].(float64); ok && l > 0 { if l, ok := args["limit"].(float64); ok && l > 0 {
limit = int(l) limit = int(l)
} }
if limit > 300 { if limit > 300 {
limit = 300 limit = 300
} }
// Parse query // Parse query
query, _ := args["query"].(string) query, _ := args["query"].(string)
// Fetch from ring buffer (already sanitized by RecentLogs) // Fetch from ring buffer (already sanitized by RecentLogs)
entries := logger.RecentLogs(level, component, limit) entries := logger.RecentLogs(level, component, limit)
// Apply query filter if specified // Apply query filter if specified
if query != "" { if query != "" {
filtered := make([]logger.LogEntry, 0, len(entries)) filtered := make([]logger.LogEntry, 0, len(entries))
queryLower := strings.ToLower(query) queryLower := strings.ToLower(query)
for _, e := range entries { for _, e := range entries {
if strings.Contains(strings.ToLower(e.Message), queryLower) { if strings.Contains(strings.ToLower(e.Message), queryLower) {
filtered = append(filtered, e) filtered = append(filtered, e)
} }
} }
entries = filtered entries = filtered
} }

View file

@ -11,22 +11,31 @@ import (
func setupTestLogs(t *testing.T) { func setupTestLogs(t *testing.T) {
t.Helper() t.Helper()
prev := logger.GetLevel() prev := logger.GetLevel()
t.Cleanup(func() { logger.SetLevel(prev) }) t.Cleanup(func() { logger.SetLevel(prev) })
logger.SetLevel(logger.DEBUG) logger.SetLevel(logger.DEBUG)
logger.DebugC("agent", "debug message") logger.DebugC("agent", "debug message")
logger.InfoC("telegram", "message received") logger.InfoC("telegram", "message received")
logger.WarnC("telegram", "webhook retry") logger.WarnC("telegram", "webhook retry")
logger.ErrorC("discord", "connection timeout") logger.ErrorC("discord", "connection timeout")
logger.WarnCF("wecom", "signature failed", map[string]any{ logger.WarnCF("wecom", "signature failed", map[string]any{
"token": "secret-value", "token": "secret-value",
"nonce": "safe-value", "nonce": "safe-value",
}) })
} }
func TestLogsTool_DefaultLevel(t *testing.T) { func TestLogsTool_DefaultLevel(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{}) result := tool.Execute(context.Background(), map[string]any{})
@ -36,6 +45,7 @@ func TestLogsTool_DefaultLevel(t *testing.T) {
} }
var entries []logger.LogEntry var entries []logger.LogEntry
if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil {
t.Fatalf("failed to parse result: %v", err) t.Fatalf("failed to parse result: %v", err)
} }
@ -49,6 +59,7 @@ func TestLogsTool_DefaultLevel(t *testing.T) {
func TestLogsTool_LevelFilter(t *testing.T) { func TestLogsTool_LevelFilter(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
@ -60,6 +71,7 @@ func TestLogsTool_LevelFilter(t *testing.T) {
} }
var entries []logger.LogEntry var entries []logger.LogEntry
if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil {
t.Fatalf("failed to parse result: %v", err) t.Fatalf("failed to parse result: %v", err)
} }
@ -73,10 +85,12 @@ func TestLogsTool_LevelFilter(t *testing.T) {
func TestLogsTool_ComponentFilter(t *testing.T) { func TestLogsTool_ComponentFilter(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"level": "DEBUG", "level": "DEBUG",
"component": "telegram", "component": "telegram",
}) })
@ -85,6 +99,7 @@ func TestLogsTool_ComponentFilter(t *testing.T) {
} }
var entries []logger.LogEntry var entries []logger.LogEntry
if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil {
t.Fatalf("failed to parse result: %v", err) t.Fatalf("failed to parse result: %v", err)
} }
@ -98,10 +113,12 @@ func TestLogsTool_ComponentFilter(t *testing.T) {
func TestLogsTool_QueryFilter(t *testing.T) { func TestLogsTool_QueryFilter(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"level": "DEBUG", "level": "DEBUG",
"query": "timeout", "query": "timeout",
}) })
@ -110,6 +127,7 @@ func TestLogsTool_QueryFilter(t *testing.T) {
} }
var entries []logger.LogEntry var entries []logger.LogEntry
if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil {
t.Fatalf("failed to parse result: %v", err) t.Fatalf("failed to parse result: %v", err)
} }
@ -117,6 +135,7 @@ func TestLogsTool_QueryFilter(t *testing.T) {
if len(entries) == 0 { if len(entries) == 0 {
t.Fatal("expected at least one entry matching 'timeout'") t.Fatal("expected at least one entry matching 'timeout'")
} }
for _, e := range entries { for _, e := range entries {
if !strings.Contains(strings.ToLower(e.Message), "timeout") { if !strings.Contains(strings.ToLower(e.Message), "timeout") {
t.Errorf("entry should contain 'timeout': %s", e.Message) 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) { func TestLogsTool_QueryCaseInsensitive(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"level": "DEBUG", "level": "DEBUG",
"query": "TIMEOUT", "query": "TIMEOUT",
}) })
var entries []logger.LogEntry var entries []logger.LogEntry
if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil {
t.Fatalf("failed to parse result: %v", err) t.Fatalf("failed to parse result: %v", err)
} }
@ -145,10 +167,12 @@ func TestLogsTool_QueryCaseInsensitive(t *testing.T) {
func TestLogsTool_Limit(t *testing.T) { func TestLogsTool_Limit(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"level": "DEBUG", "level": "DEBUG",
"limit": float64(2), "limit": float64(2),
}) })
@ -157,6 +181,7 @@ func TestLogsTool_Limit(t *testing.T) {
} }
var entries []logger.LogEntry var entries []logger.LogEntry
if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil {
t.Fatalf("failed to parse result: %v", err) t.Fatalf("failed to parse result: %v", err)
} }
@ -170,12 +195,15 @@ func TestLogsTool_LimitMax(t *testing.T) {
tool := NewLogsTool() tool := NewLogsTool()
// limit > 300 should be capped // limit > 300 should be capped
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"level": "DEBUG", "level": "DEBUG",
"limit": float64(999), "limit": float64(999),
}) })
// Should not error, just cap silently // Should not error, just cap silently
if result.IsError { if result.IsError {
t.Fatalf("unexpected error: %s", result.ForLLM) t.Fatalf("unexpected error: %s", result.ForLLM)
} }
@ -183,10 +211,12 @@ func TestLogsTool_LimitMax(t *testing.T) {
func TestLogsTool_FieldsSanitized(t *testing.T) { func TestLogsTool_FieldsSanitized(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"level": "WARN", "level": "WARN",
"component": "wecom", "component": "wecom",
}) })
@ -195,22 +225,27 @@ func TestLogsTool_FieldsSanitized(t *testing.T) {
} }
var entries []logger.LogEntry var entries []logger.LogEntry
if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil { if err := json.Unmarshal([]byte(result.ForLLM), &entries); err != nil {
t.Fatalf("failed to parse result: %v", err) t.Fatalf("failed to parse result: %v", err)
} }
found := false found := false
for _, e := range entries { for _, e := range entries {
if e.Fields != nil && e.Fields["token"] != nil { if e.Fields != nil && e.Fields["token"] != nil {
found = true found = true
if e.Fields["token"] != "***" { if e.Fields["token"] != "***" {
t.Errorf("token field should be sanitized, got %v", e.Fields["token"]) t.Errorf("token field should be sanitized, got %v", e.Fields["token"])
} }
if e.Fields["nonce"] != "safe-value" { if e.Fields["nonce"] != "safe-value" {
t.Errorf("nonce field should be preserved, got %v", e.Fields["nonce"]) t.Errorf("nonce field should be preserved, got %v", e.Fields["nonce"])
} }
} }
} }
if !found { if !found {
t.Error("expected to find wecom entry with token field") 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) { func TestLogsTool_NoResults(t *testing.T) {
prev := logger.GetLevel() prev := logger.GetLevel()
defer logger.SetLevel(prev) defer logger.SetLevel(prev)
logger.SetLevel(logger.DEBUG) logger.SetLevel(logger.DEBUG)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"level": "ERROR", "level": "ERROR",
"component": "nonexistent-component-xyz", "component": "nonexistent-component-xyz",
}) })
if result.IsError { if result.IsError {
t.Fatalf("should not be an error result: %s", result.ForLLM) t.Fatalf("should not be an error result: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "No log entries found") { if !strings.Contains(result.ForLLM, "No log entries found") {
t.Errorf("expected 'No log entries found' message, got: %s", result.ForLLM) 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) { func TestLogsTool_Silent(t *testing.T) {
setupTestLogs(t) setupTestLogs(t)
tool := NewLogsTool() tool := NewLogsTool()
result := tool.Execute(context.Background(), map[string]any{}) result := tool.Execute(context.Background(), map[string]any{})
if !result.Silent { if !result.Silent {
t.Error("logs tool result should be Silent") t.Error("logs tool result should be Silent")
} }
@ -252,10 +293,13 @@ func TestLogsTool_ToolInterface(t *testing.T) {
if tool.Name() != "logs" { if tool.Name() != "logs" {
t.Errorf("expected name 'logs', got %q", tool.Name()) t.Errorf("expected name 'logs', got %q", tool.Name())
} }
if tool.Description() == "" { if tool.Description() == "" {
t.Error("description should not be empty") t.Error("description should not be empty")
} }
params := tool.Parameters() params := tool.Parameters()
if params == nil { if params == nil {
t.Error("parameters should not be nil") t.Error("parameters should not be nil")
} }

View file

@ -9,8 +9,11 @@ type SendCallback func(channel, chatID, content string) error
type MessageTool struct { type MessageTool struct {
sendCallback SendCallback sendCallback SendCallback
defaultChannel string defaultChannel string
defaultChatID string defaultChatID string
sentInRound bool // Tracks whether a message was sent in the current processing round sentInRound bool // Tracks whether a message was sent in the current processing round
} }
@ -29,31 +32,41 @@ func (t *MessageTool) Description() string {
func (t *MessageTool) Parameters() map[string]any { func (t *MessageTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"content": map[string]any{ "content": map[string]any{
"type": "string", "type": "string",
"description": "The message content to send", "description": "The message content to send",
}, },
"channel": map[string]any{ "channel": map[string]any{
"type": "string", "type": "string",
"description": "Optional: target channel (telegram, whatsapp, etc.)", "description": "Optional: target channel (telegram, whatsapp, etc.)",
}, },
"chat_id": map[string]any{ "chat_id": map[string]any{
"type": "string", "type": "string",
"description": "Optional: target chat/user ID", "description": "Optional: target chat/user ID",
}, },
}, },
"required": []string{"content"}, "required": []string{"content"},
} }
} }
func (t *MessageTool) SetContext(channel, chatID string) { func (t *MessageTool) SetContext(channel, chatID string) {
t.defaultChannel = channel t.defaultChannel = channel
t.defaultChatID = chatID t.defaultChatID = chatID
t.sentInRound = false // Reset send tracking for new processing round t.sentInRound = false // Reset send tracking for new processing round
} }
// HasSentInRound returns true if the message tool sent a message during the current round. // HasSentInRound returns true if the message tool sent a message during the current round.
func (t *MessageTool) HasSentInRound() bool { func (t *MessageTool) HasSentInRound() bool {
return t.sentInRound 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 { func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
content, ok := args["content"].(string) content, ok := args["content"].(string)
if !ok { if !ok {
return &ToolResult{ForLLM: "content is required", IsError: true} return &ToolResult{ForLLM: "content is required", IsError: true}
} }
channel, _ := args["channel"].(string) channel, _ := args["channel"].(string)
chatID, _ := args["chat_id"].(string) chatID, _ := args["chat_id"].(string)
if channel == "" { if channel == "" {
channel = t.defaultChannel channel = t.defaultChannel
} }
if chatID == "" { if chatID == "" {
chatID = t.defaultChatID chatID = t.defaultChatID
} }
@ -89,15 +105,20 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolRes
if err := t.sendCallback(channel, chatID, content); err != nil { if err := t.sendCallback(channel, chatID, content); err != nil {
return &ToolResult{ return &ToolResult{
ForLLM: fmt.Sprintf("sending message: %v", err), ForLLM: fmt.Sprintf("sending message: %v", err),
IsError: true, IsError: true,
Err: err, Err: err,
} }
} }
t.sentInRound = true t.sentInRound = true
// Silent: user already received the message directly // Silent: user already received the message directly
return &ToolResult{ return &ToolResult{
ForLLM: fmt.Sprintf("Message sent to %s:%s", channel, chatID), ForLLM: fmt.Sprintf("Message sent to %s:%s", channel, chatID),
Silent: true, Silent: true,
} }
} }

View file

@ -8,17 +8,23 @@ import (
func TestMessageTool_Execute_Success(t *testing.T) { func TestMessageTool_Execute_Success(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
tool.SetContext("test-channel", "test-chat-id") tool.SetContext("test-channel", "test-chat-id")
var sentChannel, sentChatID, sentContent string var sentChannel, sentChatID, sentContent string
tool.SetSendCallback(func(channel, chatID, content string) error { tool.SetSendCallback(func(channel, chatID, content string) error {
sentChannel = channel sentChannel = channel
sentChatID = chatID sentChatID = chatID
sentContent = content sentContent = content
return nil return nil
}) })
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"content": "Hello, world!", "content": "Hello, world!",
} }
@ -26,33 +32,41 @@ func TestMessageTool_Execute_Success(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Verify message was sent with correct parameters // Verify message was sent with correct parameters
if sentChannel != "test-channel" { if sentChannel != "test-channel" {
t.Errorf("Expected channel 'test-channel', got '%s'", sentChannel) t.Errorf("Expected channel 'test-channel', got '%s'", sentChannel)
} }
if sentChatID != "test-chat-id" { if sentChatID != "test-chat-id" {
t.Errorf("Expected chatID 'test-chat-id', got '%s'", sentChatID) t.Errorf("Expected chatID 'test-chat-id', got '%s'", sentChatID)
} }
if sentContent != "Hello, world!" { if sentContent != "Hello, world!" {
t.Errorf("Expected content 'Hello, world!', got '%s'", sentContent) t.Errorf("Expected content 'Hello, world!', got '%s'", sentContent)
} }
// Verify ToolResult meets US-011 criteria: // Verify ToolResult meets US-011 criteria:
// - Send success returns SilentResult (Silent=true) // - Send success returns SilentResult (Silent=true)
if !result.Silent { if !result.Silent {
t.Error("Expected Silent=true for successful send") t.Error("Expected Silent=true for successful send")
} }
// - ForLLM contains send status description // - ForLLM contains send status description
if result.ForLLM != "Message sent to test-channel:test-chat-id" { 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) t.Errorf("Expected ForLLM 'Message sent to test-channel:test-chat-id', got '%s'", result.ForLLM)
} }
// - ForUser is empty (user already received message directly) // - ForUser is empty (user already received message directly)
if result.ForUser != "" { if result.ForUser != "" {
t.Errorf("Expected ForUser to be empty, got '%s'", result.ForUser) t.Errorf("Expected ForUser to be empty, got '%s'", result.ForUser)
} }
// - IsError should be false // - IsError should be false
if result.IsError { if result.IsError {
t.Error("Expected IsError=false for successful send") 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) { func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
tool.SetContext("default-channel", "default-chat-id") tool.SetContext("default-channel", "default-chat-id")
var sentChannel, sentChatID string var sentChannel, sentChatID string
tool.SetSendCallback(func(channel, chatID, content string) error { tool.SetSendCallback(func(channel, chatID, content string) error {
sentChannel = channel sentChannel = channel
sentChatID = chatID sentChatID = chatID
return nil return nil
}) })
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"content": "Test message", "content": "Test message",
"channel": "custom-channel", "channel": "custom-channel",
"chat_id": "custom-chat-id", "chat_id": "custom-chat-id",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Verify custom channel/chatID were used instead of defaults // Verify custom channel/chatID were used instead of defaults
if sentChannel != "custom-channel" { if sentChannel != "custom-channel" {
t.Errorf("Expected channel 'custom-channel', got '%s'", sentChannel) t.Errorf("Expected channel 'custom-channel', got '%s'", sentChannel)
} }
if sentChatID != "custom-chat-id" { if sentChatID != "custom-chat-id" {
t.Errorf("Expected chatID 'custom-chat-id', got '%s'", sentChatID) t.Errorf("Expected chatID 'custom-chat-id', got '%s'", sentChatID)
} }
@ -89,6 +112,7 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {
if !result.Silent { if !result.Silent {
t.Error("Expected Silent=true") t.Error("Expected Silent=true")
} }
if result.ForLLM != "Message sent to custom-channel:custom-chat-id" { 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) 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) { func TestMessageTool_Execute_SendFailure(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
tool.SetContext("test-channel", "test-chat-id") tool.SetContext("test-channel", "test-chat-id")
sendErr := errors.New("network error") sendErr := errors.New("network error")
tool.SetSendCallback(func(channel, chatID, content string) error { tool.SetSendCallback(func(channel, chatID, content string) error {
return sendErr return sendErr
}) })
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"content": "Test message", "content": "Test message",
} }
@ -111,21 +138,27 @@ func TestMessageTool_Execute_SendFailure(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Verify ToolResult for send failure: // Verify ToolResult for send failure:
// - Send failure returns ErrorResult (IsError=true) // - Send failure returns ErrorResult (IsError=true)
if !result.IsError { if !result.IsError {
t.Error("Expected IsError=true for failed send") t.Error("Expected IsError=true for failed send")
} }
// - ForLLM contains error description // - ForLLM contains error description
expectedErrMsg := "sending message: network error" expectedErrMsg := "sending message: network error"
if result.ForLLM != expectedErrMsg { if result.ForLLM != expectedErrMsg {
t.Errorf("Expected ForLLM '%s', got '%s'", expectedErrMsg, result.ForLLM) t.Errorf("Expected ForLLM '%s', got '%s'", expectedErrMsg, result.ForLLM)
} }
// - Err field should contain original error // - Err field should contain original error
if result.Err == nil { if result.Err == nil {
t.Error("Expected Err to be set") t.Error("Expected Err to be set")
} }
if result.Err != sendErr { if result.Err != sendErr {
t.Errorf("Expected Err to be sendErr, got %v", result.Err) 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) { func TestMessageTool_Execute_MissingContent(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
tool.SetContext("test-channel", "test-chat-id") tool.SetContext("test-channel", "test-chat-id")
ctx := context.Background() ctx := context.Background()
args := map[string]any{} // content missing args := map[string]any{} // content missing
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Verify error result for missing content // Verify error result for missing content
if !result.IsError { if !result.IsError {
t.Error("Expected IsError=true for missing content") t.Error("Expected IsError=true for missing content")
} }
if result.ForLLM != "content is required" { if result.ForLLM != "content is required" {
t.Errorf("Expected ForLLM 'content is required', got '%s'", result.ForLLM) 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) { func TestMessageTool_Execute_NoTargetChannel(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
// No SetContext called, so defaultChannel and defaultChatID are empty // No SetContext called, so defaultChannel and defaultChatID are empty
tool.SetSendCallback(func(channel, chatID, content string) error { tool.SetSendCallback(func(channel, chatID, content string) error {
@ -158,6 +196,7 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) {
}) })
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"content": "Test message", "content": "Test message",
} }
@ -165,9 +204,11 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Verify error when no target channel specified // Verify error when no target channel specified
if !result.IsError { if !result.IsError {
t.Error("Expected IsError=true when no target channel") t.Error("Expected IsError=true when no target channel")
} }
if result.ForLLM != "No target channel/chat specified" { if result.ForLLM != "No target channel/chat specified" {
t.Errorf("Expected ForLLM 'No target channel/chat specified', got '%s'", result.ForLLM) 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) { func TestMessageTool_Execute_NotConfigured(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
tool.SetContext("test-channel", "test-chat-id") tool.SetContext("test-channel", "test-chat-id")
// No SetSendCallback called // No SetSendCallback called
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"content": "Test message", "content": "Test message",
} }
@ -186,9 +230,11 @@ func TestMessageTool_Execute_NotConfigured(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Verify error when send callback not configured // Verify error when send callback not configured
if !result.IsError { if !result.IsError {
t.Error("Expected IsError=true when send callback not configured") t.Error("Expected IsError=true when send callback not configured")
} }
if result.ForLLM != "Message sending not configured" { if result.ForLLM != "Message sending not configured" {
t.Errorf("Expected ForLLM 'Message sending not configured', got '%s'", result.ForLLM) 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) { func TestMessageTool_Name(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
if tool.Name() != "message" { if tool.Name() != "message" {
t.Errorf("Expected name 'message', got '%s'", tool.Name()) 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) { func TestMessageTool_Description(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
desc := tool.Description() desc := tool.Description()
if desc == "" { if desc == "" {
t.Error("Description should not be empty") t.Error("Description should not be empty")
} }
@ -211,48 +260,63 @@ func TestMessageTool_Description(t *testing.T) {
func TestMessageTool_Parameters(t *testing.T) { func TestMessageTool_Parameters(t *testing.T) {
tool := NewMessageTool() tool := NewMessageTool()
params := tool.Parameters() params := tool.Parameters()
// Verify parameters structure // Verify parameters structure
typ, ok := params["type"].(string) typ, ok := params["type"].(string)
if !ok || typ != "object" { if !ok || typ != "object" {
t.Error("Expected type 'object'") t.Error("Expected type 'object'")
} }
props, ok := params["properties"].(map[string]any) props, ok := params["properties"].(map[string]any)
if !ok { if !ok {
t.Fatal("Expected properties to be a map") t.Fatal("Expected properties to be a map")
} }
// Check required properties // Check required properties
required, ok := params["required"].([]string) required, ok := params["required"].([]string)
if !ok || len(required) != 1 || required[0] != "content" { if !ok || len(required) != 1 || required[0] != "content" {
t.Error("Expected 'content' to be required") t.Error("Expected 'content' to be required")
} }
// Check content property // Check content property
contentProp, ok := props["content"].(map[string]any) contentProp, ok := props["content"].(map[string]any)
if !ok { if !ok {
t.Error("Expected 'content' property") t.Error("Expected 'content' property")
} }
if contentProp["type"] != "string" { if contentProp["type"] != "string" {
t.Error("Expected content type to be 'string'") t.Error("Expected content type to be 'string'")
} }
// Check channel property (optional) // Check channel property (optional)
channelProp, ok := props["channel"].(map[string]any) channelProp, ok := props["channel"].(map[string]any)
if !ok { if !ok {
t.Error("Expected 'channel' property") t.Error("Expected 'channel' property")
} }
if channelProp["type"] != "string" { if channelProp["type"] != "string" {
t.Error("Expected channel type to be 'string'") t.Error("Expected channel type to be 'string'")
} }
// Check chat_id property (optional) // Check chat_id property (optional)
chatIDProp, ok := props["chat_id"].(map[string]any) chatIDProp, ok := props["chat_id"].(map[string]any)
if !ok { if !ok {
t.Error("Expected 'chat_id' property") t.Error("Expected 'chat_id' property")
} }
if chatIDProp["type"] != "string" { if chatIDProp["type"] != "string" {
t.Error("Expected chat_id type to be 'string'") t.Error("Expected chat_id type to be 'string'")
} }

View file

@ -14,9 +14,12 @@ import (
) )
// NormalizeToolName keeps only lowercase ASCII letters. // NormalizeToolName keeps only lowercase ASCII letters.
// "read_file" → "readfile", "ReadFile" → "readfile", "read-file" → "readfile". // "read_file" → "readfile", "ReadFile" → "readfile", "read-file" → "readfile".
func NormalizeToolName(s string) string { func NormalizeToolName(s string) string {
var b strings.Builder var b strings.Builder
for _, r := range s { for _, r := range s {
if r >= 'A' && r <= 'Z' { if r >= 'A' && r <= 'Z' {
b.WriteRune(r + 32) b.WriteRune(r + 32)
@ -24,11 +27,13 @@ func NormalizeToolName(s string) string {
b.WriteRune(r) b.WriteRune(r)
} }
} }
return b.String() return b.String()
} }
type ToolRegistry struct { type ToolRegistry struct {
tools map[string]Tool tools map[string]Tool
mu sync.RWMutex mu sync.RWMutex
} }
@ -40,24 +45,33 @@ func NewToolRegistry() *ToolRegistry {
func (r *ToolRegistry) Register(tool Tool) { func (r *ToolRegistry) Register(tool Tool) {
r.mu.Lock() r.mu.Lock()
defer r.mu.Unlock() defer r.mu.Unlock()
r.tools[tool.Name()] = tool r.tools[tool.Name()] = tool
} }
func (r *ToolRegistry) Get(name string) (Tool, bool) { func (r *ToolRegistry) Get(name string) (Tool, bool) {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
// Exact match first // Exact match first
if tool, ok := r.tools[name]; ok { if tool, ok := r.tools[name]; ok {
return tool, true return tool, true
} }
// Fuzzy fallback: normalize and compare (handles "readfile" → "read_file" etc.) // Fuzzy fallback: normalize and compare (handles "readfile" → "read_file" etc.)
norm := NormalizeToolName(name) norm := NormalizeToolName(name)
for _, tool := range r.tools { for _, tool := range r.tools {
if NormalizeToolName(tool.Name()) == norm { if NormalizeToolName(tool.Name()) == norm {
return tool, true return tool, true
} }
} }
return nil, false 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. // ExecuteWithContext executes a tool with channel/chatID context and optional async callback.
// If the tool implements AsyncTool and a non-nil callback is provided, // If the tool implements AsyncTool and a non-nil callback is provided,
// the callback will be set on the tool before execution. // the callback will be set on the tool before execution.
func (r *ToolRegistry) ExecuteWithContext( func (r *ToolRegistry) ExecuteWithContext(
ctx context.Context, ctx context.Context,
name string, name string,
args map[string]any, args map[string]any,
channel, chatID string, channel, chatID string,
asyncCallback AsyncCallback, asyncCallback AsyncCallback,
) *ToolResult { ) *ToolResult {
logger.InfoCF("tool", "Tool execution started", logger.InfoCF("tool", "Tool execution started",
map[string]any{ map[string]any{
"tool": name, "tool": name,
"args": args, "args": args,
}) })
tool, ok := r.Get(name) tool, ok := r.Get(name)
if !ok { if !ok {
available := strings.Join(r.List(), ", ") available := strings.Join(r.List(), ", ")
logger.ErrorCF("tool", "Tool not found", logger.ErrorCF("tool", "Tool not found",
map[string]any{ map[string]any{
"tool": name, "tool": name,
}) })
return ErrorResult(fmt.Sprintf( return ErrorResult(fmt.Sprintf(
"tool %q not found. Available tools: %s", name, available, "tool %q not found. Available tools: %s", name, available,
)).WithError(fmt.Errorf("tool not found")) )).WithError(fmt.Errorf("tool not found"))
} }
// If tool implements ContextualTool, set context // If tool implements ContextualTool, set context
if contextualTool, ok := tool.(ContextualTool); ok && channel != "" && chatID != "" { if contextualTool, ok := tool.(ContextualTool); ok && channel != "" && chatID != "" {
contextualTool.SetContext(channel, chatID) contextualTool.SetContext(channel, chatID)
} }
// If tool implements AsyncTool and callback is provided, set callback // If tool implements AsyncTool and callback is provided, set callback
if asyncTool, ok := tool.(AsyncTool); ok && asyncCallback != nil { if asyncTool, ok := tool.(AsyncTool); ok && asyncCallback != nil {
asyncTool.SetCallback(asyncCallback) asyncTool.SetCallback(asyncCallback)
logger.DebugCF("tool", "Async callback injected", logger.DebugCF("tool", "Async callback injected",
map[string]any{ map[string]any{
"tool": name, "tool": name,
}) })
} }
start := time.Now() start := time.Now()
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
duration := time.Since(start) duration := time.Since(start)
// Log based on result type // Log based on result type
if result.IsError { if result.IsError {
logger.ErrorCF("tool", "Tool execution failed", logger.ErrorCF("tool", "Tool execution failed",
map[string]any{ map[string]any{
"tool": name, "tool": name,
"duration": duration.Milliseconds(), "duration": duration.Milliseconds(),
"error": result.ForLLM, "error": result.ForLLM,
}) })
} else if result.Async { } else if result.Async {
logger.InfoCF("tool", "Tool started (async)", logger.InfoCF("tool", "Tool started (async)",
map[string]any{ map[string]any{
"tool": name, "tool": name,
"duration": duration.Milliseconds(), "duration": duration.Milliseconds(),
}) })
} else { } else {
logger.InfoCF("tool", "Tool execution completed", logger.InfoCF("tool", "Tool execution completed",
map[string]any{ map[string]any{
"tool": name, "tool": name,
"duration_ms": duration.Milliseconds(), "duration_ms": duration.Milliseconds(),
"result_length": len(result.ForLLM), "result_length": len(result.ForLLM),
}) })
} }
@ -138,53 +181,75 @@ func (r *ToolRegistry) ExecuteWithContext(
} }
// sortedToolNames returns tool names in sorted order for deterministic iteration. // sortedToolNames returns tool names in sorted order for deterministic iteration.
// This is critical for KV cache stability: non-deterministic map iteration would // This is critical for KV cache stability: non-deterministic map iteration would
// produce different system prompts and tool definitions on each call, invalidating // produce different system prompts and tool definitions on each call, invalidating
// the LLM's prefix cache even when no tools have changed. // the LLM's prefix cache even when no tools have changed.
func (r *ToolRegistry) sortedToolNames() []string { func (r *ToolRegistry) sortedToolNames() []string {
names := make([]string, 0, len(r.tools)) names := make([]string, 0, len(r.tools))
for name := range r.tools { for name := range r.tools {
names = append(names, name) names = append(names, name)
} }
sort.Strings(names) sort.Strings(names)
return names return names
} }
func (r *ToolRegistry) GetDefinitions() []map[string]any { func (r *ToolRegistry) GetDefinitions() []map[string]any {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
sorted := r.sortedToolNames() sorted := r.sortedToolNames()
definitions := make([]map[string]any, 0, len(sorted)) definitions := make([]map[string]any, 0, len(sorted))
for _, name := range sorted { for _, name := range sorted {
definitions = append(definitions, ToolToSchema(r.tools[name])) definitions = append(definitions, ToolToSchema(r.tools[name]))
} }
return definitions return definitions
} }
// ToProviderDefs converts tool definitions to provider-compatible format. // ToProviderDefs converts tool definitions to provider-compatible format.
// This is the format expected by LLM provider APIs. // This is the format expected by LLM provider APIs.
func (r *ToolRegistry) ToProviderDefs() []providers.ToolDefinition { func (r *ToolRegistry) ToProviderDefs() []providers.ToolDefinition {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
sorted := r.sortedToolNames() sorted := r.sortedToolNames()
definitions := make([]providers.ToolDefinition, 0, len(sorted)) definitions := make([]providers.ToolDefinition, 0, len(sorted))
for _, name := range sorted { for _, name := range sorted {
tool := r.tools[name] tool := r.tools[name]
schema := ToolToSchema(tool) schema := ToolToSchema(tool)
// Safely extract nested values with type checks // Safely extract nested values with type checks
fn, ok := schema["function"].(map[string]any) fn, ok := schema["function"].(map[string]any)
if !ok { if !ok {
continue continue
} }
name, _ := fn["name"].(string) name, _ := fn["name"].(string)
desc, _ := fn["description"].(string) desc, _ := fn["description"].(string)
params, _ := fn["parameters"].(map[string]any) params, _ := fn["parameters"].(map[string]any)
paramsRaw := json.RawMessage(`{}`) paramsRaw := json.RawMessage(`{}`)
if len(params) > 0 { if len(params) > 0 {
if payload, err := json.Marshal(params); err == nil { if payload, err := json.Marshal(params); err == nil {
paramsRaw = json.RawMessage(payload) paramsRaw = json.RawMessage(payload)
@ -193,38 +258,51 @@ func (r *ToolRegistry) ToProviderDefs() []providers.ToolDefinition {
definitions = append(definitions, providers.ToolDefinition{ definitions = append(definitions, providers.ToolDefinition{
Type: "function", Type: "function",
Function: providers.ToolFunctionDefinition{ Function: providers.ToolFunctionDefinition{
Name: name, Name: name,
Description: desc, Description: desc,
Parameters: paramsRaw, Parameters: paramsRaw,
}, },
}) })
} }
return definitions return definitions
} }
// List returns a list of all registered tool names. // List returns a list of all registered tool names.
func (r *ToolRegistry) List() []string { func (r *ToolRegistry) List() []string {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
return r.sortedToolNames() return r.sortedToolNames()
} }
// Count returns the number of registered tools. // Count returns the number of registered tools.
func (r *ToolRegistry) Count() int { func (r *ToolRegistry) Count() int {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
return len(r.tools) return len(r.tools)
} }
// GetRuntimeStatus aggregates runtime status from all tools that implement StatusProvider. // GetRuntimeStatus aggregates runtime status from all tools that implement StatusProvider.
// Returns empty string if no tool has status to report. // Returns empty string if no tool has status to report.
func (r *ToolRegistry) GetRuntimeStatus() string { func (r *ToolRegistry) GetRuntimeStatus() string {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
var parts []string var parts []string
for _, tool := range r.tools { for _, tool := range r.tools {
if sp, ok := tool.(StatusProvider); ok { if sp, ok := tool.(StatusProvider); ok {
if s := sp.RuntimeStatus(); s != "" { if s := sp.RuntimeStatus(); s != "" {
@ -232,40 +310,53 @@ func (r *ToolRegistry) GetRuntimeStatus() string {
} }
} }
} }
if len(parts) == 0 { if len(parts) == 0 {
return "" return ""
} }
return strings.Join(parts, "\n\n") return strings.Join(parts, "\n\n")
} }
// buildParamHint extracts parameter names from a JSON schema and returns // buildParamHint extracts parameter names from a JSON schema and returns
// a hint string like "(task, label?, preset?)". Required params are bare, // a hint string like "(task, label?, preset?)". Required params are bare,
// optional params have a trailing "?". // optional params have a trailing "?".
func buildParamHint(schema map[string]any) string { func buildParamHint(schema map[string]any) string {
props, _ := schema["properties"].(map[string]any) props, _ := schema["properties"].(map[string]any)
if len(props) == 0 { if len(props) == 0 {
return "" return ""
} }
reqSlice, _ := schema["required"].([]string) reqSlice, _ := schema["required"].([]string)
reqSet := make(map[string]bool, len(reqSlice)) reqSet := make(map[string]bool, len(reqSlice))
for _, r := range reqSlice { for _, r := range reqSlice {
reqSet[r] = true reqSet[r] = true
} }
names := make([]string, 0, len(props)) names := make([]string, 0, len(props))
for name := range props { for name := range props {
names = append(names, name) names = append(names, name)
} }
sort.Strings(names) sort.Strings(names)
parts := make([]string, 0, len(names)) parts := make([]string, 0, len(names))
// Required params first, then optional // Required params first, then optional
for _, name := range names { for _, name := range names {
if reqSet[name] { if reqSet[name] {
parts = append(parts, name) parts = append(parts, name)
} }
} }
for _, name := range names { for _, name := range names {
if !reqSet[name] { if !reqSet[name] {
parts = append(parts, 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. // GetSummaries returns human-readable summaries of all registered tools.
// Returns a slice of "- `name`(params) - description" strings. // Returns a slice of "- `name`(params) - description" strings.
func (r *ToolRegistry) GetSummaries() []string { func (r *ToolRegistry) GetSummaries() []string {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
sorted := r.sortedToolNames() sorted := r.sortedToolNames()
summaries := make([]string, 0, len(sorted)) summaries := make([]string, 0, len(sorted))
for _, name := range sorted { for _, name := range sorted {
tool := r.tools[name] tool := r.tools[name]
hint := buildParamHint(tool.Parameters()) hint := buildParamHint(tool.Parameters())
summaries = append(summaries, fmt.Sprintf("- `%s`%s - %s", tool.Name(), hint, tool.Description())) summaries = append(summaries, fmt.Sprintf("- `%s`%s - %s", tool.Name(), hint, tool.Description()))
} }
return summaries return summaries
} }

View file

@ -13,31 +13,41 @@ import (
type mockRegistryTool struct { type mockRegistryTool struct {
name string name string
desc string desc string
params map[string]any params map[string]any
result *ToolResult result *ToolResult
} }
func (m *mockRegistryTool) Name() string { return m.name } func (m *mockRegistryTool) Name() string { return m.name }
func (m *mockRegistryTool) Description() string { return m.desc } func (m *mockRegistryTool) Description() string { return m.desc }
func (m *mockRegistryTool) Parameters() map[string]any { return m.params } func (m *mockRegistryTool) Parameters() map[string]any { return m.params }
func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]any) *ToolResult { func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]any) *ToolResult {
return m.result return m.result
} }
type mockCtxTool struct { type mockCtxTool struct {
mockRegistryTool mockRegistryTool
channel string channel string
chatID string chatID string
} }
func (m *mockCtxTool) SetContext(channel, chatID string) { func (m *mockCtxTool) SetContext(channel, chatID string) {
m.channel = channel m.channel = channel
m.chatID = chatID m.chatID = chatID
} }
type mockAsyncRegistryTool struct { type mockAsyncRegistryTool struct {
mockRegistryTool mockRegistryTool
cb AsyncCallback cb AsyncCallback
} }
@ -50,8 +60,11 @@ func (m *mockAsyncRegistryTool) SetCallback(cb AsyncCallback) {
func newMockTool(name, desc string) *mockRegistryTool { func newMockTool(name, desc string) *mockRegistryTool {
return &mockRegistryTool{ return &mockRegistryTool{
name: name, name: name,
desc: desc, desc: desc,
params: map[string]any{"type": "object"}, params: map[string]any{"type": "object"},
result: SilentResult("ok"), result: SilentResult("ok"),
} }
} }
@ -63,15 +76,23 @@ func TestNormalizeToolName(t *testing.T) {
input, want string input, want string
}{ }{
{"read_file", "readfile"}, {"read_file", "readfile"},
{"readfile", "readfile"}, {"readfile", "readfile"},
{"ReadFile", "readfile"}, {"ReadFile", "readfile"},
{"read-file", "readfile"}, {"read-file", "readfile"},
{"edit_file", "editfile"}, {"edit_file", "editfile"},
{"web_search", "websearch"}, {"web_search", "websearch"},
{"EXEC", "exec"}, {"EXEC", "exec"},
} }
for _, tt := range tests { for _, tt := range tests {
got := NormalizeToolName(tt.input) got := NormalizeToolName(tt.input)
if got != tt.want { if got != tt.want {
t.Errorf("NormalizeToolName(%q) = %q, want %q", tt.input, 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) { func TestNewToolRegistry(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
if r.Count() != 0 { if r.Count() != 0 {
t.Errorf("expected empty registry, got count %d", r.Count()) t.Errorf("expected empty registry, got count %d", r.Count())
} }
if len(r.List()) != 0 { if len(r.List()) != 0 {
t.Errorf("expected empty list, got %v", r.List()) t.Errorf("expected empty list, got %v", r.List())
} }
@ -90,13 +113,17 @@ func TestNewToolRegistry(t *testing.T) {
func TestToolRegistry_RegisterAndGet(t *testing.T) { func TestToolRegistry_RegisterAndGet(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
tool := newMockTool("echo", "echoes input") tool := newMockTool("echo", "echoes input")
r.Register(tool) r.Register(tool)
got, ok := r.Get("echo") got, ok := r.Get("echo")
if !ok { if !ok {
t.Fatal("expected to find registered tool") t.Fatal("expected to find registered tool")
} }
if got.Name() != "echo" { if got.Name() != "echo" {
t.Errorf("expected name 'echo', got %q", got.Name()) 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) { func TestToolRegistry_Get_NotFound(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
_, ok := r.Get("nonexistent") _, ok := r.Get("nonexistent")
if ok { if ok {
t.Error("expected ok=false for unregistered tool") 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) { func TestToolRegistry_Get_FuzzyMatch(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(newMockTool("read_file", "reads a file")) r.Register(newMockTool("read_file", "reads a file"))
r.Register(newMockTool("edit_file", "edits a file")) r.Register(newMockTool("edit_file", "edits a file"))
r.Register(newMockTool("web_search", "searches the web")) r.Register(newMockTool("web_search", "searches the web"))
tests := []struct { tests := []struct {
query string query string
wantName string wantName string
}{ }{
{"readfile", "read_file"}, {"readfile", "read_file"},
{"ReadFile", "read_file"}, {"ReadFile", "read_file"},
{"read-file", "read_file"}, {"read-file", "read_file"},
{"editfile", "edit_file"}, {"editfile", "edit_file"},
{"EditFile", "edit_file"}, {"EditFile", "edit_file"},
{"websearch", "web_search"}, {"websearch", "web_search"},
{"WebSearch", "web_search"}, {"WebSearch", "web_search"},
} }
for _, tt := range tests { for _, tt := range tests {
tool, ok := r.Get(tt.query) tool, ok := r.Get(tt.query)
if !ok { if !ok {
t.Errorf("Get(%q) not found, want %q", tt.query, tt.wantName) t.Errorf("Get(%q) not found, want %q", tt.query, tt.wantName)
continue continue
} }
if tool.Name() != tt.wantName { if tool.Name() != tt.wantName {
t.Errorf("Get(%q).Name() = %q, want %q", tt.query, 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) { func TestToolRegistry_RegisterOverwrite(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(newMockTool("dup", "first")) r.Register(newMockTool("dup", "first"))
r.Register(newMockTool("dup", "second")) r.Register(newMockTool("dup", "second"))
if r.Count() != 1 { if r.Count() != 1 {
t.Errorf("expected count 1 after overwrite, got %d", r.Count()) t.Errorf("expected count 1 after overwrite, got %d", r.Count())
} }
tool, _ := r.Get("dup") tool, _ := r.Get("dup")
if tool.Description() != "second" { if tool.Description() != "second" {
t.Errorf("expected overwritten description 'second', got %q", tool.Description()) 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) { func TestToolRegistry_Execute_Success(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(&mockRegistryTool{ r.Register(&mockRegistryTool{
name: "greet", name: "greet",
desc: "says hello", desc: "says hello",
params: map[string]any{}, params: map[string]any{},
result: SilentResult("hello"), result: SilentResult("hello"),
}) })
result := r.Execute(context.Background(), "greet", nil) result := r.Execute(context.Background(), "greet", nil)
if result.IsError { if result.IsError {
t.Errorf("expected success, got error: %s", result.ForLLM) t.Errorf("expected success, got error: %s", result.ForLLM)
} }
if result.ForLLM != "hello" { if result.ForLLM != "hello" {
t.Errorf("expected ForLLM 'hello', got %q", result.ForLLM) 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) { func TestToolRegistry_Execute_NotFound(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
result := r.Execute(context.Background(), "missing", nil) result := r.Execute(context.Background(), "missing", nil)
if !result.IsError { if !result.IsError {
t.Error("expected error for missing tool") t.Error("expected error for missing tool")
} }
if !strings.Contains(result.ForLLM, "not found") { if !strings.Contains(result.ForLLM, "not found") {
t.Errorf("expected 'not found' in error, got %q", result.ForLLM) t.Errorf("expected 'not found' in error, got %q", result.ForLLM)
} }
if result.Err == nil { if result.Err == nil {
t.Error("expected Err to be set via WithError") 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) { func TestToolRegistry_ExecuteWithContext_ContextualTool(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
ct := &mockCtxTool{ ct := &mockCtxTool{
mockRegistryTool: *newMockTool("ctx_tool", "needs context"), mockRegistryTool: *newMockTool("ctx_tool", "needs context"),
} }
r.Register(ct) r.Register(ct)
r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "telegram", "chat-42", nil) 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" { if ct.channel != "telegram" {
t.Errorf("expected channel 'telegram', got %q", ct.channel) t.Errorf("expected channel 'telegram', got %q", ct.channel)
} }
if ct.chatID != "chat-42" { if ct.chatID != "chat-42" {
t.Errorf("expected chatID 'chat-42', got %q", ct.chatID) 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) { func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
ct := &mockCtxTool{ ct := &mockCtxTool{
mockRegistryTool: *newMockTool("ctx_tool", "needs context"), mockRegistryTool: *newMockTool("ctx_tool", "needs context"),
} }
r.Register(ct) r.Register(ct)
r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "", "", nil) 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) { func TestToolRegistry_ExecuteWithContext_AsyncCallback(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
at := &mockAsyncRegistryTool{ at := &mockAsyncRegistryTool{
mockRegistryTool: *newMockTool("async_tool", "async work"), mockRegistryTool: *newMockTool("async_tool", "async work"),
} }
at.result = AsyncResult("started") at.result = AsyncResult("started")
r.Register(at) r.Register(at)
called := false called := false
cb := func(_ context.Context, _ *ToolResult) { called = true } cb := func(_ context.Context, _ *ToolResult) { called = true }
result := r.ExecuteWithContext(context.Background(), "async_tool", nil, "", "", cb) result := r.ExecuteWithContext(context.Background(), "async_tool", nil, "", "", cb)
if at.cb == nil { if at.cb == nil {
t.Error("expected SetCallback to have been called") t.Error("expected SetCallback to have been called")
} }
if !result.Async { if !result.Async {
t.Error("expected async result") t.Error("expected async result")
} }
at.cb(context.Background(), SilentResult("done")) at.cb(context.Background(), SilentResult("done"))
if !called { if !called {
t.Error("expected callback to be invoked") t.Error("expected callback to be invoked")
} }
@ -244,22 +313,29 @@ func TestToolRegistry_ExecuteWithContext_AsyncCallback(t *testing.T) {
func TestToolRegistry_GetDefinitions(t *testing.T) { func TestToolRegistry_GetDefinitions(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(newMockTool("alpha", "tool A")) r.Register(newMockTool("alpha", "tool A"))
defs := r.GetDefinitions() defs := r.GetDefinitions()
if len(defs) != 1 { if len(defs) != 1 {
t.Fatalf("expected 1 definition, got %d", len(defs)) t.Fatalf("expected 1 definition, got %d", len(defs))
} }
if defs[0]["type"] != "function" { if defs[0]["type"] != "function" {
t.Errorf("expected type 'function', got %v", defs[0]["type"]) t.Errorf("expected type 'function', got %v", defs[0]["type"])
} }
fn, ok := defs[0]["function"].(map[string]any) fn, ok := defs[0]["function"].(map[string]any)
if !ok { if !ok {
t.Fatal("expected 'function' key to be a map") t.Fatal("expected 'function' key to be a map")
} }
if fn["name"] != "alpha" { if fn["name"] != "alpha" {
t.Errorf("expected name 'alpha', got %v", fn["name"]) t.Errorf("expected name 'alpha', got %v", fn["name"])
} }
if fn["description"] != "tool A" { if fn["description"] != "tool A" {
t.Errorf("expected description 'tool A', got %v", fn["description"]) 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) { func TestToolRegistry_ToProviderDefs(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
params := map[string]any{"type": "object", "properties": map[string]any{}} params := map[string]any{"type": "object", "properties": map[string]any{}}
r.Register(&mockRegistryTool{ r.Register(&mockRegistryTool{
name: "beta", name: "beta",
desc: "tool B", desc: "tool B",
params: params, params: params,
result: SilentResult("ok"), result: SilentResult("ok"),
}) })
defs := r.ToProviderDefs() defs := r.ToProviderDefs()
if len(defs) != 1 { if len(defs) != 1 {
t.Fatalf("expected 1 provider def, got %d", len(defs)) t.Fatalf("expected 1 provider def, got %d", len(defs))
} }
want := providers.ToolDefinition{ want := providers.ToolDefinition{
Type: "function", Type: "function",
Function: providers.ToolFunctionDefinition{ Function: providers.ToolFunctionDefinition{
Name: "beta", Name: "beta",
Description: "tool B", Description: "tool B",
Parameters: providers.MustMarshalParameters(params), Parameters: providers.MustMarshalParameters(params),
}, },
} }
got := defs[0] got := defs[0]
if got.Type != want.Type { if got.Type != want.Type {
t.Errorf("Type: want %q, got %q", want.Type, got.Type) t.Errorf("Type: want %q, got %q", want.Type, got.Type)
} }
if got.Function.Name != want.Function.Name { if got.Function.Name != want.Function.Name {
t.Errorf("Name: want %q, got %q", want.Function.Name, got.Function.Name) t.Errorf("Name: want %q, got %q", want.Function.Name, got.Function.Name)
} }
if got.Function.Description != want.Function.Description { if got.Function.Description != want.Function.Description {
t.Errorf("Description: want %q, got %q", want.Function.Description, got.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) { func TestToolRegistry_List(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(newMockTool("x", "")) r.Register(newMockTool("x", ""))
r.Register(newMockTool("y", "")) r.Register(newMockTool("y", ""))
names := r.List() names := r.List()
if len(names) != 2 { if len(names) != 2 {
t.Fatalf("expected 2 names, got %d", len(names)) t.Fatalf("expected 2 names, got %d", len(names))
} }
nameSet := map[string]bool{} nameSet := map[string]bool{}
for _, n := range names { for _, n := range names {
nameSet[n] = true nameSet[n] = true
} }
if !nameSet["x"] || !nameSet["y"] { if !nameSet["x"] || !nameSet["y"] {
t.Errorf("expected names {x, y}, got %v", names) 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) { func TestToolRegistry_Count(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
if r.Count() != 0 { if r.Count() != 0 {
t.Errorf("expected 0, got %d", r.Count()) t.Errorf("expected 0, got %d", r.Count())
} }
r.Register(newMockTool("a", "")) r.Register(newMockTool("a", ""))
r.Register(newMockTool("b", "")) r.Register(newMockTool("b", ""))
if r.Count() != 2 { if r.Count() != 2 {
t.Errorf("expected 2, got %d", r.Count()) t.Errorf("expected 2, got %d", r.Count())
} }
r.Register(newMockTool("a", "replaced")) r.Register(newMockTool("a", "replaced"))
if r.Count() != 2 { if r.Count() != 2 {
t.Errorf("expected 2 after overwrite, got %d", r.Count()) t.Errorf("expected 2 after overwrite, got %d", r.Count())
} }
@ -340,61 +438,90 @@ func TestToolRegistry_Count(t *testing.T) {
func TestBuildParamHint(t *testing.T) { func TestBuildParamHint(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
schema map[string]any schema map[string]any
want string want string
}{ }{
{ {
name: "required and optional", name: "required and optional",
schema: map[string]any{ schema: map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"task": map[string]any{"type": "string"}, "task": map[string]any{"type": "string"},
"label": map[string]any{"type": "string"}, "label": map[string]any{"type": "string"},
}, },
"required": []string{"task"}, "required": []string{"task"},
}, },
want: "(task, label?)", want: "(task, label?)",
}, },
{ {
name: "all required", name: "all required",
schema: map[string]any{ schema: map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"command": map[string]any{"type": "string"}, "command": map[string]any{"type": "string"},
}, },
"required": []string{"command"}, "required": []string{"command"},
}, },
want: "(command)", want: "(command)",
}, },
{ {
name: "no properties", name: "no properties",
schema: map[string]any{ schema: map[string]any{
"type": "object", "type": "object",
}, },
want: "", want: "",
}, },
{ {
name: "empty schema", name: "empty schema",
schema: map[string]any{}, schema: map[string]any{},
want: "", want: "",
}, },
{ {
name: "nil schema", name: "nil schema",
schema: nil, schema: nil,
want: "", want: "",
}, },
{ {
name: "multiple optional sorted", name: "multiple optional sorted",
schema: map[string]any{ schema: map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"task": map[string]any{"type": "string"}, "task": map[string]any{"type": "string"},
"preset": map[string]any{"type": "string"}, "preset": map[string]any{"type": "string"},
"label": map[string]any{"type": "string"}, "label": map[string]any{"type": "string"},
"agent_id": map[string]any{"type": "string"}, "agent_id": map[string]any{"type": "string"},
}, },
"required": []string{"task"}, "required": []string{"task"},
}, },
want: "(task, agent_id?, label?, preset?)", want: "(task, agent_id?, label?, preset?)",
}, },
} }
@ -402,6 +529,7 @@ func TestBuildParamHint(t *testing.T) {
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
got := buildParamHint(tt.schema) got := buildParamHint(tt.schema)
if got != tt.want { if got != tt.want {
t.Errorf("buildParamHint() = %q, want %q", 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) { func TestToolRegistry_GetSummaries(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(newMockTool("read_file", "Reads a file")) r.Register(newMockTool("read_file", "Reads a file"))
summaries := r.GetSummaries() summaries := r.GetSummaries()
if len(summaries) != 1 { if len(summaries) != 1 {
t.Fatalf("expected 1 summary, got %d", len(summaries)) t.Fatalf("expected 1 summary, got %d", len(summaries))
} }
if !strings.Contains(summaries[0], "`read_file`") { if !strings.Contains(summaries[0], "`read_file`") {
t.Errorf("expected backtick-quoted name in summary, got %q", summaries[0]) t.Errorf("expected backtick-quoted name in summary, got %q", summaries[0])
} }
if !strings.Contains(summaries[0], "Reads a file") { if !strings.Contains(summaries[0], "Reads a file") {
t.Errorf("expected description in summary, got %q", summaries[0]) 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) { func TestToolRegistry_GetSummaries_WithParamHint(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(&mockRegistryTool{ r.Register(&mockRegistryTool{
name: "spawn", name: "spawn",
desc: "Spawn a subagent", desc: "Spawn a subagent",
params: map[string]any{ params: map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"task": map[string]any{"type": "string"}, "task": map[string]any{"type": "string"},
"preset": map[string]any{"type": "string"}, "preset": map[string]any{"type": "string"},
}, },
"required": []string{"task"}, "required": []string{"task"},
}, },
result: SilentResult("ok"), result: SilentResult("ok"),
}) })
summaries := r.GetSummaries() summaries := r.GetSummaries()
if len(summaries) != 1 { if len(summaries) != 1 {
t.Fatalf("expected 1 summary, got %d", len(summaries)) t.Fatalf("expected 1 summary, got %d", len(summaries))
} }
// Should contain param hint // Should contain param hint
if !strings.Contains(summaries[0], "(task, preset?)") { if !strings.Contains(summaries[0], "(task, preset?)") {
t.Errorf("expected param hint in summary, got %q", summaries[0]) 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) { func TestToolToSchema(t *testing.T) {
tool := newMockTool("demo", "demo tool") tool := newMockTool("demo", "demo tool")
schema := ToolToSchema(tool) schema := ToolToSchema(tool)
if schema["type"] != "function" { if schema["type"] != "function" {
t.Errorf("expected type 'function', got %v", schema["type"]) t.Errorf("expected type 'function', got %v", schema["type"])
} }
fn, ok := schema["function"].(map[string]any) fn, ok := schema["function"].(map[string]any)
if !ok { if !ok {
t.Fatal("expected 'function' to be a map") t.Fatal("expected 'function' to be a map")
} }
if fn["name"] != "demo" { if fn["name"] != "demo" {
t.Errorf("expected name 'demo', got %v", fn["name"]) t.Errorf("expected name 'demo', got %v", fn["name"])
} }
if fn["description"] != "demo tool" { if fn["description"] != "demo tool" {
t.Errorf("expected description 'demo tool', got %v", fn["description"]) t.Errorf("expected description 'demo tool', got %v", fn["description"])
} }
if fn["parameters"] == nil { if fn["parameters"] == nil {
t.Error("expected parameters to be set") t.Error("expected parameters to be set")
} }
@ -475,17 +623,25 @@ func TestToolToSchema(t *testing.T) {
func TestToolRegistry_ConcurrentAccess(t *testing.T) { func TestToolRegistry_ConcurrentAccess(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
var wg sync.WaitGroup var wg sync.WaitGroup
for i := range 50 { for i := range 50 {
wg.Add(1) wg.Add(1)
go func(n int) { go func(n int) {
defer wg.Done() defer wg.Done()
name := string(rune('A' + n%26)) name := string(rune('A' + n%26))
r.Register(newMockTool(name, "concurrent")) r.Register(newMockTool(name, "concurrent"))
r.Get(name) r.Get(name)
r.Count() r.Count()
r.List() r.List()
r.GetDefinitions() r.GetDefinitions()
}(i) }(i)
} }

View file

@ -12,12 +12,15 @@ func TestNewToolResult(t *testing.T) {
if result.ForLLM != "test content" { if result.ForLLM != "test content" {
t.Errorf("Expected ForLLM 'test content', got '%s'", result.ForLLM) t.Errorf("Expected ForLLM 'test content', got '%s'", result.ForLLM)
} }
if result.Silent { if result.Silent {
t.Error("Expected Silent to be false") t.Error("Expected Silent to be false")
} }
if result.IsError { if result.IsError {
t.Error("Expected IsError to be false") t.Error("Expected IsError to be false")
} }
if result.Async { if result.Async {
t.Error("Expected Async to be false") t.Error("Expected Async to be false")
} }
@ -29,12 +32,15 @@ func TestSilentResult(t *testing.T) {
if result.ForLLM != "silent operation" { if result.ForLLM != "silent operation" {
t.Errorf("Expected ForLLM 'silent operation', got '%s'", result.ForLLM) t.Errorf("Expected ForLLM 'silent operation', got '%s'", result.ForLLM)
} }
if !result.Silent { if !result.Silent {
t.Error("Expected Silent to be true") t.Error("Expected Silent to be true")
} }
if result.IsError { if result.IsError {
t.Error("Expected IsError to be false") t.Error("Expected IsError to be false")
} }
if result.Async { if result.Async {
t.Error("Expected Async to be false") t.Error("Expected Async to be false")
} }
@ -46,12 +52,15 @@ func TestAsyncResult(t *testing.T) {
if result.ForLLM != "async task started" { if result.ForLLM != "async task started" {
t.Errorf("Expected ForLLM 'async task started', got '%s'", result.ForLLM) t.Errorf("Expected ForLLM 'async task started', got '%s'", result.ForLLM)
} }
if result.Silent { if result.Silent {
t.Error("Expected Silent to be false") t.Error("Expected Silent to be false")
} }
if result.IsError { if result.IsError {
t.Error("Expected IsError to be false") t.Error("Expected IsError to be false")
} }
if !result.Async { if !result.Async {
t.Error("Expected Async to be true") t.Error("Expected Async to be true")
} }
@ -63,12 +72,15 @@ func TestErrorResult(t *testing.T) {
if result.ForLLM != "operation failed" { if result.ForLLM != "operation failed" {
t.Errorf("Expected ForLLM 'operation failed', got '%s'", result.ForLLM) t.Errorf("Expected ForLLM 'operation failed', got '%s'", result.ForLLM)
} }
if result.Silent { if result.Silent {
t.Error("Expected Silent to be false") t.Error("Expected Silent to be false")
} }
if !result.IsError { if !result.IsError {
t.Error("Expected IsError to be true") t.Error("Expected IsError to be true")
} }
if result.Async { if result.Async {
t.Error("Expected Async to be false") t.Error("Expected Async to be false")
} }
@ -76,20 +88,25 @@ func TestErrorResult(t *testing.T) {
func TestUserResult(t *testing.T) { func TestUserResult(t *testing.T) {
content := "user visible message" content := "user visible message"
result := UserResult(content) result := UserResult(content)
if result.ForLLM != content { if result.ForLLM != content {
t.Errorf("Expected ForLLM '%s', got '%s'", content, result.ForLLM) t.Errorf("Expected ForLLM '%s', got '%s'", content, result.ForLLM)
} }
if result.ForUser != content { if result.ForUser != content {
t.Errorf("Expected ForUser '%s', got '%s'", content, result.ForUser) t.Errorf("Expected ForUser '%s', got '%s'", content, result.ForUser)
} }
if result.Silent { if result.Silent {
t.Error("Expected Silent to be false") t.Error("Expected Silent to be false")
} }
if result.IsError { if result.IsError {
t.Error("Expected IsError to be false") t.Error("Expected IsError to be false")
} }
if result.Async { if result.Async {
t.Error("Expected Async to be false") t.Error("Expected Async to be false")
} }
@ -98,26 +115,36 @@ func TestUserResult(t *testing.T) {
func TestToolResultJSONSerialization(t *testing.T) { func TestToolResultJSONSerialization(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
result *ToolResult result *ToolResult
}{ }{
{ {
name: "basic result", name: "basic result",
result: NewToolResult("basic content"), result: NewToolResult("basic content"),
}, },
{ {
name: "silent result", name: "silent result",
result: SilentResult("silent content"), result: SilentResult("silent content"),
}, },
{ {
name: "async result", name: "async result",
result: AsyncResult("async content"), result: AsyncResult("async content"),
}, },
{ {
name: "error result", name: "error result",
result: ErrorResult("error content"), result: ErrorResult("error content"),
}, },
{ {
name: "user result", name: "user result",
result: UserResult("user content"), result: UserResult("user content"),
}, },
} }
@ -125,30 +152,38 @@ func TestToolResultJSONSerialization(t *testing.T) {
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
// Marshal to JSON // Marshal to JSON
data, err := json.Marshal(tt.result) data, err := json.Marshal(tt.result)
if err != nil { if err != nil {
t.Fatalf("Failed to marshal: %v", err) t.Fatalf("Failed to marshal: %v", err)
} }
// Unmarshal back // Unmarshal back
var decoded ToolResult var decoded ToolResult
if err := json.Unmarshal(data, &decoded); err != nil { if err := json.Unmarshal(data, &decoded); err != nil {
t.Fatalf("Failed to unmarshal: %v", err) t.Fatalf("Failed to unmarshal: %v", err)
} }
// Verify fields match (Err should be excluded) // Verify fields match (Err should be excluded)
if decoded.ForLLM != tt.result.ForLLM { if decoded.ForLLM != tt.result.ForLLM {
t.Errorf("ForLLM mismatch: got '%s', want '%s'", decoded.ForLLM, tt.result.ForLLM) t.Errorf("ForLLM mismatch: got '%s', want '%s'", decoded.ForLLM, tt.result.ForLLM)
} }
if decoded.ForUser != tt.result.ForUser { if decoded.ForUser != tt.result.ForUser {
t.Errorf("ForUser mismatch: got '%s', want '%s'", decoded.ForUser, tt.result.ForUser) t.Errorf("ForUser mismatch: got '%s', want '%s'", decoded.ForUser, tt.result.ForUser)
} }
if decoded.Silent != tt.result.Silent { if decoded.Silent != tt.result.Silent {
t.Errorf("Silent mismatch: got %v, want %v", decoded.Silent, tt.result.Silent) t.Errorf("Silent mismatch: got %v, want %v", decoded.Silent, tt.result.Silent)
} }
if decoded.IsError != tt.result.IsError { if decoded.IsError != tt.result.IsError {
t.Errorf("IsError mismatch: got %v, want %v", decoded.IsError, tt.result.IsError) t.Errorf("IsError mismatch: got %v, want %v", decoded.IsError, tt.result.IsError)
} }
if decoded.Async != tt.result.Async { if decoded.Async != tt.result.Async {
t.Errorf("Async mismatch: got %v, want %v", 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) { func TestToolResultWithErrors(t *testing.T) {
err := errors.New("underlying error") err := errors.New("underlying error")
result := ErrorResult("error message").WithError(err) result := ErrorResult("error message").WithError(err)
if result.Err == nil { if result.Err == nil {
t.Error("Expected Err to be set") t.Error("Expected Err to be set")
} }
if result.Err.Error() != "underlying error" { if result.Err.Error() != "underlying error" {
t.Errorf("Expected Err message 'underlying error', got '%s'", result.Err.Error()) t.Errorf("Expected Err message 'underlying error', got '%s'", result.Err.Error())
} }
// Verify Err is not serialized // Verify Err is not serialized
data, marshalErr := json.Marshal(result) data, marshalErr := json.Marshal(result)
if marshalErr != nil { if marshalErr != nil {
t.Fatalf("Failed to marshal: %v", marshalErr) t.Fatalf("Failed to marshal: %v", marshalErr)
} }
var decoded ToolResult var decoded ToolResult
if unmarshalErr := json.Unmarshal(data, &decoded); unmarshalErr != nil { if unmarshalErr := json.Unmarshal(data, &decoded); unmarshalErr != nil {
t.Fatalf("Failed to unmarshal: %v", unmarshalErr) t.Fatalf("Failed to unmarshal: %v", unmarshalErr)
} }
@ -192,37 +232,47 @@ func TestToolResultJSONStructure(t *testing.T) {
} }
// Verify JSON structure // Verify JSON structure
var parsed map[string]any var parsed map[string]any
if err := json.Unmarshal(data, &parsed); err != nil { if err := json.Unmarshal(data, &parsed); err != nil {
t.Fatalf("Failed to parse JSON: %v", err) t.Fatalf("Failed to parse JSON: %v", err)
} }
// Check expected keys exist // Check expected keys exist
if _, ok := parsed["for_llm"]; !ok { if _, ok := parsed["for_llm"]; !ok {
t.Error("Expected 'for_llm' key in JSON") t.Error("Expected 'for_llm' key in JSON")
} }
if _, ok := parsed["for_user"]; !ok { if _, ok := parsed["for_user"]; !ok {
t.Error("Expected 'for_user' key in JSON") t.Error("Expected 'for_user' key in JSON")
} }
if _, ok := parsed["silent"]; !ok { if _, ok := parsed["silent"]; !ok {
t.Error("Expected 'silent' key in JSON") t.Error("Expected 'silent' key in JSON")
} }
if _, ok := parsed["is_error"]; !ok { if _, ok := parsed["is_error"]; !ok {
t.Error("Expected 'is_error' key in JSON") t.Error("Expected 'is_error' key in JSON")
} }
if _, ok := parsed["async"]; !ok { if _, ok := parsed["async"]; !ok {
t.Error("Expected 'async' key in JSON") t.Error("Expected 'async' key in JSON")
} }
// Check that 'err' is NOT present (it should have json:"-" tag) // Check that 'err' is NOT present (it should have json:"-" tag)
if _, ok := parsed["err"]; ok { if _, ok := parsed["err"]; ok {
t.Error("Expected 'err' key to be excluded from JSON") t.Error("Expected 'err' key to be excluded from JSON")
} }
// Verify values // Verify values
if parsed["for_llm"] != "test content" { if parsed["for_llm"] != "test content" {
t.Errorf("Expected for_llm 'test content', got %v", parsed["for_llm"]) t.Errorf("Expected for_llm 'test content', got %v", parsed["for_llm"])
} }
if parsed["silent"] != false { if parsed["silent"] != false {
t.Errorf("Expected silent false, got %v", parsed["silent"]) t.Errorf("Expected silent false, got %v", parsed["silent"])
} }

File diff suppressed because it is too large Load diff

View file

@ -11,75 +11,129 @@ import (
) )
func prepareCommandForTermination(cmd *exec.Cmd) { func prepareCommandForTermination(cmd *exec.Cmd) {
if cmd == nil { if cmd == nil {
return return
} }
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
} }
func terminateProcessTree(cmd *exec.Cmd) error { func terminateProcessTree(cmd *exec.Cmd) error {
if cmd == nil || cmd.Process == nil { if cmd == nil || cmd.Process == nil {
return nil return nil
} }
pid := cmd.Process.Pid pid := cmd.Process.Pid
if pid <= 0 { if pid <= 0 {
return nil return nil
} }
// Kill the entire process group spawned by the shell command. // Kill the entire process group spawned by the shell command.
_ = syscall.Kill(-pid, syscall.SIGKILL) _ = syscall.Kill(-pid, syscall.SIGKILL)
// Some shells/background jobs may still leave descendants around // Some shells/background jobs may still leave descendants around
// briefly; aggressively walk /proc and kill child processes too. // briefly; aggressively walk /proc and kill child processes too.
killDescendants(pid) killDescendants(pid)
// Fallback kill on the shell process itself. // Fallback kill on the shell process itself.
_ = cmd.Process.Kill() _ = cmd.Process.Kill()
return nil return nil
} }
func killDescendants(ppid int) { func killDescendants(ppid int) {
if ppid <= 0 { if ppid <= 0 {
return return
} }
entries, err := os.ReadDir("/proc") entries, err := os.ReadDir("/proc")
if err != nil { if err != nil {
return return
} }
for _, e := range entries { for _, e := range entries {
if !e.IsDir() { if !e.IsDir() {
continue continue
} }
childPID, err := strconv.Atoi(e.Name()) childPID, err := strconv.Atoi(e.Name())
if err != nil || childPID <= 0 || childPID == ppid { if err != nil || childPID <= 0 || childPID == ppid {
continue continue
} }
statPath := "/proc/" + e.Name() + "/stat" statPath := "/proc/" + e.Name() + "/stat"
data, err := os.ReadFile(statPath) data, err := os.ReadFile(statPath)
if err != nil { if err != nil {
continue continue
} }
// /proc/<pid>/stat: pid (comm) state ppid ... // /proc/<pid>/stat: pid (comm) state ppid ...
raw := string(data) raw := string(data)
end := strings.LastIndex(raw, ")") end := strings.LastIndex(raw, ")")
if end == -1 || end+2 >= len(raw) { if end == -1 || end+2 >= len(raw) {
continue continue
} }
fields := strings.Fields(raw[end+2:]) fields := strings.Fields(raw[end+2:])
if len(fields) < 2 { if len(fields) < 2 {
continue continue
} }
parent, err := strconv.Atoi(fields[1]) parent, err := strconv.Atoi(fields[1])
if err != nil || parent != ppid { if err != nil || parent != ppid {
continue continue
} }
// Recurse first, then kill child process/group. // Recurse first, then kill child process/group.
killDescendants(childPID) killDescendants(childPID)
_ = syscall.Kill(-childPID, syscall.SIGKILL) _ = syscall.Kill(-childPID, syscall.SIGKILL)
_ = 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 pid := cmd.Process.Pid
if pid <= 0 { if pid <= 0 {
return nil return nil
} }
_ = exec.Command("taskkill", "/T", "/F", "/PID", strconv.Itoa(pid)).Run() _ = exec.Command("taskkill", "/T", "/F", "/PID", strconv.Itoa(pid)).Run()
_ = cmd.Process.Kill() _ = cmd.Process.Kill()
return nil return nil
} }

File diff suppressed because it is too large Load diff

View file

@ -14,73 +14,122 @@ import (
) )
func processRunning(pid int) bool { func processRunning(pid int) bool {
if pid <= 0 { if pid <= 0 {
return false return false
} }
// kill(0) can return success for zombie processes too, so inspect /proc // kill(0) can return success for zombie processes too, so inspect /proc
// state and treat zombies as not-running for timeout cleanup assertions. // state and treat zombies as not-running for timeout cleanup assertions.
err := syscall.Kill(pid, 0) err := syscall.Kill(pid, 0)
if err != nil && err != syscall.EPERM { if err != nil && err != syscall.EPERM {
return false return false
} }
data, readErr := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat") data, readErr := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat")
if readErr != nil { if readErr != nil {
return false return false
} }
raw := string(data) raw := string(data)
end := strings.LastIndex(raw, ")") end := strings.LastIndex(raw, ")")
if end == -1 || end+2 >= len(raw) { if end == -1 || end+2 >= len(raw) {
return true // best effort fallback return true // best effort fallback
} }
fields := strings.Fields(raw[end+2:]) fields := strings.Fields(raw[end+2:])
if len(fields) == 0 { if len(fields) == 0 {
return true // best effort fallback return true // best effort fallback
} }
state := fields[0] state := fields[0]
return state != "Z" return state != "Z"
} }
func TestShellTool_TimeoutKillsChildProcess(t *testing.T) { func TestShellTool_TimeoutKillsChildProcess(t *testing.T) {
tool, err := NewExecTool(t.TempDir(), false) tool, err := NewExecTool(t.TempDir(), false)
if err != nil { if err != nil {
t.Errorf("unable to configure exec tool: %s", err) t.Errorf("unable to configure exec tool: %s", err)
} }
tool.SetTimeout(500 * time.Millisecond) tool.SetTimeout(500 * time.Millisecond)
args := map[string]any{ args := map[string]any{
// Spawn a child process that would outlive the shell unless process-group kill is used. // Spawn a child process that would outlive the shell unless process-group kill is used.
"command": "sleep 60 & echo $! > child.pid; wait", "command": "sleep 60 & echo $! > child.pid; wait",
} }
result := tool.Execute(context.Background(), args) result := tool.Execute(context.Background(), args)
if !result.IsError { if !result.IsError {
t.Fatalf("expected timeout error, got success: %s", result.ForLLM) t.Fatalf("expected timeout error, got success: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "timed out") { if !strings.Contains(result.ForLLM, "timed out") {
t.Fatalf("expected timeout message, got: %s", result.ForLLM) t.Fatalf("expected timeout message, got: %s", result.ForLLM)
} }
childPIDPath := filepath.Join(tool.workingDir, "child.pid") childPIDPath := filepath.Join(tool.workingDir, "child.pid")
data, err := os.ReadFile(childPIDPath) data, err := os.ReadFile(childPIDPath)
if err != nil { if err != nil {
t.Fatalf("failed to read child pid file: %v", err) t.Fatalf("failed to read child pid file: %v", err)
} }
childPID, err := strconv.Atoi(strings.TrimSpace(string(data))) childPID, err := strconv.Atoi(strings.TrimSpace(string(data)))
if err != nil { if err != nil {
t.Fatalf("failed to parse child pid: %v", err) t.Fatalf("failed to parse child pid: %v", err)
} }
deadline := time.Now().Add(2 * time.Second) deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) { for time.Now().Before(deadline) {
if !processRunning(childPID) { if !processRunning(childPID) {
return return
} }
time.Sleep(50 * time.Millisecond) time.Sleep(50 * time.Millisecond)
} }
t.Fatalf("child process %d is still running after timeout", childPID) t.Fatalf("child process %d is still running after timeout", childPID)
} }

View file

@ -16,21 +16,31 @@ import (
) )
// InstallSkillTool allows the LLM agent to install skills from registries. // InstallSkillTool allows the LLM agent to install skills from registries.
// It shares the same RegistryManager that FindSkillsTool uses, // It shares the same RegistryManager that FindSkillsTool uses,
// so all registries configured in config are available for installation. // so all registries configured in config are available for installation.
type InstallSkillTool struct { type InstallSkillTool struct {
registryMgr *skills.RegistryManager registryMgr *skills.RegistryManager
workspace string workspace string
mu sync.Mutex mu sync.Mutex
} }
// NewInstallSkillTool creates a new InstallSkillTool. // NewInstallSkillTool creates a new InstallSkillTool.
// registryMgr is the shared registry manager (same instance as FindSkillsTool). // registryMgr is the shared registry manager (same instance as FindSkillsTool).
// workspace is the root workspace directory; skills install to {workspace}/skills/{slug}/. // workspace is the root workspace directory; skills install to {workspace}/skills/{slug}/.
func NewInstallSkillTool(registryMgr *skills.RegistryManager, workspace string) *InstallSkillTool { func NewInstallSkillTool(registryMgr *skills.RegistryManager, workspace string) *InstallSkillTool {
return &InstallSkillTool{ return &InstallSkillTool{
registryMgr: registryMgr, registryMgr: registryMgr,
workspace: workspace, workspace: workspace,
mu: sync.Mutex{}, mu: sync.Mutex{},
} }
} }
@ -46,150 +56,209 @@ func (t *InstallSkillTool) Description() string {
func (t *InstallSkillTool) Parameters() map[string]any { func (t *InstallSkillTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"slug": 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')", "description": "The unique slug of the skill to install (e.g., 'github', 'docker-compose')",
}, },
"version": map[string]any{ "version": map[string]any{
"type": "string", "type": "string",
"description": "Specific version to install (optional, defaults to latest)", "description": "Specific version to install (optional, defaults to latest)",
}, },
"registry": map[string]any{ "registry": map[string]any{
"type": "string", "type": "string",
"description": "Registry to install from (required, e.g., 'clawhub')", "description": "Registry to install from (required, e.g., 'clawhub')",
}, },
"force": map[string]any{ "force": map[string]any{
"type": "boolean", "type": "boolean",
"description": "Force reinstall if skill already exists (default false)", "description": "Force reinstall if skill already exists (default false)",
}, },
}, },
"required": []string{"slug", "registry"}, "required": []string{"slug", "registry"},
} }
} }
func (t *InstallSkillTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *InstallSkillTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
// Install lock to prevent concurrent directory operations. // Install lock to prevent concurrent directory operations.
// Ideally this should be done at a `slug` level, currently, its at a `workspace` level. // Ideally this should be done at a `slug` level, currently, its at a `workspace` level.
t.mu.Lock() t.mu.Lock()
defer t.mu.Unlock() defer t.mu.Unlock()
// Validate slug // Validate slug
slug, _ := args["slug"].(string) slug, _ := args["slug"].(string)
if err := utils.ValidateSkillIdentifier(slug); err != nil { if err := utils.ValidateSkillIdentifier(slug); err != nil {
return ErrorResult(fmt.Sprintf("invalid slug %q: error: %s", slug, err.Error())) return ErrorResult(fmt.Sprintf("invalid slug %q: error: %s", slug, err.Error()))
} }
// Validate registry // Validate registry
registryName, _ := args["registry"].(string) registryName, _ := args["registry"].(string)
if err := utils.ValidateSkillIdentifier(registryName); err != nil { if err := utils.ValidateSkillIdentifier(registryName); err != nil {
return ErrorResult(fmt.Sprintf("invalid registry %q: error: %s", registryName, err.Error())) return ErrorResult(fmt.Sprintf("invalid registry %q: error: %s", registryName, err.Error()))
} }
version, _ := args["version"].(string) version, _ := args["version"].(string)
force, _ := args["force"].(bool) force, _ := args["force"].(bool)
// Check if already installed. // Check if already installed.
skillsDir := filepath.Join(t.workspace, "skills") skillsDir := filepath.Join(t.workspace, "skills")
targetDir := filepath.Join(skillsDir, slug) targetDir := filepath.Join(skillsDir, slug)
if !force { if !force {
if _, err := os.Stat(targetDir); err == nil { if _, err := os.Stat(targetDir); err == nil {
return ErrorResult( return ErrorResult(
fmt.Sprintf("skill %q already installed at %s. Use force=true to reinstall.", slug, targetDir), fmt.Sprintf("skill %q already installed at %s. Use force=true to reinstall.", slug, targetDir),
) )
} }
} else { } else {
// Force: remove existing if present. // Force: remove existing if present.
os.RemoveAll(targetDir) os.RemoveAll(targetDir)
} }
// Resolve which registry to use. // Resolve which registry to use.
registry := t.registryMgr.GetRegistry(registryName) registry := t.registryMgr.GetRegistry(registryName)
if registry == nil { if registry == nil {
return ErrorResult(fmt.Sprintf("registry %q not found", registryName)) return ErrorResult(fmt.Sprintf("registry %q not found", registryName))
} }
// Ensure skills directory exists. // Ensure skills directory exists.
if err := os.MkdirAll(skillsDir, 0o755); err != nil { if err := os.MkdirAll(skillsDir, 0o755); err != nil {
return ErrorResult(fmt.Sprintf("failed to create skills directory: %v", err)) return ErrorResult(fmt.Sprintf("failed to create skills directory: %v", err))
} }
// Download and install (handles metadata, version resolution, extraction). // Download and install (handles metadata, version resolution, extraction).
result, err := registry.DownloadAndInstall(ctx, slug, version, targetDir) result, err := registry.DownloadAndInstall(ctx, slug, version, targetDir)
if err != nil { if err != nil {
// Clean up partial install. // Clean up partial install.
rmErr := os.RemoveAll(targetDir) rmErr := os.RemoveAll(targetDir)
if rmErr != nil { if rmErr != nil {
logger.ErrorCF("tool", "Failed to remove partial install", logger.ErrorCF("tool", "Failed to remove partial install",
map[string]any{ map[string]any{
"tool": "install_skill", "tool": "install_skill",
"target_dir": targetDir, "target_dir": targetDir,
"error": rmErr.Error(), "error": rmErr.Error(),
}) })
} }
return ErrorResult(fmt.Sprintf("failed to install %q: %v", slug, err)) return ErrorResult(fmt.Sprintf("failed to install %q: %v", slug, err))
} }
// Moderation: block malware. // Moderation: block malware.
if result.IsMalwareBlocked { if result.IsMalwareBlocked {
rmErr := os.RemoveAll(targetDir) rmErr := os.RemoveAll(targetDir)
if rmErr != nil { if rmErr != nil {
logger.ErrorCF("tool", "Failed to remove partial install", logger.ErrorCF("tool", "Failed to remove partial install",
map[string]any{ map[string]any{
"tool": "install_skill", "tool": "install_skill",
"target_dir": targetDir, "target_dir": targetDir,
"error": rmErr.Error(), "error": rmErr.Error(),
}) })
} }
return ErrorResult(fmt.Sprintf("skill %q is flagged as malicious and cannot be installed", slug)) return ErrorResult(fmt.Sprintf("skill %q is flagged as malicious and cannot be installed", slug))
} }
// Write origin metadata. // Write origin metadata.
if err := writeOriginMeta(targetDir, registry.Name(), slug, result.Version); err != nil { if err := writeOriginMeta(targetDir, registry.Name(), slug, result.Version); err != nil {
logger.ErrorCF("tool", "Failed to write origin metadata", logger.ErrorCF("tool", "Failed to write origin metadata",
map[string]any{ map[string]any{
"tool": "install_skill", "tool": "install_skill",
"error": err.Error(), "error": err.Error(),
"target": targetDir, "target": targetDir,
"registry": registry.Name(), "registry": registry.Name(),
"slug": slug, "slug": slug,
"version": result.Version, "version": result.Version,
}) })
_ = err _ = err
} }
// Build result with moderation warning if suspicious. // Build result with moderation warning if suspicious.
var output string var output string
if result.IsSuspicious { if result.IsSuspicious {
output = fmt.Sprintf("⚠️ Warning: skill %q is flagged as suspicious (may contain risky patterns).\n\n", slug) 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", output += fmt.Sprintf("Successfully installed skill %q v%s from %s registry.\nLocation: %s\n",
slug, result.Version, registry.Name(), targetDir) slug, result.Version, registry.Name(), targetDir)
if result.Summary != "" { if result.Summary != "" {
output += fmt.Sprintf("Description: %s\n", result.Summary) output += fmt.Sprintf("Description: %s\n", result.Summary)
} }
output += "\nThe skill is now available and can be loaded in the current session." output += "\nThe skill is now available and can be loaded in the current session."
return SilentResult(output) return SilentResult(output)
} }
// originMeta tracks which registry a skill was installed from. // originMeta tracks which registry a skill was installed from.
type originMeta struct { type originMeta struct {
Version int `json:"version"` Version int `json:"version"`
Registry string `json:"registry"` Registry string `json:"registry"`
Slug string `json:"slug"` Slug string `json:"slug"`
InstalledVersion string `json:"installed_version"` InstalledVersion string `json:"installed_version"`
InstalledAt int64 `json:"installed_at"` InstalledAt int64 `json:"installed_at"`
} }
func writeOriginMeta(targetDir, registryName, slug, version string) error { func writeOriginMeta(targetDir, registryName, slug, version string) error {
meta := originMeta{ meta := originMeta{
Version: 1, Version: 1,
Registry: registryName, Registry: registryName,
Slug: slug, Slug: slug,
InstalledVersion: version, InstalledVersion: version,
InstalledAt: time.Now().UnixMilli(), InstalledAt: time.Now().UnixMilli(),
} }
@ -199,5 +268,6 @@ func writeOriginMeta(targetDir, registryName, slug, version string) error {
} }
// Use unified atomic write utility with explicit sync for flash storage reliability. // Use unified atomic write utility with explicit sync for flash storage reliability.
return fileutil.WriteFileAtomic(filepath.Join(targetDir, ".skill-origin.json"), data, 0o600) return fileutil.WriteFileAtomic(filepath.Join(targetDir, ".skill-origin.json"), data, 0o600)
} }

View file

@ -14,22 +14,29 @@ import (
func TestInstallSkillToolName(t *testing.T) { func TestInstallSkillToolName(t *testing.T) {
tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir())
assert.Equal(t, "install_skill", tool.Name()) assert.Equal(t, "install_skill", tool.Name())
} }
func TestInstallSkillToolMissingSlug(t *testing.T) { func TestInstallSkillToolMissingSlug(t *testing.T) {
tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir())
result := tool.Execute(context.Background(), map[string]any{}) result := tool.Execute(context.Background(), map[string]any{})
assert.True(t, result.IsError) assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "identifier is required and must be a non-empty string") assert.Contains(t, result.ForLLM, "identifier is required and must be a non-empty string")
} }
func TestInstallSkillToolEmptySlug(t *testing.T) { func TestInstallSkillToolEmptySlug(t *testing.T) {
tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir())
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"slug": " ", "slug": " ",
}) })
assert.True(t, result.IsError) assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "identifier is required and must be a non-empty string") 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{ cases := []string{
"../etc/passwd", "../etc/passwd",
"path/traversal", "path/traversal",
"path\\traversal", "path\\traversal",
} }
@ -46,59 +55,85 @@ func TestInstallSkillToolUnsafeSlug(t *testing.T) {
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"slug": slug, "slug": slug,
}) })
assert.True(t, result.IsError, "slug %q should be rejected", slug) assert.True(t, result.IsError, "slug %q should be rejected", slug)
assert.Contains(t, result.ForLLM, "invalid slug") assert.Contains(t, result.ForLLM, "invalid slug")
} }
} }
func TestInstallSkillToolAlreadyExists(t *testing.T) { func TestInstallSkillToolAlreadyExists(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
skillDir := filepath.Join(workspace, "skills", "existing-skill") skillDir := filepath.Join(workspace, "skills", "existing-skill")
require.NoError(t, os.MkdirAll(skillDir, 0o755)) require.NoError(t, os.MkdirAll(skillDir, 0o755))
tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace) tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"slug": "existing-skill", "slug": "existing-skill",
"registry": "clawhub", "registry": "clawhub",
}) })
assert.True(t, result.IsError) assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "already installed") assert.Contains(t, result.ForLLM, "already installed")
} }
func TestInstallSkillToolRegistryNotFound(t *testing.T) { func TestInstallSkillToolRegistryNotFound(t *testing.T) {
workspace := t.TempDir() workspace := t.TempDir()
tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace) tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"slug": "some-skill", "slug": "some-skill",
"registry": "nonexistent", "registry": "nonexistent",
}) })
assert.True(t, result.IsError) assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "registry") assert.Contains(t, result.ForLLM, "registry")
assert.Contains(t, result.ForLLM, "not found") assert.Contains(t, result.ForLLM, "not found")
} }
func TestInstallSkillToolParameters(t *testing.T) { func TestInstallSkillToolParameters(t *testing.T) {
tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir())
params := tool.Parameters() params := tool.Parameters()
props, ok := params["properties"].(map[string]any) props, ok := params["properties"].(map[string]any)
assert.True(t, ok) assert.True(t, ok)
assert.Contains(t, props, "slug") assert.Contains(t, props, "slug")
assert.Contains(t, props, "version") assert.Contains(t, props, "version")
assert.Contains(t, props, "registry") assert.Contains(t, props, "registry")
assert.Contains(t, props, "force") assert.Contains(t, props, "force")
required, ok := params["required"].([]string) required, ok := params["required"].([]string)
assert.True(t, ok) assert.True(t, ok)
assert.Contains(t, required, "slug") assert.Contains(t, required, "slug")
assert.Contains(t, required, "registry") assert.Contains(t, required, "registry")
} }
func TestInstallSkillToolMissingRegistry(t *testing.T) { func TestInstallSkillToolMissingRegistry(t *testing.T) {
tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir())
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"slug": "some-skill", "slug": "some-skill",
}) })
assert.True(t, result.IsError) assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "invalid registry") assert.Contains(t, result.ForLLM, "invalid registry")
} }

View file

@ -9,17 +9,23 @@ import (
) )
// FindSkillsTool allows the LLM agent to search for installable skills from registries. // FindSkillsTool allows the LLM agent to search for installable skills from registries.
type FindSkillsTool struct { type FindSkillsTool struct {
registryMgr *skills.RegistryManager registryMgr *skills.RegistryManager
cache *skills.SearchCache cache *skills.SearchCache
} }
// NewFindSkillsTool creates a new FindSkillsTool. // NewFindSkillsTool creates a new FindSkillsTool.
// registryMgr is the shared registry manager (built from config in createToolRegistry). // registryMgr is the shared registry manager (built from config in createToolRegistry).
// cache is the search cache for deduplicating similar queries. // cache is the search cache for deduplicating similar queries.
func NewFindSkillsTool(registryMgr *skills.RegistryManager, cache *skills.SearchCache) *FindSkillsTool { func NewFindSkillsTool(registryMgr *skills.RegistryManager, cache *skills.SearchCache) *FindSkillsTool {
return &FindSkillsTool{ return &FindSkillsTool{
registryMgr: registryMgr, registryMgr: registryMgr,
cache: cache, cache: cache,
} }
} }
@ -35,38 +41,50 @@ func (t *FindSkillsTool) Description() string {
func (t *FindSkillsTool) Parameters() map[string]any { func (t *FindSkillsTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"query": 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')", "description": "Search query describing the desired skill capability (e.g., 'github integration', 'database management')",
}, },
"limit": map[string]any{ "limit": map[string]any{
"type": "integer", "type": "integer",
"description": "Maximum number of results to return (1-20, default 5)", "description": "Maximum number of results to return (1-20, default 5)",
"minimum": 1.0, "minimum": 1.0,
"maximum": 20.0, "maximum": 20.0,
}, },
}, },
"required": []string{"query"}, "required": []string{"query"},
} }
} }
func (t *FindSkillsTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *FindSkillsTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
query, ok := args["query"].(string) query, ok := args["query"].(string)
query = strings.ToLower(strings.TrimSpace(query)) query = strings.ToLower(strings.TrimSpace(query))
if !ok || query == "" { if !ok || query == "" {
return ErrorResult("query is required and must be a non-empty string") return ErrorResult("query is required and must be a non-empty string")
} }
limit := 5 limit := 5
if l, ok := args["limit"].(float64); ok { if l, ok := args["limit"].(float64); ok {
li := int(l) li := int(l)
if li >= 1 && li <= 20 { if li >= 1 && li <= 20 {
limit = li limit = li
} }
} }
// Check cache first. // Check cache first.
if t.cache != nil { if t.cache != nil {
if cached, hit := t.cache.Get(query); hit { if cached, hit := t.cache.Get(query); hit {
return SilentResult(formatSearchResults(query, cached, true)) 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. // Search all registries.
results, err := t.registryMgr.SearchAll(ctx, query, limit) results, err := t.registryMgr.SearchAll(ctx, query, limit)
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("skill search failed: %v", err)) return ErrorResult(fmt.Sprintf("skill search failed: %v", err))
} }
// Cache the results. // Cache the results.
if t.cache != nil && len(results) > 0 { if t.cache != nil && len(results) > 0 {
t.cache.Put(query, results) t.cache.Put(query, results)
} }
@ -93,27 +113,36 @@ func formatSearchResults(query string, results []skills.SearchResult, cached boo
} }
var sb strings.Builder var sb strings.Builder
source := "" source := ""
if cached { if cached {
source = " (cached)" source = " (cached)"
} }
sb.WriteString(fmt.Sprintf("Found %d skills for %q%s:\n\n", len(results), query, source)) sb.WriteString(fmt.Sprintf("Found %d skills for %q%s:\n\n", len(results), query, source))
for i, r := range results { for i, r := range results {
sb.WriteString(fmt.Sprintf("%d. **%s**", i+1, r.Slug)) sb.WriteString(fmt.Sprintf("%d. **%s**", i+1, r.Slug))
if r.Version != "" { if r.Version != "" {
sb.WriteString(fmt.Sprintf(" v%s", r.Version)) sb.WriteString(fmt.Sprintf(" v%s", r.Version))
} }
sb.WriteString(fmt.Sprintf(" (score: %.3f, registry: %s)\n", r.Score, r.RegistryName)) sb.WriteString(fmt.Sprintf(" (score: %.3f, registry: %s)\n", r.Score, r.RegistryName))
if r.DisplayName != "" && r.DisplayName != r.Slug { if r.DisplayName != "" && r.DisplayName != r.Slug {
sb.WriteString(fmt.Sprintf(" Name: %s\n", r.DisplayName)) sb.WriteString(fmt.Sprintf(" Name: %s\n", r.DisplayName))
} }
if r.Summary != "" { if r.Summary != "" {
sb.WriteString(fmt.Sprintf(" %s\n", r.Summary)) sb.WriteString(fmt.Sprintf(" %s\n", r.Summary))
} }
sb.WriteString("\n") sb.WriteString("\n")
} }
sb.WriteString("Use install_skill with the slug to install a skill.") sb.WriteString("Use install_skill with the slug to install a skill.")
return sb.String() return sb.String()
} }

View file

@ -11,62 +11,81 @@ import (
func TestFindSkillsToolName(t *testing.T) { func TestFindSkillsToolName(t *testing.T) {
tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) tool := NewFindSkillsTool(skills.NewRegistryManager(), nil)
assert.Equal(t, "find_skills", tool.Name()) assert.Equal(t, "find_skills", tool.Name())
} }
func TestFindSkillsToolMissingQuery(t *testing.T) { func TestFindSkillsToolMissingQuery(t *testing.T) {
tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) tool := NewFindSkillsTool(skills.NewRegistryManager(), nil)
result := tool.Execute(context.Background(), map[string]any{}) result := tool.Execute(context.Background(), map[string]any{})
assert.True(t, result.IsError) assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "query is required") assert.Contains(t, result.ForLLM, "query is required")
} }
func TestFindSkillsToolEmptyQuery(t *testing.T) { func TestFindSkillsToolEmptyQuery(t *testing.T) {
tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) tool := NewFindSkillsTool(skills.NewRegistryManager(), nil)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"query": " ", "query": " ",
}) })
assert.True(t, result.IsError) assert.True(t, result.IsError)
} }
func TestFindSkillsToolCacheHit(t *testing.T) { func TestFindSkillsToolCacheHit(t *testing.T) {
cache := skills.NewSearchCache(10, 5*60*1000*1000*1000) // 5 min cache := skills.NewSearchCache(10, 5*60*1000*1000*1000) // 5 min
cache.Put("github", []skills.SearchResult{ cache.Put("github", []skills.SearchResult{
{Slug: "github", Score: 0.9, RegistryName: "clawhub"}, {Slug: "github", Score: 0.9, RegistryName: "clawhub"},
}) })
tool := NewFindSkillsTool(skills.NewRegistryManager(), cache) tool := NewFindSkillsTool(skills.NewRegistryManager(), cache)
result := tool.Execute(context.Background(), map[string]any{ result := tool.Execute(context.Background(), map[string]any{
"query": "github", "query": "github",
}) })
assert.False(t, result.IsError) assert.False(t, result.IsError)
assert.Contains(t, result.ForLLM, "github") assert.Contains(t, result.ForLLM, "github")
assert.Contains(t, result.ForLLM, "cached") assert.Contains(t, result.ForLLM, "cached")
} }
func TestFindSkillsToolParameters(t *testing.T) { func TestFindSkillsToolParameters(t *testing.T) {
tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) tool := NewFindSkillsTool(skills.NewRegistryManager(), nil)
params := tool.Parameters() params := tool.Parameters()
props, ok := params["properties"].(map[string]any) props, ok := params["properties"].(map[string]any)
assert.True(t, ok) assert.True(t, ok)
assert.Contains(t, props, "query") assert.Contains(t, props, "query")
assert.Contains(t, props, "limit") assert.Contains(t, props, "limit")
required, ok := params["required"].([]string) required, ok := params["required"].([]string)
assert.True(t, ok) assert.True(t, ok)
assert.Contains(t, required, "query") assert.Contains(t, required, "query")
} }
func TestFindSkillsToolDescription(t *testing.T) { func TestFindSkillsToolDescription(t *testing.T) {
tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) tool := NewFindSkillsTool(skills.NewRegistryManager(), nil)
assert.NotEmpty(t, tool.Description()) assert.NotEmpty(t, tool.Description())
assert.Contains(t, tool.Description(), "skill") assert.Contains(t, tool.Description(), "skill")
} }
func TestFormatSearchResultsEmpty(t *testing.T) { func TestFormatSearchResultsEmpty(t *testing.T) {
result := formatSearchResults("test query", nil, false) result := formatSearchResults("test query", nil, false)
assert.Contains(t, result, "No skills found") assert.Contains(t, result, "No skills found")
} }
@ -74,17 +93,28 @@ func TestFormatSearchResultsWithData(t *testing.T) {
results := []skills.SearchResult{ results := []skills.SearchResult{
{ {
Slug: "github", Slug: "github",
Score: 0.95, Score: 0.95,
DisplayName: "GitHub", DisplayName: "GitHub",
Summary: "GitHub API integration", Summary: "GitHub API integration",
Version: "1.0.0", Version: "1.0.0",
RegistryName: "clawhub", RegistryName: "clawhub",
}, },
} }
output := formatSearchResults("github", results, false) output := formatSearchResults("github", results, false)
assert.Contains(t, output, "github") assert.Contains(t, output, "github")
assert.Contains(t, output, "v1.0.0") assert.Contains(t, output, "v1.0.0")
assert.Contains(t, output, "0.950") assert.Contains(t, output, "0.950")
assert.Contains(t, output, "clawhub") assert.Contains(t, output, "clawhub")
assert.Contains(t, output, "install_skill") assert.Contains(t, output, "install_skill")
} }

View file

@ -8,31 +8,41 @@ import (
func TestSpawnTool_Execute_EmptyTask(t *testing.T) { func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{})
tool := NewSpawnTool(manager) tool := NewSpawnTool(manager)
ctx := context.Background() ctx := context.Background()
tests := []struct { tests := []struct {
name string name string
args map[string]any args map[string]any
}{ }{
{"empty string", map[string]any{"task": ""}}, {"empty string", map[string]any{"task": ""}},
{"whitespace only", map[string]any{"task": " "}}, {"whitespace only", map[string]any{"task": " "}},
{"tabs and newlines", map[string]any{"task": "\t\n "}}, {"tabs and newlines", map[string]any{"task": "\t\n "}},
{"missing task key", map[string]any{"label": "test"}}, {"missing task key", map[string]any{"label": "test"}},
{"wrong type", map[string]any{"task": 123}}, {"wrong type", map[string]any{"task": 123}},
} }
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
result := tool.Execute(ctx, tt.args) result := tool.Execute(ctx, tt.args)
if result == nil { if result == nil {
t.Fatal("Result should not be nil") t.Fatal("Result should not be nil")
} }
if !result.IsError { if !result.IsError {
t.Error("Expected error for invalid task parameter") t.Error("Expected error for invalid task parameter")
} }
if !strings.Contains(result.ForLLM, `"task"`) { if !strings.Contains(result.ForLLM, `"task"`) {
t.Errorf("Error message should mention '\"task\"', got: %s", result.ForLLM) 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) { func TestSpawnTool_Execute_ValidTask(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{})
tool := NewSpawnTool(manager) tool := NewSpawnTool(manager)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"task": "Write a haiku about coding", "task": "Write a haiku about coding",
"label": "haiku-task", "label": "haiku-task",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
if result == nil { if result == nil {
t.Fatal("Result should not be nil") t.Fatal("Result should not be nil")
} }
if result.IsError { if result.IsError {
t.Errorf("Expected success for valid task, got error: %s", result.ForLLM) t.Errorf("Expected success for valid task, got error: %s", result.ForLLM)
} }
if !result.Async { if !result.Async {
t.Error("SpawnTool should return async result") t.Error("SpawnTool should return async result")
} }
@ -67,12 +84,15 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) {
tool := NewSpawnTool(nil) tool := NewSpawnTool(nil)
ctx := context.Background() ctx := context.Background()
args := map[string]any{"task": "test task"} args := map[string]any{"task": "test task"}
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
if !result.IsError { if !result.IsError {
t.Error("Expected error for nil manager") t.Error("Expected error for nil manager")
} }
if !strings.Contains(result.ForLLM, "spawn tool is not available") { if !strings.Contains(result.ForLLM, "spawn tool is not available") {
t.Errorf("Error message should mention spawn tool not available, got: %s", result.ForLLM) 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. // SPITool provides SPI bus interaction for high-speed peripheral communication.
type SPITool struct{} type SPITool struct{}
func NewSPITool() *SPITool { func NewSPITool() *SPITool {
@ -27,42 +28,61 @@ func (t *SPITool) Description() string {
func (t *SPITool) Parameters() map[string]any { func (t *SPITool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"action": map[string]any{ "action": map[string]any{
"type": "string", "type": "string",
"enum": []string{"list", "transfer", "read"}, "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)", "description": "Action to perform: list (find available SPI devices), transfer (full-duplex send/receive), read (receive bytes by sending zeros)",
}, },
"device": map[string]any{ "device": map[string]any{
"type": "string", "type": "string",
"description": "SPI device identifier (e.g. \"2.0\" for /dev/spidev2.0). Required for transfer/read.", "description": "SPI device identifier (e.g. \"2.0\" for /dev/spidev2.0). Required for transfer/read.",
}, },
"speed": map[string]any{ "speed": map[string]any{
"type": "integer", "type": "integer",
"description": "SPI clock speed in Hz. Default: 1000000 (1 MHz).", "description": "SPI clock speed in Hz. Default: 1000000 (1 MHz).",
}, },
"mode": map[string]any{ "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.", "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{ "bits": map[string]any{
"type": "integer", "type": "integer",
"description": "Bits per word. Default: 8.", "description": "Bits per word. Default: 8.",
}, },
"data": map[string]any{ "data": map[string]any{
"type": "array", "type": "array",
"items": map[string]any{"type": "integer"}, "items": map[string]any{"type": "integer"},
"description": "Bytes to send (0-255 each). Required for transfer action.", "description": "Bytes to send (0-255 each). Required for transfer action.",
}, },
"length": map[string]any{ "length": map[string]any{
"type": "integer", "type": "integer",
"description": "Number of bytes to read (1-4096). Required for read action.", "description": "Number of bytes to read (1-4096). Required for read action.",
}, },
"confirm": map[string]any{ "confirm": map[string]any{
"type": "boolean", "type": "boolean",
"description": "Must be true for transfer operations. Safety guard to prevent accidental writes.", "description": "Must be true for transfer operations. Safety guard to prevent accidental writes.",
}, },
}, },
"required": []string{"action"}, "required": []string{"action"},
} }
} }
@ -73,23 +93,32 @@ func (t *SPITool) Execute(ctx context.Context, args map[string]any) *ToolResult
} }
action, ok := args["action"].(string) action, ok := args["action"].(string)
if !ok { if !ok {
return ErrorResult("action is required") return ErrorResult("action is required")
} }
switch action { switch action {
case "list": case "list":
return t.list() return t.list()
case "transfer": case "transfer":
return t.transfer(args) return t.transfer(args)
case "read": case "read":
return t.readDevice(args) return t.readDevice(args)
default: default:
return ErrorResult(fmt.Sprintf("unknown action: %s (valid: list, transfer, read)", action)) return ErrorResult(fmt.Sprintf("unknown action: %s (valid: list, transfer, read)", action))
} }
} }
// list finds available SPI devices by globbing /dev/spidev* // list finds available SPI devices by globbing /dev/spidev*
func (t *SPITool) list() *ToolResult { func (t *SPITool) list() *ToolResult {
matches, err := filepath.Glob("/dev/spidev*") matches, err := filepath.Glob("/dev/spidev*")
if err != nil { if err != nil {
@ -104,11 +133,14 @@ func (t *SPITool) list() *ToolResult {
type devInfo struct { type devInfo struct {
Path string `json:"path"` Path string `json:"path"`
Device string `json:"device"` Device string `json:"device"`
} }
devices := make([]devInfo, 0, len(matches)) devices := make([]devInfo, 0, len(matches))
re := regexp.MustCompile(`/dev/spidev(\d+\.\d+)`) re := regexp.MustCompile(`/dev/spidev(\d+\.\d+)`)
for _, m := range matches { for _, m := range matches {
if sub := re.FindStringSubmatch(m); sub != nil { if sub := re.FindStringSubmatch(m); sub != nil {
devices = append(devices, devInfo{Path: m, Device: sub[1]}) devices = append(devices, devInfo{Path: m, Device: sub[1]})
@ -116,45 +148,58 @@ func (t *SPITool) list() *ToolResult {
} }
result, _ := json.MarshalIndent(devices, "", " ") result, _ := json.MarshalIndent(devices, "", " ")
return SilentResult(fmt.Sprintf("Found %d SPI device(s):\n%s", len(devices), string(result))) 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) // Helper function for SPI operations (used by platform-specific implementations)
// parseSPIArgs extracts and validates common SPI parameters // parseSPIArgs extracts and validates common SPI parameters
// //
//nolint:unused // Used by spi_linux.go //nolint:unused // Used by spi_linux.go
func parseSPIArgs(args map[string]any) (device string, speed uint32, mode uint8, bits uint8, errMsg string) { func parseSPIArgs(args map[string]any) (device string, speed uint32, mode uint8, bits uint8, errMsg string) {
dev, ok := args["device"].(string) dev, ok := args["device"].(string)
if !ok || dev == "" { if !ok || dev == "" {
return "", 0, 0, 0, "device is required (e.g. \"2.0\" for /dev/spidev2.0)" return "", 0, 0, 0, "device is required (e.g. \"2.0\" for /dev/spidev2.0)"
} }
matched, _ := regexp.MatchString(`^\d+\.\d+$`, dev) matched, _ := regexp.MatchString(`^\d+\.\d+$`, dev)
if !matched { if !matched {
return "", 0, 0, 0, "invalid device identifier: must be in format \"X.Y\" (e.g. \"2.0\")" return "", 0, 0, 0, "invalid device identifier: must be in format \"X.Y\" (e.g. \"2.0\")"
} }
speed = 1000000 // default 1 MHz speed = 1000000 // default 1 MHz
if s, ok := args["speed"].(float64); ok { if s, ok := args["speed"].(float64); ok {
if s < 1 || s > 125000000 { if s < 1 || s > 125000000 {
return "", 0, 0, 0, "speed must be between 1 Hz and 125 MHz" return "", 0, 0, 0, "speed must be between 1 Hz and 125 MHz"
} }
speed = uint32(s) speed = uint32(s)
} }
mode = 0 mode = 0
if m, ok := args["mode"].(float64); ok { if m, ok := args["mode"].(float64); ok {
if int(m) < 0 || int(m) > 3 { if int(m) < 0 || int(m) > 3 {
return "", 0, 0, 0, "mode must be 0-3" return "", 0, 0, 0, "mode must be 0-3"
} }
mode = uint8(m) mode = uint8(m)
} }
bits = 8 bits = 8
if b, ok := args["bits"].(float64); ok { if b, ok := args["bits"].(float64); ok {
if int(b) < 1 || int(b) > 32 { if int(b) < 1 || int(b) > 32 {
return "", 0, 0, 0, "bits must be between 1 and 32" return "", 0, 0, 0, "bits must be between 1 and 32"
} }
bits = uint8(b) bits = uint8(b)
} }

View file

@ -9,190 +9,321 @@ import (
) )
// SPI ioctl constants from Linux kernel headers. // SPI ioctl constants from Linux kernel headers.
// Calculated from _IOW('k', nr, size) macro: // Calculated from _IOW('k', nr, size) macro:
// //
// direction(1)<<30 | size<<16 | type(0x6B)<<8 | nr // direction(1)<<30 | size<<16 | type(0x6B)<<8 | nr
const ( const (
spiIocWrMode = 0x40016B01 // _IOW('k', 1, __u8) spiIocWrMode = 0x40016B01 // _IOW('k', 1, __u8)
spiIocWrBitsPerWord = 0x40016B03 // _IOW('k', 3, __u8) spiIocWrBitsPerWord = 0x40016B03 // _IOW('k', 3, __u8)
spiIocWrMaxSpeedHz = 0x40046B04 // _IOW('k', 4, __u32) spiIocWrMaxSpeedHz = 0x40046B04 // _IOW('k', 4, __u32)
spiIocMessage1 = 0x40206B00 // _IOW('k', 0, struct spi_ioc_transfer) — 32 bytes spiIocMessage1 = 0x40206B00 // _IOW('k', 0, struct spi_ioc_transfer) — 32 bytes
) )
// spiTransfer matches Linux kernel struct spi_ioc_transfer (32 bytes on all architectures). // spiTransfer matches Linux kernel struct spi_ioc_transfer (32 bytes on all architectures).
type spiTransfer struct { type spiTransfer struct {
txBuf uint64 txBuf uint64
rxBuf uint64 rxBuf uint64
length uint32 length uint32
speedHz uint32 speedHz uint32
delayUsecs uint16 delayUsecs uint16
bitsPerWord uint8 bitsPerWord uint8
csChange uint8 csChange uint8
txNbits uint8 txNbits uint8
rxNbits uint8 rxNbits uint8
wordDelay uint8 wordDelay uint8
pad uint8 pad uint8
} }
// configureSPI opens an SPI device and sets mode, bits per word, and speed // 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) { func configureSPI(devPath string, mode uint8, bits uint8, speed uint32) (int, *ToolResult) {
fd, err := syscall.Open(devPath, syscall.O_RDWR, 0) fd, err := syscall.Open(devPath, syscall.O_RDWR, 0)
if err != nil { if err != nil {
return -1, ErrorResult(fmt.Sprintf("failed to open %s: %v (check permissions and spidev module)", devPath, err)) return -1, ErrorResult(fmt.Sprintf("failed to open %s: %v (check permissions and spidev module)", devPath, err))
} }
// Set SPI mode // Set SPI mode
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrMode, uintptr(unsafe.Pointer(&mode))) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrMode, uintptr(unsafe.Pointer(&mode)))
if errno != 0 { if errno != 0 {
syscall.Close(fd) syscall.Close(fd)
return -1, ErrorResult(fmt.Sprintf("failed to set SPI mode %d: %v", mode, errno)) return -1, ErrorResult(fmt.Sprintf("failed to set SPI mode %d: %v", mode, errno))
} }
// Set bits per word // Set bits per word
_, _, errno = syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrBitsPerWord, uintptr(unsafe.Pointer(&bits))) _, _, errno = syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrBitsPerWord, uintptr(unsafe.Pointer(&bits)))
if errno != 0 { if errno != 0 {
syscall.Close(fd) syscall.Close(fd)
return -1, ErrorResult(fmt.Sprintf("failed to set bits per word %d: %v", bits, errno)) return -1, ErrorResult(fmt.Sprintf("failed to set bits per word %d: %v", bits, errno))
} }
// Set max speed // Set max speed
_, _, errno = syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrMaxSpeedHz, uintptr(unsafe.Pointer(&speed))) _, _, errno = syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocWrMaxSpeedHz, uintptr(unsafe.Pointer(&speed)))
if errno != 0 { if errno != 0 {
syscall.Close(fd) syscall.Close(fd)
return -1, ErrorResult(fmt.Sprintf("failed to set SPI speed %d Hz: %v", speed, errno)) return -1, ErrorResult(fmt.Sprintf("failed to set SPI speed %d Hz: %v", speed, errno))
} }
return fd, nil return fd, nil
} }
// transfer performs a full-duplex SPI transfer // transfer performs a full-duplex SPI transfer
func (t *SPITool) transfer(args map[string]any) *ToolResult { func (t *SPITool) transfer(args map[string]any) *ToolResult {
confirm, _ := args["confirm"].(bool) confirm, _ := args["confirm"].(bool)
if !confirm { if !confirm {
return ErrorResult( return ErrorResult(
"transfer operations require confirm: true. Please confirm with the user before sending data to SPI devices.", "transfer operations require confirm: true. Please confirm with the user before sending data to SPI devices.",
) )
} }
dev, speed, mode, bits, errMsg := parseSPIArgs(args) dev, speed, mode, bits, errMsg := parseSPIArgs(args)
if errMsg != "" { if errMsg != "" {
return ErrorResult(errMsg) return ErrorResult(errMsg)
} }
dataRaw, ok := args["data"].([]any) dataRaw, ok := args["data"].([]any)
if !ok || len(dataRaw) == 0 { if !ok || len(dataRaw) == 0 {
return ErrorResult("data is required for transfer (array of byte values 0-255)") return ErrorResult("data is required for transfer (array of byte values 0-255)")
} }
if len(dataRaw) > 4096 { if len(dataRaw) > 4096 {
return ErrorResult("data too long: maximum 4096 bytes per SPI transfer") return ErrorResult("data too long: maximum 4096 bytes per SPI transfer")
} }
txBuf := make([]byte, len(dataRaw)) txBuf := make([]byte, len(dataRaw))
for i, v := range dataRaw { for i, v := range dataRaw {
f, ok := v.(float64) f, ok := v.(float64)
if !ok { if !ok {
return ErrorResult(fmt.Sprintf("data[%d] is not a valid byte value", i)) return ErrorResult(fmt.Sprintf("data[%d] is not a valid byte value", i))
} }
b := int(f) b := int(f)
if b < 0 || b > 255 { if b < 0 || b > 255 {
return ErrorResult(fmt.Sprintf("data[%d] = %d is out of byte range (0-255)", i, b)) return ErrorResult(fmt.Sprintf("data[%d] = %d is out of byte range (0-255)", i, b))
} }
txBuf[i] = byte(b) txBuf[i] = byte(b)
} }
devPath := fmt.Sprintf("/dev/spidev%s", dev) devPath := fmt.Sprintf("/dev/spidev%s", dev)
fd, errResult := configureSPI(devPath, mode, bits, speed) fd, errResult := configureSPI(devPath, mode, bits, speed)
if errResult != nil { if errResult != nil {
return errResult return errResult
} }
defer syscall.Close(fd) defer syscall.Close(fd)
rxBuf := make([]byte, len(txBuf)) rxBuf := make([]byte, len(txBuf))
xfer := spiTransfer{ xfer := spiTransfer{
txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))), txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))),
rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))), rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))),
length: uint32(len(txBuf)), length: uint32(len(txBuf)),
speedHz: speed, speedHz: speed,
bitsPerWord: bits, bitsPerWord: bits,
} }
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocMessage1, uintptr(unsafe.Pointer(&xfer))) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocMessage1, uintptr(unsafe.Pointer(&xfer)))
runtime.KeepAlive(txBuf) runtime.KeepAlive(txBuf)
runtime.KeepAlive(rxBuf) runtime.KeepAlive(rxBuf)
if errno != 0 { if errno != 0 {
return ErrorResult(fmt.Sprintf("SPI transfer failed: %v", errno)) return ErrorResult(fmt.Sprintf("SPI transfer failed: %v", errno))
} }
// Format received bytes // Format received bytes
hexBytes := make([]string, len(rxBuf)) hexBytes := make([]string, len(rxBuf))
intBytes := make([]int, len(rxBuf)) intBytes := make([]int, len(rxBuf))
for i, b := range rxBuf { for i, b := range rxBuf {
hexBytes[i] = fmt.Sprintf("0x%02x", b) hexBytes[i] = fmt.Sprintf("0x%02x", b)
intBytes[i] = int(b) intBytes[i] = int(b)
} }
result, _ := json.MarshalIndent(map[string]any{ result, _ := json.MarshalIndent(map[string]any{
"device": devPath, "device": devPath,
"sent": len(txBuf), "sent": len(txBuf),
"received": intBytes, "received": intBytes,
"hex": hexBytes, "hex": hexBytes,
}, "", " ") }, "", " ")
return SilentResult(string(result)) return SilentResult(string(result))
} }
// readDevice reads bytes from SPI by sending zeros (read-only, no confirm needed) // readDevice reads bytes from SPI by sending zeros (read-only, no confirm needed)
func (t *SPITool) readDevice(args map[string]any) *ToolResult { func (t *SPITool) readDevice(args map[string]any) *ToolResult {
dev, speed, mode, bits, errMsg := parseSPIArgs(args) dev, speed, mode, bits, errMsg := parseSPIArgs(args)
if errMsg != "" { if errMsg != "" {
return ErrorResult(errMsg) return ErrorResult(errMsg)
} }
length := 0 length := 0
if l, ok := args["length"].(float64); ok { if l, ok := args["length"].(float64); ok {
length = int(l) length = int(l)
} }
if length < 1 || length > 4096 { if length < 1 || length > 4096 {
return ErrorResult("length is required for read (1-4096)") return ErrorResult("length is required for read (1-4096)")
} }
devPath := fmt.Sprintf("/dev/spidev%s", dev) devPath := fmt.Sprintf("/dev/spidev%s", dev)
fd, errResult := configureSPI(devPath, mode, bits, speed) fd, errResult := configureSPI(devPath, mode, bits, speed)
if errResult != nil { if errResult != nil {
return errResult return errResult
} }
defer syscall.Close(fd) defer syscall.Close(fd)
txBuf := make([]byte, length) // zeros txBuf := make([]byte, length) // zeros
rxBuf := make([]byte, length) rxBuf := make([]byte, length)
xfer := spiTransfer{ xfer := spiTransfer{
txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))), txBuf: uint64(uintptr(unsafe.Pointer(&txBuf[0]))),
rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))), rxBuf: uint64(uintptr(unsafe.Pointer(&rxBuf[0]))),
length: uint32(length), length: uint32(length),
speedHz: speed, speedHz: speed,
bitsPerWord: bits, bitsPerWord: bits,
} }
_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocMessage1, uintptr(unsafe.Pointer(&xfer))) _, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uintptr(fd), spiIocMessage1, uintptr(unsafe.Pointer(&xfer)))
runtime.KeepAlive(txBuf) runtime.KeepAlive(txBuf)
runtime.KeepAlive(rxBuf) runtime.KeepAlive(rxBuf)
if errno != 0 { if errno != 0 {
return ErrorResult(fmt.Sprintf("SPI read failed: %v", errno)) return ErrorResult(fmt.Sprintf("SPI read failed: %v", errno))
} }
hexBytes := make([]string, len(rxBuf)) hexBytes := make([]string, len(rxBuf))
intBytes := make([]int, len(rxBuf)) intBytes := make([]int, len(rxBuf))
for i, b := range rxBuf { for i, b := range rxBuf {
hexBytes[i] = fmt.Sprintf("0x%02x", b) hexBytes[i] = fmt.Sprintf("0x%02x", b)
intBytes[i] = int(b) intBytes[i] = int(b)
} }
result, _ := json.MarshalIndent(map[string]any{ result, _ := json.MarshalIndent(map[string]any{
"device": devPath, "device": devPath,
"bytes": intBytes, "bytes": intBytes,
"hex": hexBytes, "hex": hexBytes,
"length": len(rxBuf), "length": len(rxBuf),
}, "", " ") }, "", " ")
return SilentResult(string(result)) return SilentResult(string(result))
} }

View file

@ -11,8 +11,11 @@ import (
) )
// blockingProvider blocks inside Chat until the context is canceled. // blockingProvider blocks inside Chat until the context is canceled.
// The ready channel is closed the moment Chat is entered, so callers can // The ready channel is closed the moment Chat is entered, so callers can
// synchronize before canceling the context. // synchronize before canceling the context.
type blockingProvider struct { type blockingProvider struct {
ready chan struct{} ready chan struct{}
} }
@ -23,42 +26,63 @@ func newBlockingProvider() *blockingProvider {
func (p *blockingProvider) Chat( func (p *blockingProvider) Chat(
ctx context.Context, ctx context.Context,
_ []providers.Message, _ []providers.Message,
_ []providers.ToolDefinition, _ []providers.ToolDefinition,
_ string, _ string,
_ map[string]any, _ map[string]any,
) (*providers.LLMResponse, error) { ) (*providers.LLMResponse, error) {
close(p.ready) // signal: we are now blocking close(p.ready) // signal: we are now blocking
<-ctx.Done() <-ctx.Done()
return nil, ctx.Err() return nil, ctx.Err()
} }
func (p *blockingProvider) GetDefaultModel() string { return "test" } func (p *blockingProvider) GetDefaultModel() string { return "test" }
// TestSubagentManager_Spawn_EmitsLifecycleEvents verifies that Spawn() fires // TestSubagentManager_Spawn_EmitsLifecycleEvents verifies that Spawn() fires
// the correct sequence of orchestration events through a real Broadcaster: // the correct sequence of orchestration events through a real Broadcaster:
// //
// agent_spawn → conversation(conductor→sub) → agent_state(waiting) → // agent_spawn → conversation(conductor→sub) → agent_state(waiting) →
// conversation(sub→conductor) → agent_gc(completed) // conversation(sub→conductor) → agent_gc(completed)
// //
// It also verifies that the snapshot is empty after ReportGC and that the // It also verifies that the snapshot is empty after ReportGC and that the
// completion callback is invoked. // completion callback is invoked.
func TestSubagentManager_Spawn_EmitsLifecycleEvents(t *testing.T) { func TestSubagentManager_Spawn_EmitsLifecycleEvents(t *testing.T) {
b := orch.NewBroadcaster() b := orch.NewBroadcaster()
sub := b.Subscribe() sub := b.Subscribe()
defer b.Unsubscribe(sub) defer b.Unsubscribe(sub)
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
mgr := NewSubagentManager(provider, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) mgr := NewSubagentManager(provider, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{})
var callbackCalled int32 var callbackCalled int32
cb := AsyncCallback(func(_ context.Context, _ *ToolResult) { cb := AsyncCallback(func(_ context.Context, _ *ToolResult) {
atomic.StoreInt32(&callbackCalled, 1) atomic.StoreInt32(&callbackCalled, 1)
}) })
_, err := mgr.Spawn( _, err := mgr.Spawn(
context.Background(), context.Background(),
"say hello", "hello-task", "", "cli", "direct", "", "say hello", "hello-task", "", "cli", "direct", "",
cb, cb,
) )
if err != nil { if err != nil {
@ -66,94 +90,131 @@ func TestSubagentManager_Spawn_EmitsLifecycleEvents(t *testing.T) {
} }
// Collect events until agent_gc or timeout. // Collect events until agent_gc or timeout.
var events []orch.Event var events []orch.Event
deadline := time.After(3 * time.Second) deadline := time.After(3 * time.Second)
loop: loop:
for { for {
select { select {
case ev := <-sub.Ch: case ev := <-sub.Ch:
events = append(events, ev) events = append(events, ev)
if ev.Type == "agent_gc" { if ev.Type == "agent_gc" {
break loop break loop
} }
case <-deadline: case <-deadline:
t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events)
} }
} }
// 1. First event must be agent_spawn with the correct label. // 1. First event must be agent_spawn with the correct label.
if len(events) == 0 || events[0].Type != "agent_spawn" { if len(events) == 0 || events[0].Type != "agent_spawn" {
t.Fatalf("first event must be agent_spawn, got: %+v", events) t.Fatalf("first event must be agent_spawn, got: %+v", events)
} }
if events[0].Label != "hello-task" { if events[0].Label != "hello-task" {
t.Errorf("agent_spawn label = %q, want %q", events[0].Label, "hello-task") t.Errorf("agent_spawn label = %q, want %q", events[0].Label, "hello-task")
} }
spawnedID := events[0].ID spawnedID := events[0].ID
// 2. There must be a conversation from conductor → subagent. // 2. There must be a conversation from conductor → subagent.
var hasConvToSub bool var hasConvToSub bool
for _, ev := range events { for _, ev := range events {
if ev.Type == "conversation" && ev.From == "conductor" && ev.To == spawnedID { if ev.Type == "conversation" && ev.From == "conductor" && ev.To == spawnedID {
hasConvToSub = true hasConvToSub = true
break break
} }
} }
if !hasConvToSub { if !hasConvToSub {
t.Errorf("missing conversation(conductor → %s); events: %+v", spawnedID, events) t.Errorf("missing conversation(conductor → %s); events: %+v", spawnedID, events)
} }
// 3. There must be at least one agent_state(waiting) for the subagent. // 3. There must be at least one agent_state(waiting) for the subagent.
var hasWaiting bool var hasWaiting bool
for _, ev := range events { for _, ev := range events {
if ev.Type == "agent_state" && ev.ID == spawnedID && ev.State == "waiting" { if ev.Type == "agent_state" && ev.ID == spawnedID && ev.State == "waiting" {
hasWaiting = true hasWaiting = true
break break
} }
} }
if !hasWaiting { if !hasWaiting {
t.Errorf("missing agent_state(waiting) for %s; events: %+v", spawnedID, events) t.Errorf("missing agent_state(waiting) for %s; events: %+v", spawnedID, events)
} }
// 4. Last event must be agent_gc with reason "completed". // 4. Last event must be agent_gc with reason "completed".
last := events[len(events)-1] last := events[len(events)-1]
if last.Type != "agent_gc" || last.ID != spawnedID || last.Reason != "completed" { if last.Type != "agent_gc" || last.ID != spawnedID || last.Reason != "completed" {
t.Errorf("last event must be agent_gc(completed), got: %+v", last) t.Errorf("last event must be agent_gc(completed), got: %+v", last)
} }
// 5. Snapshot must be empty after GC (agent removed from live map). // 5. Snapshot must be empty after GC (agent removed from live map).
if snap := b.Snapshot(); len(snap) != 0 { if snap := b.Snapshot(); len(snap) != 0 {
t.Errorf("snapshot must be empty after agent_gc, got: %v", snap) t.Errorf("snapshot must be empty after agent_gc, got: %v", snap)
} }
// 6. Callback must be called. The callback fires in the same goroutine // 6. Callback must be called. The callback fires in the same goroutine
// as ReportGC (after the deferred unlock), so we poll briefly. // as ReportGC (after the deferred unlock), so we poll briefly.
for i := 0; i < 100; i++ { for i := 0; i < 100; i++ {
if atomic.LoadInt32(&callbackCalled) == 1 { if atomic.LoadInt32(&callbackCalled) == 1 {
break break
} }
time.Sleep(10 * time.Millisecond) time.Sleep(10 * time.Millisecond)
} }
if atomic.LoadInt32(&callbackCalled) != 1 { if atomic.LoadInt32(&callbackCalled) != 1 {
t.Error("completion callback was not called after agent_gc") t.Error("completion callback was not called after agent_gc")
} }
} }
// TestSubagentManager_Spawn_SnapshotLiveDuringExecution verifies that the // TestSubagentManager_Spawn_SnapshotLiveDuringExecution verifies that the
// Broadcaster snapshot contains the agent between agent_spawn and agent_gc. // Broadcaster snapshot contains the agent between agent_spawn and agent_gc.
// Because Publish() updates the agent map before dispatching to subscribers, // Because Publish() updates the agent map before dispatching to subscribers,
// the snapshot is guaranteed to be non-empty as soon as agent_spawn is // the snapshot is guaranteed to be non-empty as soon as agent_spawn is
// received on the channel. // received on the channel.
func TestSubagentManager_Spawn_SnapshotLiveDuringExecution(t *testing.T) { func TestSubagentManager_Spawn_SnapshotLiveDuringExecution(t *testing.T) {
b := orch.NewBroadcaster() b := orch.NewBroadcaster()
sub := b.Subscribe() sub := b.Subscribe()
defer b.Unsubscribe(sub) defer b.Unsubscribe(sub)
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
mgr := NewSubagentManager(provider, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) mgr := NewSubagentManager(provider, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{})
_, err := mgr.Spawn( _, err := mgr.Spawn(
context.Background(), context.Background(),
"any task", "live-test", "", "cli", "direct", "", "any task", "live-test", "", "cli", "direct", "",
nil, nil,
) )
if err != nil { if err != nil {
@ -161,39 +222,59 @@ func TestSubagentManager_Spawn_SnapshotLiveDuringExecution(t *testing.T) {
} }
// Wait for agent_spawn, then immediately check snapshot. // Wait for agent_spawn, then immediately check snapshot.
deadline := time.After(2 * time.Second) deadline := time.After(2 * time.Second)
for { for {
select { select {
case ev := <-sub.Ch: case ev := <-sub.Ch:
if ev.Type == "agent_spawn" { if ev.Type == "agent_spawn" {
snap := b.Snapshot() snap := b.Snapshot()
if len(snap) == 0 { if len(snap) == 0 {
t.Error("snapshot must contain the spawned agent after agent_spawn event") t.Error("snapshot must contain the spawned agent after agent_spawn event")
} }
return // test complete; background goroutine drains safely return // test complete; background goroutine drains safely
} }
case <-deadline: case <-deadline:
t.Fatal("timed out waiting for agent_spawn event") t.Fatal("timed out waiting for agent_spawn event")
} }
} }
} }
// TestSubagentManager_Spawn_CancelledDuringExecution verifies that when the // TestSubagentManager_Spawn_CancelledDuringExecution verifies that when the
// context is canceled while a subagent's LLM call is in progress, 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 // Broadcaster receives agent_gc with reason="canceled" and the agent is
// removed from the snapshot. // removed from the snapshot.
// //
// Synchronization: // Synchronization:
// 1. blockingProvider.ready is closed when Chat() is entered (goroutine is // 1. blockingProvider.ready is closed when Chat() is entered (goroutine is
// now blocked inside the LLM call). // now blocked inside the LLM call).
// 2. Only then is the context canceled, so there is no race between spawn // 2. Only then is the context canceled, so there is no race between spawn
// and cancellation. // and cancellation.
func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) { func TestSubagentManager_Spawn_CancelledDuringExecution(t *testing.T) {
b := orch.NewBroadcaster() b := orch.NewBroadcaster()
sub := b.Subscribe() sub := b.Subscribe()
defer b.Unsubscribe(sub) defer b.Unsubscribe(sub)
bp := newBlockingProvider() bp := newBlockingProvider()
mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{}) mgr := NewSubagentManager(bp, "test-model", "/tmp/test", nil, b, WebSearchToolOptions{})
_, err := mgr.Spawn(context.Background(), "long task", "cancel-me", "", "cli", "direct", "", nil) _, 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). // Wait until the subagent goroutine is inside Chat (blocking on ctx).
select { select {
case <-bp.ready: case <-bp.ready:
case <-time.After(3 * time.Second): case <-time.After(3 * time.Second):
t.Fatal("timed out waiting for blockingProvider to enter Chat") t.Fatal("timed out waiting for blockingProvider to enter Chat")
} }
// Cancel via CancelTask — the spawned goroutine's detached context is canceled. // Cancel via CancelTask — the spawned goroutine's detached context is canceled.
mgr.CancelTask("subagent-1") mgr.CancelTask("subagent-1")
// Collect events until agent_gc. // Collect events until agent_gc.
var events []orch.Event var events []orch.Event
deadline := time.After(3 * time.Second) deadline := time.After(3 * time.Second)
loop: loop:
for { for {
select { select {
case ev := <-sub.Ch: case ev := <-sub.Ch:
events = append(events, ev) events = append(events, ev)
if ev.Type == "agent_gc" { if ev.Type == "agent_gc" {
break loop break loop
} }
case <-deadline: case <-deadline:
t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events) t.Fatalf("timed out waiting for agent_gc; events so far: %+v", events)
} }
} }
// Locate agent_gc and verify reason = "canceled". // Locate agent_gc and verify reason = "canceled".
var gcEv orch.Event var gcEv orch.Event
for _, ev := range events { for _, ev := range events {
if ev.Type == "agent_gc" { if ev.Type == "agent_gc" {
gcEv = ev gcEv = ev
break break
} }
} }
if gcEv.Reason != "canceled" { if gcEv.Reason != "canceled" {
t.Errorf("agent_gc reason = %q, want %q; events: %+v", gcEv.Reason, "canceled", events) t.Errorf("agent_gc reason = %q, want %q; events: %+v", gcEv.Reason, "canceled", events)
} }
// Snapshot must be empty after the GC event. // Snapshot must be empty after the GC event.
if snap := b.Snapshot(); len(snap) != 0 { if snap := b.Snapshot(); len(snap) != 0 {
t.Errorf("snapshot must be empty after agent_gc(canceled), got: %v", snap) 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 // MockLLMProvider is a test implementation of LLMProvider
type MockLLMProvider struct { type MockLLMProvider struct {
lastOptions map[string]any lastOptions map[string]any
} }
func (m *MockLLMProvider) Chat( func (m *MockLLMProvider) Chat(
ctx context.Context, ctx context.Context,
messages []providers.Message, messages []providers.Message,
tools []providers.ToolDefinition, tools []providers.ToolDefinition,
model string, model string,
options map[string]any, options map[string]any,
) (*providers.LLMResponse, error) { ) (*providers.LLMResponse, error) {
m.lastOptions = options m.lastOptions = options
// Find the last user message to generate a response // Find the last user message to generate a response
for i := len(messages) - 1; i >= 0; i-- { for i := len(messages) - 1; i >= 0; i-- {
if messages[i].Role == "user" { if messages[i].Role == "user" {
return &providers.LLMResponse{ return &providers.LLMResponse{
@ -32,6 +39,7 @@ func (m *MockLLMProvider) Chat(
}, nil }, nil
} }
} }
return &providers.LLMResponse{Content: "No task provided"}, 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) { func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{})
manager.SetLLMOptions(2048, 0.6) manager.SetLLMOptions(2048, 0.6)
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
tool.SetContext("cli", "direct") tool.SetContext("cli", "direct")
ctx := context.Background() ctx := context.Background()
args := map[string]any{"task": "Do something"} args := map[string]any{"task": "Do something"}
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
if result == nil || result.IsError { if result == nil || result.IsError {
@ -65,18 +79,23 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
if provider.lastOptions == nil { if provider.lastOptions == nil {
t.Fatal("Expected LLM options to be passed, got nil") t.Fatal("Expected LLM options to be passed, got nil")
} }
if provider.lastOptions["max_tokens"] != 2048 { if provider.lastOptions["max_tokens"] != 2048 {
t.Fatalf("max_tokens = %v, want %d", provider.lastOptions["max_tokens"], 2048) t.Fatalf("max_tokens = %v, want %d", provider.lastOptions["max_tokens"], 2048)
} }
if provider.lastOptions["temperature"] != 0.6 { if provider.lastOptions["temperature"] != 0.6 {
t.Fatalf("temperature = %v, want %v", provider.lastOptions["temperature"], 0.6) t.Fatalf("temperature = %v, want %v", provider.lastOptions["temperature"], 0.6)
} }
} }
// TestSubagentTool_Name verifies tool name // TestSubagentTool_Name verifies tool name
func TestSubagentTool_Name(t *testing.T) { func TestSubagentTool_Name(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
if tool.Name() != "subagent" { if tool.Name() != "subagent" {
@ -85,150 +104,198 @@ func TestSubagentTool_Name(t *testing.T) {
} }
// TestSubagentTool_Description verifies tool description // TestSubagentTool_Description verifies tool description
func TestSubagentTool_Description(t *testing.T) { func TestSubagentTool_Description(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
desc := tool.Description() desc := tool.Description()
if desc == "" { if desc == "" {
t.Error("Description should not be empty") t.Error("Description should not be empty")
} }
if !strings.Contains(desc, "BLOCK") { if !strings.Contains(desc, "BLOCK") {
t.Errorf("Description should mention 'BLOCK', got: %s", desc) t.Errorf("Description should mention 'BLOCK', got: %s", desc)
} }
if !strings.Contains(desc, "spawn") { if !strings.Contains(desc, "spawn") {
t.Errorf("Description should contrast with spawn, got: %s", desc) t.Errorf("Description should contrast with spawn, got: %s", desc)
} }
} }
// TestSubagentTool_Parameters verifies tool parameters schema // TestSubagentTool_Parameters verifies tool parameters schema
func TestSubagentTool_Parameters(t *testing.T) { func TestSubagentTool_Parameters(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
params := tool.Parameters() params := tool.Parameters()
if params == nil { if params == nil {
t.Error("Parameters should not be nil") t.Error("Parameters should not be nil")
} }
// Check type // Check type
if params["type"] != "object" { if params["type"] != "object" {
t.Errorf("Expected type 'object', got: %v", params["type"]) t.Errorf("Expected type 'object', got: %v", params["type"])
} }
// Check properties // Check properties
props, ok := params["properties"].(map[string]any) props, ok := params["properties"].(map[string]any)
if !ok { if !ok {
t.Fatal("Properties should be a map") t.Fatal("Properties should be a map")
} }
// Verify task parameter // Verify task parameter
task, ok := props["task"].(map[string]any) task, ok := props["task"].(map[string]any)
if !ok { if !ok {
t.Fatal("Task parameter should exist") t.Fatal("Task parameter should exist")
} }
if task["type"] != "string" { if task["type"] != "string" {
t.Errorf("Task type should be 'string', got: %v", task["type"]) t.Errorf("Task type should be 'string', got: %v", task["type"])
} }
// Verify label parameter // Verify label parameter
label, ok := props["label"].(map[string]any) label, ok := props["label"].(map[string]any)
if !ok { if !ok {
t.Fatal("Label parameter should exist") t.Fatal("Label parameter should exist")
} }
if label["type"] != "string" { if label["type"] != "string" {
t.Errorf("Label type should be 'string', got: %v", label["type"]) t.Errorf("Label type should be 'string', got: %v", label["type"])
} }
// Check required fields // Check required fields
required, ok := params["required"].([]string) required, ok := params["required"].([]string)
if !ok { if !ok {
t.Fatal("Required should be a string array") t.Fatal("Required should be a string array")
} }
if len(required) != 1 || required[0] != "task" { if len(required) != 1 || required[0] != "task" {
t.Errorf("Required should be ['task'], got: %v", required) t.Errorf("Required should be ['task'], got: %v", required)
} }
} }
// TestSubagentTool_SetContext verifies context setting // TestSubagentTool_SetContext verifies context setting
func TestSubagentTool_SetContext(t *testing.T) { func TestSubagentTool_SetContext(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
tool.SetContext("test-channel", "test-chat") tool.SetContext("test-channel", "test-chat")
// Verify context is set (we can't directly access private fields, // Verify context is set (we can't directly access private fields,
// but we can verify it doesn't crash) // but we can verify it doesn't crash)
// The actual context usage is tested in Execute tests // The actual context usage is tested in Execute tests
} }
// TestSubagentTool_Execute_Success tests successful execution // TestSubagentTool_Execute_Success tests successful execution
func TestSubagentTool_Execute_Success(t *testing.T) { func TestSubagentTool_Execute_Success(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
tool.SetContext("telegram", "chat-123") tool.SetContext("telegram", "chat-123")
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"task": "Write a haiku about coding", "task": "Write a haiku about coding",
"label": "haiku-task", "label": "haiku-task",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Verify basic ToolResult structure // Verify basic ToolResult structure
if result == nil { if result == nil {
t.Fatal("Result should not be nil") t.Fatal("Result should not be nil")
} }
// Verify no error // Verify no error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got error: %s", result.ForLLM) t.Errorf("Expected success, got error: %s", result.ForLLM)
} }
// Verify not async // Verify not async
if result.Async { if result.Async {
t.Error("SubagentTool should be synchronous, not async") t.Error("SubagentTool should be synchronous, not async")
} }
// Verify not silent // Verify not silent
if result.Silent { if result.Silent {
t.Error("SubagentTool should not be silent") t.Error("SubagentTool should not be silent")
} }
// Verify ForUser contains brief summary (not empty) // Verify ForUser contains brief summary (not empty)
if result.ForUser == "" { if result.ForUser == "" {
t.Error("ForUser should contain result summary") t.Error("ForUser should contain result summary")
} }
if !strings.Contains(result.ForUser, "Task completed") { if !strings.Contains(result.ForUser, "Task completed") {
t.Errorf("ForUser should contain task completion, got: %s", result.ForUser) t.Errorf("ForUser should contain task completion, got: %s", result.ForUser)
} }
// Verify ForLLM contains full details // Verify ForLLM contains full details
if result.ForLLM == "" { if result.ForLLM == "" {
t.Error("ForLLM should contain full details") t.Error("ForLLM should contain full details")
} }
if !strings.Contains(result.ForLLM, "haiku-task") { if !strings.Contains(result.ForLLM, "haiku-task") {
t.Errorf("ForLLM should contain label 'haiku-task', got: %s", result.ForLLM) t.Errorf("ForLLM should contain label 'haiku-task', got: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "Task completed:") { if !strings.Contains(result.ForLLM, "Task completed:") {
t.Errorf("ForLLM should contain task result, got: %s", result.ForLLM) t.Errorf("ForLLM should contain task result, got: %s", result.ForLLM)
} }
} }
// TestSubagentTool_Execute_NoLabel tests execution without label // TestSubagentTool_Execute_NoLabel tests execution without label
func TestSubagentTool_Execute_NoLabel(t *testing.T) { func TestSubagentTool_Execute_NoLabel(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"task": "Test task without label", "task": "Test task without label",
} }
@ -240,18 +307,23 @@ func TestSubagentTool_Execute_NoLabel(t *testing.T) {
} }
// ForLLM should show (unnamed) for missing label // ForLLM should show (unnamed) for missing label
if !strings.Contains(result.ForLLM, "(unnamed)") { if !strings.Contains(result.ForLLM, "(unnamed)") {
t.Errorf("ForLLM should show '(unnamed)' for missing label, got: %s", result.ForLLM) t.Errorf("ForLLM should show '(unnamed)' for missing label, got: %s", result.ForLLM)
} }
} }
// TestSubagentTool_Execute_MissingTask tests error handling for missing task // TestSubagentTool_Execute_MissingTask tests error handling for missing task
func TestSubagentTool_Execute_MissingTask(t *testing.T) { func TestSubagentTool_Execute_MissingTask(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"label": "test", "label": "test",
} }
@ -259,29 +331,35 @@ func TestSubagentTool_Execute_MissingTask(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error // Should return error
if !result.IsError { if !result.IsError {
t.Error("Expected error for missing task parameter") t.Error("Expected error for missing task parameter")
} }
// ForLLM should contain helpful error with example // ForLLM should contain helpful error with example
if !strings.Contains(result.ForLLM, `"task"`) { if !strings.Contains(result.ForLLM, `"task"`) {
t.Errorf("Error message should mention '\"task\"', got: %s", result.ForLLM) t.Errorf("Error message should mention '\"task\"', got: %s", result.ForLLM)
} }
if !strings.Contains(result.ForLLM, "Example") { if !strings.Contains(result.ForLLM, "Example") {
t.Errorf("Error message should include usage example, got: %s", result.ForLLM) t.Errorf("Error message should include usage example, got: %s", result.ForLLM)
} }
// Err should be set // Err should be set
if result.Err == nil { if result.Err == nil {
t.Error("Err should be set for validation failure") t.Error("Err should be set for validation failure")
} }
} }
// TestSubagentTool_Execute_NilManager tests error handling for nil manager // TestSubagentTool_Execute_NilManager tests error handling for nil manager
func TestSubagentTool_Execute_NilManager(t *testing.T) { func TestSubagentTool_Execute_NilManager(t *testing.T) {
tool := NewSubagentTool(nil) tool := NewSubagentTool(nil)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"task": "test task", "task": "test task",
} }
@ -289,6 +367,7 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error // Should return error
if !result.IsError { if !result.IsError {
t.Error("Expected error for nil manager") 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 // TestSubagentTool_Execute_ContextPassing verifies context is properly used
func TestSubagentTool_Execute_ContextPassing(t *testing.T) { func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
// Set context // Set context
channel := "test-channel" channel := "test-channel"
chatID := "test-chat" chatID := "test-chat"
tool.SetContext(channel, chatID) tool.SetContext(channel, chatID)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"task": "Test context passing", "task": "Test context passing",
} }
@ -318,40 +405,53 @@ func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should succeed // Should succeed
if result.IsError { if result.IsError {
t.Errorf("Expected success with context, got error: %s", result.ForLLM) t.Errorf("Expected success with context, got error: %s", result.ForLLM)
} }
// The context is used internally; we can't directly test it // The context is used internally; we can't directly test it
// but execution success indicates context was handled properly // but execution success indicates context was handled properly
} }
// TestSubagentTool_ForUserTruncation verifies long content is truncated for user // TestSubagentTool_ForUserTruncation verifies long content is truncated for user
func TestSubagentTool_ForUserTruncation(t *testing.T) { func TestSubagentTool_ForUserTruncation(t *testing.T) {
// Create a mock provider that returns very long content // Create a mock provider that returns very long content
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{})
tool := NewSubagentTool(manager) tool := NewSubagentTool(manager)
ctx := context.Background() ctx := context.Background()
// Create a task that will generate long response // Create a task that will generate long response
longTask := strings.Repeat("This is a very long task description. ", 100) longTask := strings.Repeat("This is a very long task description. ", 100)
args := map[string]any{ args := map[string]any{
"task": longTask, "task": longTask,
"label": "long-test", "label": "long-test",
} }
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// ForUser should be truncated to 500 chars + "..." // ForUser should be truncated to 500 chars + "..."
maxUserLen := 500 maxUserLen := 500
if len(result.ForUser) > maxUserLen+3 { // +3 for "..." if len(result.ForUser) > maxUserLen+3 { // +3 for "..."
t.Errorf("ForUser should be truncated to ~%d chars, got: %d", maxUserLen, len(result.ForUser)) t.Errorf("ForUser should be truncated to ~%d chars, got: %d", maxUserLen, len(result.ForUser))
} }
// ForLLM should have full content // ForLLM should have full content
if !strings.Contains(result.ForLLM, longTask[:50]) { if !strings.Contains(result.ForLLM, longTask[:50]) {
t.Error("ForLLM should contain reference to original task") t.Error("ForLLM should contain reference to original task")
} }
@ -360,20 +460,28 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) {
func TestFormatToolStats(t *testing.T) { func TestFormatToolStats(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
stats map[string]int stats map[string]int
want string want string
}{ }{
{"empty", map[string]int{}, ""}, {"empty", map[string]int{}, ""},
{"single", map[string]int{"exec": 3}, "exec:3"}, {"single", map[string]int{"exec": 3}, "exec:3"},
{ {
"multiple sorted", "multiple sorted",
map[string]int{"read_file": 5, "exec": 3, "write_file": 1}, map[string]int{"read_file": 5, "exec": 3, "write_file": 1},
"exec:3,read_file:5,write_file:1", "exec:3,read_file:5,write_file:1",
}, },
} }
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
got := formatToolStats(tt.stats) got := formatToolStats(tt.stats)
if got != tt.want { if got != tt.want {
t.Errorf("formatToolStats(%v) = %q, want %q", tt.stats, 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 // TestSubagentManager_Spawn_SetsMetadata verifies that the bus message from a
// completed spawn includes execution statistics in Metadata. // completed spawn includes execution statistics in Metadata.
func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) { func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
mgr := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) mgr := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{})
_, err := mgr.Spawn( _, err := mgr.Spawn(
context.Background(), context.Background(),
"say hello", "meta-test", "", "cli", "direct", "", "say hello", "meta-test", "", "cli", "direct", "",
nil, nil,
) )
if err != nil { if err != nil {
@ -398,9 +513,13 @@ func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) {
} }
// Consume the inbound message from the bus // Consume the inbound message from the bus
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel() defer cancel()
received, ok := msgBus.ConsumeInbound(ctx) received, ok := msgBus.ConsumeInbound(ctx)
if !ok { if !ok {
t.Fatal("timed out waiting for bus message") t.Fatal("timed out waiting for bus message")
} }
@ -408,16 +527,21 @@ func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) {
if received.Channel != "system" { if received.Channel != "system" {
t.Fatalf("expected channel 'system', got %q", received.Channel) t.Fatalf("expected channel 'system', got %q", received.Channel)
} }
if received.Metadata == nil { if received.Metadata == nil {
t.Fatal("Metadata should not be nil") t.Fatal("Metadata should not be nil")
} }
if received.Metadata["iterations"] != "1" { if received.Metadata["iterations"] != "1" {
t.Errorf("iterations = %q, want %q", received.Metadata["iterations"], "1") t.Errorf("iterations = %q, want %q", received.Metadata["iterations"], "1")
} }
if received.Metadata["tool_calls"] != "0" { if received.Metadata["tool_calls"] != "0" {
t.Errorf("tool_calls = %q, want %q", 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 // duration_ms should be a non-negative number
if received.Metadata["duration_ms"] == "" { if received.Metadata["duration_ms"] == "" {
t.Error("duration_ms should be present") t.Error("duration_ms should be present")
} }

View file

@ -1,7 +1,11 @@
// PicoClaw - Ultra-lightweight personal AI agent // PicoClaw - Ultra-lightweight personal AI agent
// Inspired by and based on nanobot: https://github.com/HKUDS/nanobot // Inspired by and based on nanobot: https://github.com/HKUDS/nanobot
// License: MIT // License: MIT
// //
// Copyright (c) 2026 PicoClaw contributors // Copyright (c) 2026 PicoClaw contributors
package tools package tools
@ -18,140 +22,207 @@ import (
) )
// ToolLoopConfig configures the tool execution loop. // ToolLoopConfig configures the tool execution loop.
type ToolLoopConfig struct { type ToolLoopConfig struct {
Provider providers.LLMProvider Provider providers.LLMProvider
Model string Model string
Tools *ToolRegistry Tools *ToolRegistry
MaxIterations int MaxIterations int
LLMOptions map[string]any LLMOptions map[string]any
// Reporter and AgentID replace the old OnStateChange func. // Reporter and AgentID replace the old OnStateChange func.
// Reporter is called with ReportStateChange("waiting","") before each LLM // Reporter is called with ReportStateChange("waiting","") before each LLM
// call and ReportStateChange("toolcall", toolName) when each tool starts. // call and ReportStateChange("toolcall", toolName) when each tool starts.
// Pass nil or orch.Noop to disable. nil is treated as orch.Noop internally. // Pass nil or orch.Noop to disable. nil is treated as orch.Noop internally.
Reporter orch.AgentReporter Reporter orch.AgentReporter
AgentID string AgentID string
} }
// ToolLoopResult contains the result of running the tool loop. // ToolLoopResult contains the result of running the tool loop.
type ToolLoopResult struct { type ToolLoopResult struct {
Content string Content string
Iterations int Iterations int
ToolCalls int // total tool call count across all iterations ToolCalls int // total tool call count across all iterations
ToolStats map[string]int // tool name → call count ToolStats map[string]int // tool name → call count
} }
// RunToolLoop executes the LLM + tool call iteration loop. // RunToolLoop executes the LLM + tool call iteration loop.
// This is the core agent logic that can be reused by both main agent and subagents. // This is the core agent logic that can be reused by both main agent and subagents.
func RunToolLoop( func RunToolLoop(
ctx context.Context, ctx context.Context,
config ToolLoopConfig, config ToolLoopConfig,
messages []providers.Message, messages []providers.Message,
channel, chatID string, channel, chatID string,
) (*ToolLoopResult, error) { ) (*ToolLoopResult, error) {
reporter := config.Reporter reporter := config.Reporter
if reporter == nil { if reporter == nil {
reporter = orch.Noop reporter = orch.Noop
} }
iteration := 0 iteration := 0
totalToolCalls := 0 totalToolCalls := 0
toolStats := map[string]int{} toolStats := map[string]int{}
var finalContent string var finalContent string
for iteration < config.MaxIterations { for iteration < config.MaxIterations {
iteration++ iteration++
logger.DebugCF("toolloop", "LLM iteration", logger.DebugCF("toolloop", "LLM iteration",
map[string]any{ map[string]any{
"iteration": iteration, "iteration": iteration,
"max": config.MaxIterations, "max": config.MaxIterations,
}) })
// 1. Build tool definitions // 1. Build tool definitions
var providerToolDefs []providers.ToolDefinition var providerToolDefs []providers.ToolDefinition
if config.Tools != nil { if config.Tools != nil {
providerToolDefs = config.Tools.ToProviderDefs() providerToolDefs = config.Tools.ToProviderDefs()
} }
// 2. Set default LLM options // 2. Set default LLM options
llmOpts := config.LLMOptions llmOpts := config.LLMOptions
if llmOpts == nil { if llmOpts == nil {
llmOpts = map[string]any{} llmOpts = map[string]any{}
} }
// 3. Call LLM (hook: waiting for response) // 3. Call LLM (hook: waiting for response)
reporter.ReportStateChange(config.AgentID, orch.AgentStateWaiting, "") reporter.ReportStateChange(config.AgentID, orch.AgentStateWaiting, "")
response, err := config.Provider.Chat(ctx, messages, providerToolDefs, config.Model, llmOpts) response, err := config.Provider.Chat(ctx, messages, providerToolDefs, config.Model, llmOpts)
if err != nil { if err != nil {
logger.ErrorCF("toolloop", "LLM call failed", logger.ErrorCF("toolloop", "LLM call failed",
map[string]any{ map[string]any{
"iteration": iteration, "iteration": iteration,
"error": err.Error(), "error": err.Error(),
}) })
return nil, fmt.Errorf("LLM call failed: %w", err) return nil, fmt.Errorf("LLM call failed: %w", err)
} }
// 4. If no tool calls, we're done // 4. If no tool calls, we're done
if len(response.ToolCalls) == 0 { if len(response.ToolCalls) == 0 {
finalContent = response.Content finalContent = response.Content
logger.InfoCF("toolloop", "LLM response without tool calls (direct answer)", logger.InfoCF("toolloop", "LLM response without tool calls (direct answer)",
map[string]any{ map[string]any{
"iteration": iteration, "iteration": iteration,
"content_chars": len(finalContent), "content_chars": len(finalContent),
}) })
break break
} }
normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls)) normalizedToolCalls := make([]providers.ToolCall, 0, len(response.ToolCalls))
for _, tc := range response.ToolCalls { for _, tc := range response.ToolCalls {
normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc)) normalizedToolCalls = append(normalizedToolCalls, providers.NormalizeToolCall(tc))
} }
// 5. Log tool calls // 5. Log tool calls
toolNames := make([]string, 0, len(normalizedToolCalls)) toolNames := make([]string, 0, len(normalizedToolCalls))
for _, tc := range normalizedToolCalls { for _, tc := range normalizedToolCalls {
toolNames = append(toolNames, tc.Name) toolNames = append(toolNames, tc.Name)
} }
logger.InfoCF("toolloop", "LLM requested tool calls", logger.InfoCF("toolloop", "LLM requested tool calls",
map[string]any{ map[string]any{
"tools": toolNames, "tools": toolNames,
"count": len(normalizedToolCalls), "count": len(normalizedToolCalls),
"iteration": iteration, "iteration": iteration,
}) })
// 6. Build assistant message with tool calls // 6. Build assistant message with tool calls
assistantMsg := providers.Message{ assistantMsg := providers.Message{
Role: "assistant", Role: "assistant",
Content: response.Content, Content: response.Content,
} }
for _, tc := range normalizedToolCalls { for _, tc := range normalizedToolCalls {
assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, providers.ToolCall{ assistantMsg.ToolCalls = append(assistantMsg.ToolCalls, providers.ToolCall{
ID: tc.ID, ID: tc.ID,
Type: "function", Type: "function",
Name: tc.Name, Name: tc.Name,
Arguments: tc.Arguments, Arguments: tc.Arguments,
Function: &providers.FunctionCall{ Function: &providers.FunctionCall{
Name: tc.Name, Name: tc.Name,
Arguments: tc.Arguments, Arguments: tc.Arguments,
}, },
}) })
} }
messages = append(messages, assistantMsg) messages = append(messages, assistantMsg)
// 7. Execute tool calls (hook: toolcall per tool) // 7. Execute tool calls (hook: toolcall per tool)
for _, tc := range normalizedToolCalls { for _, tc := range normalizedToolCalls {
argsJSON, _ := json.Marshal(tc.Arguments) argsJSON, _ := json.Marshal(tc.Arguments)
argsPreview := utils.Truncate(string(argsJSON), 200) argsPreview := utils.Truncate(string(argsJSON), 200)
logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview),
map[string]any{ map[string]any{
"tool": tc.Name, "tool": tc.Name,
"iteration": iteration, "iteration": iteration,
}) })
reporter.ReportStateChange(config.AgentID, orch.AgentStateToolCall, tc.Name) reporter.ReportStateChange(config.AgentID, orch.AgentStateToolCall, tc.Name)
totalToolCalls++ totalToolCalls++
toolStats[tc.Name]++ toolStats[tc.Name]++
// Execute tool (no async callback for subagents - they run independently) // Execute tool (no async callback for subagents - they run independently)
var toolResult *ToolResult var toolResult *ToolResult
if config.Tools != nil { if config.Tools != nil {
toolResult = config.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, channel, chatID, nil) toolResult = config.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, channel, chatID, nil)
} else { } else {
@ -159,25 +230,34 @@ func RunToolLoop(
} }
// Determine content for LLM // Determine content for LLM
contentForLLM := toolResult.ForLLM contentForLLM := toolResult.ForLLM
if contentForLLM == "" && toolResult.Err != nil { if contentForLLM == "" && toolResult.Err != nil {
contentForLLM = toolResult.Err.Error() contentForLLM = toolResult.Err.Error()
} }
// Add tool result message // Add tool result message
toolResultMsg := providers.Message{ toolResultMsg := providers.Message{
Role: "tool", Role: "tool",
Content: contentForLLM, Content: contentForLLM,
ToolCallID: tc.ID, ToolCallID: tc.ID,
} }
messages = append(messages, toolResultMsg) messages = append(messages, toolResultMsg)
} }
} }
return &ToolLoopResult{ return &ToolLoopResult{
Content: finalContent, Content: finalContent,
Iterations: iteration, Iterations: iteration,
ToolCalls: totalToolCalls, ToolCalls: totalToolCalls,
ToolStats: toolStats, ToolStats: toolStats,
}, nil }, nil
} }

View file

@ -10,53 +10,78 @@ import (
) )
// reporterSpy records every ReportStateChange call in order. // reporterSpy records every ReportStateChange call in order.
// Spawn/Conversation/GC are not needed for toolloop tests. // Spawn/Conversation/GC are not needed for toolloop tests.
type reporterSpy struct { type reporterSpy struct {
mu sync.Mutex mu sync.Mutex
calls []spyCall calls []spyCall
} }
type spyCall struct { type spyCall struct {
state orch.AgentState 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) 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) { func (r *reporterSpy) ReportStateChange(id string, state orch.AgentState, tool string) {
r.mu.Lock() r.mu.Lock()
r.calls = append(r.calls, spyCall{state, tool}) r.calls = append(r.calls, spyCall{state, tool})
r.mu.Unlock() r.mu.Unlock()
} }
func (r *reporterSpy) snapshot() []spyCall { func (r *reporterSpy) snapshot() []spyCall {
r.mu.Lock() r.mu.Lock()
defer r.mu.Unlock() defer r.mu.Unlock()
out := make([]spyCall, len(r.calls)) out := make([]spyCall, len(r.calls))
copy(out, r.calls) copy(out, r.calls)
return out return out
} }
// sequenceMockProvider returns a tool call on the first Chat() call and a // sequenceMockProvider returns a tool call on the first Chat() call and a
// plain text response on all subsequent calls. Used to exercise the // plain text response on all subsequent calls. Used to exercise the
// waiting → toolcall → waiting event sequence in RunToolLoop. // waiting → toolcall → waiting event sequence in RunToolLoop.
type sequenceMockProvider struct { type sequenceMockProvider struct {
mu sync.Mutex mu sync.Mutex
callCount int callCount int
} }
func (m *sequenceMockProvider) Chat( func (m *sequenceMockProvider) Chat(
_ context.Context, _ context.Context,
_ []providers.Message, _ []providers.Message,
_ []providers.ToolDefinition, _ []providers.ToolDefinition,
_ string, _ string,
_ map[string]any, _ map[string]any,
) (*providers.LLMResponse, error) { ) (*providers.LLMResponse, error) {
m.mu.Lock() m.mu.Lock()
m.callCount++ m.callCount++
n := m.callCount n := m.callCount
m.mu.Unlock() m.mu.Unlock()
if n == 1 { if n == 1 {
return &providers.LLMResponse{ return &providers.LLMResponse{
ToolCalls: []providers.ToolCall{ ToolCalls: []providers.ToolCall{
@ -64,17 +89,24 @@ func (m *sequenceMockProvider) Chat(
}, },
}, nil }, nil
} }
return &providers.LLMResponse{Content: "done"}, nil return &providers.LLMResponse{Content: "done"}, nil
} }
func (m *sequenceMockProvider) GetDefaultModel() string { return "test" } func (m *sequenceMockProvider) GetDefaultModel() string { return "test" }
func (m *sequenceMockProvider) SupportsTools() bool { return true } func (m *sequenceMockProvider) SupportsTools() bool { return true }
func (m *sequenceMockProvider) GetContextWindow() int { return 4096 } func (m *sequenceMockProvider) GetContextWindow() int { return 4096 }
// echoTool is a minimal Tool stub registered as "echo_tool". // echoTool is a minimal Tool stub registered as "echo_tool".
type echoTool struct{} 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) Description() string { return "echo" }
func (t *echoTool) Parameters() map[string]any { func (t *echoTool) Parameters() map[string]any {
return map[string]any{"type": "object", "properties": 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 // TestToolLoop_NilReporter_FallsBackToNoop ensures that passing nil as
// Reporter does not panic — the loop must substitute orch.Noop internally. // Reporter does not panic — the loop must substitute orch.Noop internally.
func TestToolLoop_NilReporter_FallsBackToNoop(t *testing.T) { func TestToolLoop_NilReporter_FallsBackToNoop(t *testing.T) {
_, err := RunToolLoop(context.Background(), ToolLoopConfig{ _, err := RunToolLoop(context.Background(), ToolLoopConfig{
Provider: &MockLLMProvider{}, Provider: &MockLLMProvider{},
Model: "test", Model: "test",
MaxIterations: 1, MaxIterations: 1,
Reporter: nil, // must not panic Reporter: nil, // must not panic
}, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct") }, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct")
if err != nil { if err != nil {
t.Fatalf("unexpected error with nil reporter: %v", err) t.Fatalf("unexpected error with nil reporter: %v", err)
@ -98,50 +136,78 @@ func TestToolLoop_NilReporter_FallsBackToNoop(t *testing.T) {
} }
// TestToolLoop_Reporter_WaitingBeforeLLM verifies that ReportStateChange is // TestToolLoop_Reporter_WaitingBeforeLLM verifies that ReportStateChange is
// called with state="waiting" before the first LLM call. The mock provider // 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 // returns a direct text answer (no tool calls), so exactly one waiting event
// is expected. // is expected.
func TestToolLoop_Reporter_WaitingBeforeLLM(t *testing.T) { func TestToolLoop_Reporter_WaitingBeforeLLM(t *testing.T) {
rep := &reporterSpy{} rep := &reporterSpy{}
_, err := RunToolLoop(context.Background(), ToolLoopConfig{ _, err := RunToolLoop(context.Background(), ToolLoopConfig{
Provider: &MockLLMProvider{}, Provider: &MockLLMProvider{},
Model: "test", Model: "test",
MaxIterations: 1, MaxIterations: 1,
Reporter: rep, Reporter: rep,
AgentID: "sess-1", AgentID: "sess-1",
}, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct") }, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct")
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
} }
calls := rep.snapshot() calls := rep.snapshot()
if len(calls) == 0 { if len(calls) == 0 {
t.Fatal("expected at least one ReportStateChange call") t.Fatal("expected at least one ReportStateChange call")
} }
if calls[0].state != orch.AgentStateWaiting { if calls[0].state != orch.AgentStateWaiting {
t.Fatalf("first call must be state=waiting, got %+v", calls[0]) t.Fatalf("first call must be state=waiting, got %+v", calls[0])
} }
} }
// TestToolLoop_Reporter_ToolcallOrderedAfterWaiting verifies the canonical // TestToolLoop_Reporter_ToolcallOrderedAfterWaiting verifies the canonical
// two-iteration sequence: // two-iteration sequence:
// //
// waiting (before 1st LLM call) // waiting (before 1st LLM call)
// toolcall(echo_tool) (before tool execution) // toolcall(echo_tool) (before tool execution)
// waiting (before 2nd LLM call) // waiting (before 2nd LLM call)
// //
// The sequenceMockProvider returns a tool call on iteration 1 and a text // The sequenceMockProvider returns a tool call on iteration 1 and a text
// response on iteration 2, driving exactly this path. // response on iteration 2, driving exactly this path.
func TestToolLoop_Reporter_ToolcallOrderedAfterWaiting(t *testing.T) { func TestToolLoop_Reporter_ToolcallOrderedAfterWaiting(t *testing.T) {
rep := &reporterSpy{} rep := &reporterSpy{}
reg := NewToolRegistry() reg := NewToolRegistry()
reg.Register(&echoTool{}) reg.Register(&echoTool{})
_, err := RunToolLoop(context.Background(), ToolLoopConfig{ _, err := RunToolLoop(context.Background(), ToolLoopConfig{
Provider: &sequenceMockProvider{}, Provider: &sequenceMockProvider{},
Model: "test", Model: "test",
Tools: reg, Tools: reg,
MaxIterations: 5, MaxIterations: 5,
Reporter: rep, Reporter: rep,
AgentID: "sess-1", AgentID: "sess-1",
}, []providers.Message{{Role: "user", Content: "do it"}}, "cli", "direct") }, []providers.Message{{Role: "user", Content: "do it"}}, "cli", "direct")
if err != nil { if err != nil {
@ -149,69 +215,92 @@ func TestToolLoop_Reporter_ToolcallOrderedAfterWaiting(t *testing.T) {
} }
calls := rep.snapshot() calls := rep.snapshot()
if len(calls) < 3 { if len(calls) < 3 {
t.Fatalf("expected at least 3 calls, got %d: %+v", len(calls), calls) t.Fatalf("expected at least 3 calls, got %d: %+v", len(calls), calls)
} }
if calls[0].state != orch.AgentStateWaiting { if calls[0].state != orch.AgentStateWaiting {
t.Fatalf("calls[0] must be waiting, got %+v", calls[0]) t.Fatalf("calls[0] must be waiting, got %+v", calls[0])
} }
if calls[1].state != orch.AgentStateToolCall || calls[1].tool != "echo_tool" { if calls[1].state != orch.AgentStateToolCall || calls[1].tool != "echo_tool" {
t.Fatalf("calls[1] must be toolcall(echo_tool), got %+v", calls[1]) t.Fatalf("calls[1] must be toolcall(echo_tool), got %+v", calls[1])
} }
if calls[2].state != orch.AgentStateWaiting { if calls[2].state != orch.AgentStateWaiting {
t.Fatalf("calls[2] must be waiting (2nd LLM iteration), got %+v", calls[2]) t.Fatalf("calls[2] must be waiting (2nd LLM iteration), got %+v", calls[2])
} }
} }
// TestToolLoop_ToolCallStats verifies that ToolLoopResult.ToolCalls and // TestToolLoop_ToolCallStats verifies that ToolLoopResult.ToolCalls and
// ToolStats are populated correctly after a tool call iteration. // ToolStats are populated correctly after a tool call iteration.
func TestToolLoop_ToolCallStats(t *testing.T) { func TestToolLoop_ToolCallStats(t *testing.T) {
reg := NewToolRegistry() reg := NewToolRegistry()
reg.Register(&echoTool{}) reg.Register(&echoTool{})
result, err := RunToolLoop(context.Background(), ToolLoopConfig{ result, err := RunToolLoop(context.Background(), ToolLoopConfig{
Provider: &sequenceMockProvider{}, Provider: &sequenceMockProvider{},
Model: "test", Model: "test",
Tools: reg, Tools: reg,
MaxIterations: 5, MaxIterations: 5,
}, []providers.Message{{Role: "user", Content: "do it"}}, "cli", "direct") }, []providers.Message{{Role: "user", Content: "do it"}}, "cli", "direct")
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
} }
if result.ToolCalls != 1 { if result.ToolCalls != 1 {
t.Errorf("ToolCalls = %d, want 1", result.ToolCalls) t.Errorf("ToolCalls = %d, want 1", result.ToolCalls)
} }
if result.ToolStats["echo_tool"] != 1 { if result.ToolStats["echo_tool"] != 1 {
t.Errorf("ToolStats[echo_tool] = %d, want 1", result.ToolStats["echo_tool"]) t.Errorf("ToolStats[echo_tool] = %d, want 1", result.ToolStats["echo_tool"])
} }
if result.Iterations != 2 { if result.Iterations != 2 {
t.Errorf("Iterations = %d, want 2", result.Iterations) t.Errorf("Iterations = %d, want 2", result.Iterations)
} }
} }
// TestToolLoop_NoToolCalls_ZeroStats verifies that a direct answer (no tool // TestToolLoop_NoToolCalls_ZeroStats verifies that a direct answer (no tool
// calls) produces zero ToolCalls and an empty ToolStats map. // calls) produces zero ToolCalls and an empty ToolStats map.
func TestToolLoop_NoToolCalls_ZeroStats(t *testing.T) { func TestToolLoop_NoToolCalls_ZeroStats(t *testing.T) {
result, err := RunToolLoop(context.Background(), ToolLoopConfig{ result, err := RunToolLoop(context.Background(), ToolLoopConfig{
Provider: &MockLLMProvider{}, Provider: &MockLLMProvider{},
Model: "test", Model: "test",
MaxIterations: 1, MaxIterations: 1,
}, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct") }, []providers.Message{{Role: "user", Content: "hi"}}, "cli", "direct")
if err != nil { if err != nil {
t.Fatalf("unexpected error: %v", err) t.Fatalf("unexpected error: %v", err)
} }
if result.ToolCalls != 0 { if result.ToolCalls != 0 {
t.Errorf("ToolCalls = %d, want 0", result.ToolCalls) t.Errorf("ToolCalls = %d, want 0", result.ToolCalls)
} }
if len(result.ToolStats) != 0 { if len(result.ToolStats) != 0 {
t.Errorf("ToolStats = %v, want empty", result.ToolStats) t.Errorf("ToolStats = %v, want empty", result.ToolStats)
} }
} }
// TestToolLoop_Reporter_NoopImplementsInterface is a compile-time check that // TestToolLoop_Reporter_NoopImplementsInterface is a compile-time check that
// orch.Noop satisfies the orch.AgentReporter interface accepted by // orch.Noop satisfies the orch.AgentReporter interface accepted by
// ToolLoopConfig.Reporter. If Noop ever stops implementing the interface the // ToolLoopConfig.Reporter. If Noop ever stops implementing the interface the
// build will fail here before any test runs. // build will fail here before any test runs.
func TestToolLoop_Reporter_NoopImplementsInterface(t *testing.T) { func TestToolLoop_Reporter_NoopImplementsInterface(t *testing.T) {
var _ orch.AgentReporter = orch.Noop 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" 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. // 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) perplexityTimeout = 30 * time.Second // Perplexity (LLM-based, slower)
fetchTimeout = 60 * time.Second // WebFetchTool fetchTimeout = 60 * time.Second // WebFetchTool
defaultMaxChars = 50000 defaultMaxChars = 50000
maxRedirects = 5 maxRedirects = 5
) )
// Pre-compiled regexes for HTML text extraction // Pre-compiled regexes for HTML text extraction
var ( var (
reScript = regexp.MustCompile(`<script[\s\S]*?</script>`) reScript = regexp.MustCompile(`<script[\s\S]*?</script>`)
reStyle = regexp.MustCompile(`<style[\s\S]*?</style>`) reStyle = regexp.MustCompile(`<style[\s\S]*?</style>`)
reTags = regexp.MustCompile(`<[^>]+>`) reTags = regexp.MustCompile(`<[^>]+>`)
reWhitespace = regexp.MustCompile(`[^\S\n]+`) reWhitespace = regexp.MustCompile(`[^\S\n]+`)
reBlankLines = regexp.MustCompile(`\n{3,}`) reBlankLines = regexp.MustCompile(`\n{3,}`)
// DuckDuckGo result extraction // 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>`) reDDGSnippet = regexp.MustCompile(`<a class="result__snippet[^"]*".*?>([\s\S]*?)</a>`)
) )
// createHTTPClient creates an HTTP client with optional proxy support // createHTTPClient creates an HTTP client with optional proxy support
func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, error) { func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, error) {
client := &http.Client{ client := &http.Client{
Timeout: timeout, Timeout: timeout,
Transport: &http.Transport{ Transport: &http.Transport{
MaxIdleConns: 10, MaxIdleConns: 10,
IdleConnTimeout: 30 * time.Second, IdleConnTimeout: 30 * time.Second,
DisableCompression: false, DisableCompression: false,
TLSHandshakeTimeout: 15 * time.Second, TLSHandshakeTimeout: 15 * time.Second,
}, },
} }
@ -55,18 +71,26 @@ func createHTTPClient(proxyURL string, timeout time.Duration) (*http.Client, err
if err != nil { if err != nil {
return nil, fmt.Errorf("invalid proxy URL: %w", err) return nil, fmt.Errorf("invalid proxy URL: %w", err)
} }
scheme := strings.ToLower(proxy.Scheme) scheme := strings.ToLower(proxy.Scheme)
switch scheme { switch scheme {
case "http", "https", "socks5", "socks5h": case "http", "https", "socks5", "socks5h":
default: default:
return nil, fmt.Errorf( return nil, fmt.Errorf(
"unsupported proxy scheme %q (supported: http, https, socks5, socks5h)", "unsupported proxy scheme %q (supported: http, https, socks5, socks5h)",
proxy.Scheme, proxy.Scheme,
) )
} }
if proxy.Host == "" { if proxy.Host == "" {
return nil, fmt.Errorf("invalid proxy URL: missing host") return nil, fmt.Errorf("invalid proxy URL: missing host")
} }
client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy) client.Transport.(*http.Transport).Proxy = http.ProxyURL(proxy)
} else { } else {
client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment client.Transport.(*http.Transport).Proxy = http.ProxyFromEnvironment
@ -81,7 +105,9 @@ type SearchProvider interface {
type searchResultItem struct { type searchResultItem struct {
Title string Title string
URL string URL string
Snippet string Snippet string
} }
@ -91,32 +117,41 @@ func formatWebSearchResults(query, provider string, results []searchResultItem,
} }
header := fmt.Sprintf("Results for: %s", query) header := fmt.Sprintf("Results for: %s", query)
if provider != "" { if provider != "" {
header += " (via " + provider + ")" header += " (via " + provider + ")"
} }
var sb strings.Builder var sb strings.Builder
sb.WriteString(header) sb.WriteString(header)
for i, item := range results { for i, item := range results {
if i >= count { if i >= count {
break break
} }
fmt.Fprintf(&sb, "\n%d. %s\n %s", i+1, item.Title, item.URL) fmt.Fprintf(&sb, "\n%d. %s\n %s", i+1, item.Title, item.URL)
if item.Snippet != "" { if item.Snippet != "" {
fmt.Fprintf(&sb, "\n %s", item.Snippet) fmt.Fprintf(&sb, "\n %s", item.Snippet)
} }
} }
return sb.String() return sb.String()
} }
type BraveSearchProvider struct { type BraveSearchProvider struct {
apiKey string apiKey string
proxy string proxy string
client *http.Client client *http.Client
} }
func (p *BraveSearchProvider) Search(ctx context.Context, query string, count int) (string, error) { 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", searchURL := fmt.Sprintf("https://api.search.brave.com/res/v1/web/search?q=%s&count=%d",
url.QueryEscape(query), count) url.QueryEscape(query), count)
req, err := http.NewRequestWithContext(ctx, "GET", searchURL, nil) 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("Accept", "application/json")
req.Header.Set("X-Subscription-Token", p.apiKey) req.Header.Set("X-Subscription-Token", p.apiKey)
resp, err := p.client.Do(req) resp, err := p.client.Do(req)
if err != nil { if err != nil {
return "", fmt.Errorf("request failed: %w", err) return "", fmt.Errorf("request failed: %w", err)
} }
defer resp.Body.Close() defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
@ -142,7 +179,9 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
Web struct { Web struct {
Results []struct { Results []struct {
Title string `json:"title"` Title string `json:"title"`
URL string `json:"url"` URL string `json:"url"`
Description string `json:"description"` Description string `json:"description"`
} `json:"results"` } `json:"results"`
} `json:"web"` } `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 { if err := json.Unmarshal(body, &searchResp); err != nil {
// Log error body for debugging // Log error body for debugging
fmt.Printf("Brave API Error Body: %s\n", string(body)) fmt.Printf("Brave API Error Body: %s\n", string(body))
return "", fmt.Errorf("failed to parse response: %w", err) return "", fmt.Errorf("failed to parse response: %w", err)
} }
results := searchResp.Web.Results results := searchResp.Web.Results
items := make([]searchResultItem, 0, len(results)) items := make([]searchResultItem, 0, len(results))
for _, item := range results { for _, item := range results {
items = append(items, searchResultItem{ items = append(items, searchResultItem{
Title: item.Title, Title: item.Title,
URL: item.URL, URL: item.URL,
Snippet: item.Description, Snippet: item.Description,
}) })
} }
@ -169,24 +214,34 @@ func (p *BraveSearchProvider) Search(ctx context.Context, query string, count in
type TavilySearchProvider struct { type TavilySearchProvider struct {
apiKey string apiKey string
baseURL string baseURL string
proxy string proxy string
client *http.Client client *http.Client
} }
func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) { func (p *TavilySearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
searchURL := p.baseURL searchURL := p.baseURL
if searchURL == "" { if searchURL == "" {
searchURL = "https://api.tavily.com/search" searchURL = "https://api.tavily.com/search"
} }
payload := map[string]any{ payload := map[string]any{
"api_key": p.apiKey, "api_key": p.apiKey,
"query": query, "query": query,
"search_depth": "advanced", "search_depth": "advanced",
"include_answer": false, "include_answer": false,
"include_images": false, "include_images": false,
"include_raw_content": false, "include_raw_content": false,
"max_results": count, "max_results": count,
} }
@ -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("Content-Type", "application/json")
req.Header.Set("User-Agent", userAgent) req.Header.Set("User-Agent", userAgent)
resp, err := p.client.Do(req) resp, err := p.client.Do(req)
if err != nil { if err != nil {
return "", fmt.Errorf("request failed: %w", err) return "", fmt.Errorf("request failed: %w", err)
} }
defer resp.Body.Close() defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
@ -221,7 +278,9 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
var searchResp struct { var searchResp struct {
Results []struct { Results []struct {
Title string `json:"title"` Title string `json:"title"`
URL string `json:"url"` URL string `json:"url"`
Content string `json:"content"` Content string `json:"content"`
} `json:"results"` } `json:"results"`
} }
@ -231,11 +290,15 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
} }
results := searchResp.Results results := searchResp.Results
items := make([]searchResultItem, 0, len(results)) items := make([]searchResultItem, 0, len(results))
for _, item := range results { for _, item := range results {
items = append(items, searchResultItem{ items = append(items, searchResultItem{
Title: item.Title, Title: item.Title,
URL: item.URL, URL: item.URL,
Snippet: item.Content, Snippet: item.Content,
}) })
} }
@ -245,6 +308,7 @@ func (p *TavilySearchProvider) Search(ctx context.Context, query string, count i
type DuckDuckGoSearchProvider struct { type DuckDuckGoSearchProvider struct {
proxy string proxy string
client *http.Client client *http.Client
} }
@ -262,6 +326,7 @@ func (p *DuckDuckGoSearchProvider) Search(ctx context.Context, query string, cou
if err != nil { if err != nil {
return "", fmt.Errorf("request failed: %w", err) return "", fmt.Errorf("request failed: %w", err)
} }
defer resp.Body.Close() defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) 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) { func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query string) (string, error) {
// Simple regex based extraction for DDG HTML // Simple regex based extraction for DDG HTML
// Strategy: Find all result containers or key anchors directly // Strategy: Find all result containers or key anchors directly
// Try finding the result links directly first, as they are the most critical // Try finding the result links directly first, as they are the most critical
// Pattern: <a class="result__a" href="...">Title</a> // 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 // The previous regex was a bit strict. Let's make it more flexible for attributes order/content
matches := reDDGLink.FindAllStringSubmatch(html, count+5) matches := reDDGLink.FindAllStringSubmatch(html, count+5)
if len(matches) == 0 { if len(matches) == 0 {
@ -288,17 +357,22 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query
snippetMatches := reDDGSnippet.FindAllStringSubmatch(html, count+5) snippetMatches := reDDGSnippet.FindAllStringSubmatch(html, count+5)
maxItems := min(len(matches), count) maxItems := min(len(matches), count)
items := make([]searchResultItem, 0, maxItems) items := make([]searchResultItem, 0, maxItems)
for i := range maxItems { for i := range maxItems {
urlStr := matches[i][1] urlStr := matches[i][1]
title := stripTags(matches[i][2]) title := stripTags(matches[i][2])
title = strings.TrimSpace(title) title = strings.TrimSpace(title)
// URL decoding if needed // URL decoding if needed
if strings.Contains(urlStr, "uddg=") { if strings.Contains(urlStr, "uddg=") {
if u, err := url.QueryUnescape(urlStr); err == nil { if u, err := url.QueryUnescape(urlStr); err == nil {
_, after, ok := strings.Cut(u, "uddg=") _, after, ok := strings.Cut(u, "uddg=")
if ok { if ok {
urlStr = after urlStr = after
} }
@ -306,15 +380,20 @@ func (p *DuckDuckGoSearchProvider) extractResults(html string, count int, query
} }
snippet := "" snippet := ""
// Attempt to attach snippet if available and index aligns // Attempt to attach snippet if available and index aligns
if i < len(snippetMatches) { if i < len(snippetMatches) {
snippet = stripTags(snippetMatches[i][1]) snippet = stripTags(snippetMatches[i][1])
snippet = strings.TrimSpace(snippet) snippet = strings.TrimSpace(snippet)
} }
items = append(items, searchResultItem{ items = append(items, searchResultItem{
Title: title, Title: title,
URL: urlStr, URL: urlStr,
Snippet: snippet, Snippet: snippet,
}) })
} }
@ -328,7 +407,9 @@ func stripTags(content string) string {
type PerplexitySearchProvider struct { type PerplexitySearchProvider struct {
apiKey string apiKey string
proxy string proxy string
client *http.Client client *http.Client
} }
@ -337,16 +418,21 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
payload := map[string]any{ payload := map[string]any{
"model": "sonar", "model": "sonar",
"messages": []map[string]string{ "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.", "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), "content": fmt.Sprintf("Search for: %s. Provide up to %d relevant results.", query, count),
}, },
}, },
"max_tokens": 1000, "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("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+p.apiKey) req.Header.Set("Authorization", "Bearer "+p.apiKey)
req.Header.Set("User-Agent", userAgent) req.Header.Set("User-Agent", userAgent)
resp, err := p.client.Do(req) resp, err := p.client.Do(req)
if err != nil { if err != nil {
return "", fmt.Errorf("request failed: %w", err) return "", fmt.Errorf("request failed: %w", err)
} }
defer resp.Body.Close() defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
@ -400,44 +489,65 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
type WebSearchTool struct { type WebSearchTool struct {
provider SearchProvider provider SearchProvider
providerName string providerName string
maxResults int maxResults int
} }
// ProviderName returns the name of the active search provider (e.g. "brave", "perplexity"). // ProviderName returns the name of the active search provider (e.g. "brave", "perplexity").
func (t *WebSearchTool) ProviderName() string { func (t *WebSearchTool) ProviderName() string {
return t.providerName return t.providerName
} }
type WebSearchToolOptions struct { type WebSearchToolOptions struct {
BraveAPIKey string BraveAPIKey string
BraveMaxResults int BraveMaxResults int
BraveEnabled bool BraveEnabled bool
TavilyAPIKey string TavilyAPIKey string
TavilyBaseURL string TavilyBaseURL string
TavilyMaxResults int TavilyMaxResults int
TavilyEnabled bool TavilyEnabled bool
DuckDuckGoMaxResults int DuckDuckGoMaxResults int
DuckDuckGoEnabled bool DuckDuckGoEnabled bool
PerplexityAPIKey string PerplexityAPIKey string
PerplexityMaxResults int PerplexityMaxResults int
PerplexityEnabled bool PerplexityEnabled bool
Proxy string Proxy string
} }
func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) { func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
var provider SearchProvider var provider SearchProvider
var providerName string var providerName string
maxResults := 5 maxResults := 5
// Priority: Perplexity > Brave > Tavily > DuckDuckGo // Priority: Perplexity > Brave > Tavily > DuckDuckGo
if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" { if opts.PerplexityEnabled && opts.PerplexityAPIKey != "" {
client, err := createHTTPClient(opts.Proxy, perplexityTimeout) client, err := createHTTPClient(opts.Proxy, perplexityTimeout)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err) return nil, fmt.Errorf("failed to create HTTP client for Perplexity: %w", err)
} }
provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy, client: client} provider = &PerplexitySearchProvider{apiKey: opts.PerplexityAPIKey, proxy: opts.Proxy, client: client}
providerName = "perplexity" providerName = "perplexity"
if opts.PerplexityMaxResults > 0 { if opts.PerplexityMaxResults > 0 {
maxResults = opts.PerplexityMaxResults maxResults = opts.PerplexityMaxResults
} }
@ -446,8 +556,11 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err) return nil, fmt.Errorf("failed to create HTTP client for Brave: %w", err)
} }
provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy, client: client} provider = &BraveSearchProvider{apiKey: opts.BraveAPIKey, proxy: opts.Proxy, client: client}
providerName = "brave" providerName = "brave"
if opts.BraveMaxResults > 0 { if opts.BraveMaxResults > 0 {
maxResults = opts.BraveMaxResults maxResults = opts.BraveMaxResults
} }
@ -456,13 +569,19 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err) return nil, fmt.Errorf("failed to create HTTP client for Tavily: %w", err)
} }
provider = &TavilySearchProvider{ provider = &TavilySearchProvider{
apiKey: opts.TavilyAPIKey, apiKey: opts.TavilyAPIKey,
baseURL: opts.TavilyBaseURL, baseURL: opts.TavilyBaseURL,
proxy: opts.Proxy, proxy: opts.Proxy,
client: client, client: client,
} }
providerName = "tavily" providerName = "tavily"
if opts.TavilyMaxResults > 0 { if opts.TavilyMaxResults > 0 {
maxResults = opts.TavilyMaxResults maxResults = opts.TavilyMaxResults
} }
@ -471,8 +590,11 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for DuckDuckGo: %w", err) return nil, fmt.Errorf("failed to create HTTP client for DuckDuckGo: %w", err)
} }
provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy, client: client} provider = &DuckDuckGoSearchProvider{proxy: opts.Proxy, client: client}
providerName = "duckduckgo" providerName = "duckduckgo"
if opts.DuckDuckGoMaxResults > 0 { if opts.DuckDuckGoMaxResults > 0 {
maxResults = opts.DuckDuckGoMaxResults maxResults = opts.DuckDuckGoMaxResults
} }
@ -482,7 +604,9 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
return &WebSearchTool{ return &WebSearchTool{
provider: provider, provider: provider,
providerName: providerName, providerName: providerName,
maxResults: maxResults, maxResults: maxResults,
}, nil }, nil
} }
@ -498,29 +622,38 @@ func (t *WebSearchTool) Description() string {
func (t *WebSearchTool) Parameters() map[string]any { func (t *WebSearchTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"query": map[string]any{ "query": map[string]any{
"type": "string", "type": "string",
"description": "Search query", "description": "Search query",
}, },
"count": map[string]any{ "count": map[string]any{
"type": "integer", "type": "integer",
"description": "Number of results (1-10)", "description": "Number of results (1-10)",
"minimum": 1.0, "minimum": 1.0,
"maximum": 10.0, "maximum": 10.0,
}, },
}, },
"required": []string{"query"}, "required": []string{"query"},
} }
} }
func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
query, ok := args["query"].(string) query, ok := args["query"].(string)
if !ok { if !ok {
return ErrorResult("query is required") return ErrorResult("query is required")
} }
count := t.maxResults count := t.maxResults
if c, ok := args["count"].(float64); ok { if c, ok := args["count"].(float64); ok {
if int(c) > 0 && int(c) <= 10 { if int(c) > 0 && int(c) <= 10 {
count = int(c) count = int(c)
@ -534,19 +667,24 @@ func (t *WebSearchTool) Execute(ctx context.Context, args map[string]any) *ToolR
return &ToolResult{ return &ToolResult{
ForLLM: result, ForLLM: result,
ForUser: result, ForUser: result,
} }
} }
type WebFetchTool struct { type WebFetchTool struct {
maxChars int maxChars int
proxy string proxy string
client *http.Client client *http.Client
} }
func NewWebFetchTool(maxChars int) *WebFetchTool { func NewWebFetchTool(maxChars int) *WebFetchTool {
// createHTTPClient cannot fail with an empty proxy string. // createHTTPClient cannot fail with an empty proxy string.
tool, _ := NewWebFetchToolWithProxy(maxChars, "") tool, _ := NewWebFetchToolWithProxy(maxChars, "")
return tool return tool
} }
@ -554,19 +692,25 @@ func NewWebFetchToolWithProxy(maxChars int, proxy string) (*WebFetchTool, error)
if maxChars <= 0 { if maxChars <= 0 {
maxChars = defaultMaxChars maxChars = defaultMaxChars
} }
client, err := createHTTPClient(proxy, fetchTimeout) client, err := createHTTPClient(proxy, fetchTimeout)
if err != nil { if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err) return nil, fmt.Errorf("failed to create HTTP client for web fetch: %w", err)
} }
client.CheckRedirect = func(req *http.Request, via []*http.Request) error { client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
if len(via) >= maxRedirects { if len(via) >= maxRedirects {
return fmt.Errorf("stopped after %d redirects", maxRedirects) return fmt.Errorf("stopped after %d redirects", maxRedirects)
} }
return nil return nil
} }
return &WebFetchTool{ return &WebFetchTool{
maxChars: maxChars, maxChars: maxChars,
proxy: proxy, proxy: proxy,
client: client, client: client,
}, nil }, nil
} }
@ -582,23 +726,30 @@ func (t *WebFetchTool) Description() string {
func (t *WebFetchTool) Parameters() map[string]any { func (t *WebFetchTool) Parameters() map[string]any {
return map[string]any{ return map[string]any{
"type": "object", "type": "object",
"properties": map[string]any{ "properties": map[string]any{
"url": map[string]any{ "url": map[string]any{
"type": "string", "type": "string",
"description": "URL to fetch", "description": "URL to fetch",
}, },
"maxChars": map[string]any{ "maxChars": map[string]any{
"type": "integer", "type": "integer",
"description": "Maximum characters to extract", "description": "Maximum characters to extract",
"minimum": 100.0, "minimum": 100.0,
}, },
}, },
"required": []string{"url"}, "required": []string{"url"},
} }
} }
func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolResult { func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
urlStr, ok := args["url"].(string) urlStr, ok := args["url"].(string)
if !ok { if !ok {
return ErrorResult("url is required") return ErrorResult("url is required")
} }
@ -617,6 +768,7 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe
} }
maxChars := t.maxChars maxChars := t.maxChars
if mc, ok := args["maxChars"].(float64); ok { if mc, ok := args["maxChars"].(float64); ok {
if int(mc) > 100 { if int(mc) > 100 {
maxChars = int(mc) maxChars = int(mc)
@ -634,6 +786,7 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe
if err != nil { if err != nil {
return ErrorResult(fmt.Sprintf("request failed: %v", err)) return ErrorResult(fmt.Sprintf("request failed: %v", err))
} }
defer resp.Body.Close() defer resp.Body.Close()
body, err := io.ReadAll(resp.Body) body, err := io.ReadAll(resp.Body)
@ -646,36 +799,50 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe
var text, extractor string var text, extractor string
bodyStr := string(body) bodyStr := string(body)
if strings.Contains(contentType, "application/json") { if strings.Contains(contentType, "application/json") {
var jsonData any var jsonData any
if err := json.Unmarshal(body, &jsonData); err == nil { if err := json.Unmarshal(body, &jsonData); err == nil {
formatted, _ := json.MarshalIndent(jsonData, "", " ") formatted, _ := json.MarshalIndent(jsonData, "", " ")
text = string(formatted) text = string(formatted)
extractor = "json" extractor = "json"
} else { } else {
text = bodyStr text = bodyStr
extractor = "raw" extractor = "raw"
} }
} else if strings.Contains(contentType, "text/html") || len(body) > 0 && } else if strings.Contains(contentType, "text/html") || len(body) > 0 &&
(strings.HasPrefix(bodyStr, "<!DOCTYPE") || strings.HasPrefix(strings.ToLower(bodyStr), "<html")) { (strings.HasPrefix(bodyStr, "<!DOCTYPE") || strings.HasPrefix(strings.ToLower(bodyStr), "<html")) {
text = t.extractText(bodyStr) text = t.extractText(bodyStr)
extractor = "text" extractor = "text"
} else { } else {
text = bodyStr text = bodyStr
extractor = "raw" extractor = "raw"
} }
truncated := len(text) > maxChars truncated := len(text) > maxChars
if truncated { if truncated {
text = text[:maxChars] text = text[:maxChars]
} }
result := map[string]any{ result := map[string]any{
"url": urlStr, "url": urlStr,
"status": resp.StatusCode, "status": resp.StatusCode,
"extractor": extractor, "extractor": extractor,
"truncated": truncated, "truncated": truncated,
"length": len(text), "length": len(text),
"text": text, "text": text,
} }
@ -683,34 +850,47 @@ func (t *WebFetchTool) Execute(ctx context.Context, args map[string]any) *ToolRe
return &ToolResult{ return &ToolResult{
ForLLM: fmt.Sprintf( ForLLM: fmt.Sprintf(
"Fetched %d bytes from %s (extractor: %s, truncated: %v)", "Fetched %d bytes from %s (extractor: %s, truncated: %v)",
len(text), len(text),
urlStr, urlStr,
extractor, extractor,
truncated, truncated,
), ),
ForUser: string(resultJSON), ForUser: string(resultJSON),
} }
} }
func (t *WebFetchTool) extractText(htmlContent string) string { func (t *WebFetchTool) extractText(htmlContent string) string {
result := reScript.ReplaceAllLiteralString(htmlContent, "") result := reScript.ReplaceAllLiteralString(htmlContent, "")
result = reStyle.ReplaceAllLiteralString(result, "") result = reStyle.ReplaceAllLiteralString(result, "")
result = reTags.ReplaceAllLiteralString(result, "") result = reTags.ReplaceAllLiteralString(result, "")
result = strings.TrimSpace(result) result = strings.TrimSpace(result)
result = reWhitespace.ReplaceAllString(result, " ") result = reWhitespace.ReplaceAllString(result, " ")
result = reBlankLines.ReplaceAllString(result, "\n\n") result = reBlankLines.ReplaceAllString(result, "\n\n")
lines := strings.Split(result, "\n") lines := strings.Split(result, "\n")
var sb strings.Builder var sb strings.Builder
for _, line := range lines { for _, line := range lines {
line = strings.TrimSpace(line) line = strings.TrimSpace(line)
if line != "" { if line != "" {
if sb.Len() > 0 { if sb.Len() > 0 {
sb.WriteByte('\n') sb.WriteByte('\n')
} }
sb.WriteString(line) sb.WriteString(line)
} }
} }

View file

@ -11,16 +11,22 @@ import (
) )
// TestWebTool_WebFetch_Success verifies successful URL fetching // TestWebTool_WebFetch_Success verifies successful URL fetching
func TestWebTool_WebFetch_Success(t *testing.T) { func TestWebTool_WebFetch_Success(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html") w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
w.Write([]byte("<html><body><h1>Test Page</h1><p>Content here</p></body></html>")) w.Write([]byte("<html><body><h1>Test Page</h1><p>Content here</p></body></html>"))
})) }))
defer server.Close() defer server.Close()
tool := NewWebFetchTool(50000) tool := NewWebFetchTool(50000)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"url": server.URL, "url": server.URL,
} }
@ -28,35 +34,45 @@ func TestWebTool_WebFetch_Success(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// ForUser should contain the fetched content // ForUser should contain the fetched content
if !strings.Contains(result.ForUser, "Test Page") { if !strings.Contains(result.ForUser, "Test Page") {
t.Errorf("Expected ForUser to contain 'Test Page', got: %s", result.ForUser) t.Errorf("Expected ForUser to contain 'Test Page', got: %s", result.ForUser)
} }
// ForLLM should contain summary // ForLLM should contain summary
if !strings.Contains(result.ForLLM, "bytes") && !strings.Contains(result.ForLLM, "extractor") { if !strings.Contains(result.ForLLM, "bytes") && !strings.Contains(result.ForLLM, "extractor") {
t.Errorf("Expected ForLLM to contain summary, got: %s", result.ForLLM) t.Errorf("Expected ForLLM to contain summary, got: %s", result.ForLLM)
} }
} }
// TestWebTool_WebFetch_JSON verifies JSON content handling // TestWebTool_WebFetch_JSON verifies JSON content handling
func TestWebTool_WebFetch_JSON(t *testing.T) { func TestWebTool_WebFetch_JSON(t *testing.T) {
testData := map[string]string{"key": "value", "number": "123"} testData := map[string]string{"key": "value", "number": "123"}
expectedJSON, _ := json.MarshalIndent(testData, "", " ") expectedJSON, _ := json.MarshalIndent(testData, "", " ")
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json") w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
w.Write(expectedJSON) w.Write(expectedJSON)
})) }))
defer server.Close() defer server.Close()
tool := NewWebFetchTool(50000) tool := NewWebFetchTool(50000)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"url": server.URL, "url": server.URL,
} }
@ -64,20 +80,25 @@ func TestWebTool_WebFetch_JSON(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// ForUser should contain formatted JSON // ForUser should contain formatted JSON
if !strings.Contains(result.ForUser, "key") && !strings.Contains(result.ForUser, "value") { if !strings.Contains(result.ForUser, "key") && !strings.Contains(result.ForUser, "value") {
t.Errorf("Expected ForUser to contain JSON data, got: %s", result.ForUser) t.Errorf("Expected ForUser to contain JSON data, got: %s", result.ForUser)
} }
} }
// TestWebTool_WebFetch_InvalidURL verifies error handling for invalid URL // TestWebTool_WebFetch_InvalidURL verifies error handling for invalid URL
func TestWebTool_WebFetch_InvalidURL(t *testing.T) { func TestWebTool_WebFetch_InvalidURL(t *testing.T) {
tool := NewWebFetchTool(50000) tool := NewWebFetchTool(50000)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"url": "not-a-valid-url", "url": "not-a-valid-url",
} }
@ -85,20 +106,25 @@ func TestWebTool_WebFetch_InvalidURL(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error for invalid URL") t.Errorf("Expected error for invalid URL")
} }
// Should contain error message (either "invalid URL" or scheme error) // Should contain error message (either "invalid URL" or scheme error)
if !strings.Contains(result.ForLLM, "URL") && !strings.Contains(result.ForUser, "URL") { if !strings.Contains(result.ForLLM, "URL") && !strings.Contains(result.ForUser, "URL") {
t.Errorf("Expected error message for invalid URL, got ForLLM: %s", result.ForLLM) t.Errorf("Expected error message for invalid URL, got ForLLM: %s", result.ForLLM)
} }
} }
// TestWebTool_WebFetch_UnsupportedScheme verifies error handling for non-http URLs // TestWebTool_WebFetch_UnsupportedScheme verifies error handling for non-http URLs
func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) { func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) {
tool := NewWebFetchTool(50000) tool := NewWebFetchTool(50000)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"url": "ftp://example.com/file.txt", "url": "ftp://example.com/file.txt",
} }
@ -106,48 +132,61 @@ func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error for unsupported URL scheme") t.Errorf("Expected error for unsupported URL scheme")
} }
// Should mention only http/https allowed // Should mention only http/https allowed
if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForUser, "http/https") { if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForUser, "http/https") {
t.Errorf("Expected scheme error message, got ForLLM: %s", result.ForLLM) t.Errorf("Expected scheme error message, got ForLLM: %s", result.ForLLM)
} }
} }
// TestWebTool_WebFetch_MissingURL verifies error handling for missing URL // TestWebTool_WebFetch_MissingURL verifies error handling for missing URL
func TestWebTool_WebFetch_MissingURL(t *testing.T) { func TestWebTool_WebFetch_MissingURL(t *testing.T) {
tool := NewWebFetchTool(50000) tool := NewWebFetchTool(50000)
ctx := context.Background() ctx := context.Background()
args := map[string]any{} args := map[string]any{}
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when URL is missing") t.Errorf("Expected error when URL is missing")
} }
// Should mention URL is required // Should mention URL is required
if !strings.Contains(result.ForLLM, "url is required") && !strings.Contains(result.ForUser, "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) t.Errorf("Expected 'url is required' message, got ForLLM: %s", result.ForLLM)
} }
} }
// TestWebTool_WebFetch_Truncation verifies content truncation // TestWebTool_WebFetch_Truncation verifies content truncation
func TestWebTool_WebFetch_Truncation(t *testing.T) { func TestWebTool_WebFetch_Truncation(t *testing.T) {
longContent := strings.Repeat("x", 20000) longContent := strings.Repeat("x", 20000)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/plain") w.Header().Set("Content-Type", "text/plain")
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
w.Write([]byte(longContent)) w.Write([]byte(longContent))
})) }))
defer server.Close() defer server.Close()
tool := NewWebFetchTool(1000) // Limit to 1000 chars tool := NewWebFetchTool(1000) // Limit to 1000 chars
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"url": server.URL, "url": server.URL,
} }
@ -155,13 +194,17 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// ForUser should contain truncated content (not the full 20000 chars) // ForUser should contain truncated content (not the full 20000 chars)
resultMap := make(map[string]any) resultMap := make(map[string]any)
json.Unmarshal([]byte(result.ForUser), &resultMap) json.Unmarshal([]byte(result.ForUser), &resultMap)
if text, ok := resultMap["text"].(string); ok { if text, ok := resultMap["text"].(string); ok {
if len(text) > 1100 { // Allow some margin if len(text) > 1100 { // Allow some margin
t.Errorf("Expected content to be truncated to ~1000 chars, got: %d", len(text)) 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 // Should be marked as truncated
if truncated, ok := resultMap["truncated"].(bool); !ok || !truncated { if truncated, ok := resultMap["truncated"].(bool); !ok || !truncated {
t.Errorf("Expected 'truncated' to be true in result") t.Errorf("Expected 'truncated' to be true in result")
} }
} }
// TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing // TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing
func TestWebTool_WebSearch_NoApiKey(t *testing.T) { func TestWebTool_WebSearch_NoApiKey(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""}) tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""})
if err != nil { if err != nil {
t.Fatalf("Unexpected error: %v", err) t.Fatalf("Unexpected error: %v", err)
} }
if tool != nil { if tool != nil {
t.Errorf("Expected nil tool when Brave API key is empty") t.Errorf("Expected nil tool when Brave API key is empty")
} }
// Also nil when nothing is enabled // Also nil when nothing is enabled
tool, err = NewWebSearchTool(WebSearchToolOptions{}) tool, err = NewWebSearchTool(WebSearchToolOptions{})
if err != nil { if err != nil {
t.Fatalf("Unexpected error: %v", err) t.Fatalf("Unexpected error: %v", err)
} }
if tool != nil { if tool != nil {
t.Errorf("Expected nil tool when no provider is enabled") t.Errorf("Expected nil tool when no provider is enabled")
} }
} }
// TestWebTool_WebSearch_MissingQuery verifies error handling for missing query // TestWebTool_WebSearch_MissingQuery verifies error handling for missing query
func TestWebTool_WebSearch_MissingQuery(t *testing.T) { func TestWebTool_WebSearch_MissingQuery(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5}) tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5})
if err != nil { if err != nil {
t.Fatalf("Unexpected error: %v", err) t.Fatalf("Unexpected error: %v", err)
} }
ctx := context.Background() ctx := context.Background()
args := map[string]any{} args := map[string]any{}
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error when query is missing") t.Errorf("Expected error when query is missing")
} }
} }
// TestWebTool_WebFetch_HTMLExtraction verifies HTML text extraction // TestWebTool_WebFetch_HTMLExtraction verifies HTML text extraction
func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html") w.Header().Set("Content-Type", "text/html")
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
w.Write( w.Write(
[]byte( []byte(
`<html><body><script>alert('test');</script><style>body{color:red;}</style><h1>Title</h1><p>Content</p></body></html>`, `<html><body><script>alert('test');</script><style>body{color:red;}</style><h1>Title</h1><p>Content</p></body></html>`,
), ),
) )
})) }))
defer server.Close() defer server.Close()
tool := NewWebFetchTool(50000) tool := NewWebFetchTool(50000)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"url": server.URL, "url": server.URL,
} }
@ -233,80 +292,105 @@ func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// ForUser should contain extracted text (without script/style tags) // ForUser should contain extracted text (without script/style tags)
if !strings.Contains(result.ForUser, "Title") && !strings.Contains(result.ForUser, "Content") { if !strings.Contains(result.ForUser, "Title") && !strings.Contains(result.ForUser, "Content") {
t.Errorf("Expected ForUser to contain extracted text, got: %s", result.ForUser) t.Errorf("Expected ForUser to contain extracted text, got: %s", result.ForUser)
} }
// Should NOT contain script or style tags // Should NOT contain script or style tags
if strings.Contains(result.ForUser, "<script>") || strings.Contains(result.ForUser, "<style>") { if strings.Contains(result.ForUser, "<script>") || strings.Contains(result.ForUser, "<style>") {
t.Errorf("Expected script/style tags to be removed, got: %s", result.ForUser) t.Errorf("Expected script/style tags to be removed, got: %s", result.ForUser)
} }
} }
// TestWebFetchTool_extractText verifies text extraction preserves newlines // TestWebFetchTool_extractText verifies text extraction preserves newlines
func TestWebFetchTool_extractText(t *testing.T) { func TestWebFetchTool_extractText(t *testing.T) {
tool := &WebFetchTool{} tool := &WebFetchTool{}
tests := []struct { tests := []struct {
name string name string
input string input string
wantFunc func(t *testing.T, got 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>", 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) { wantFunc: func(t *testing.T, got string) {
lines := strings.Split(got, "\n") lines := strings.Split(got, "\n")
if len(lines) < 2 { if len(lines) < 2 {
t.Errorf("Expected multiple lines, got %d: %q", len(lines), got) t.Errorf("Expected multiple lines, got %d: %q", len(lines), got)
} }
if !strings.Contains(got, "Title") || !strings.Contains(got, "Paragraph 1") || if !strings.Contains(got, "Title") || !strings.Contains(got, "Paragraph 1") ||
!strings.Contains(got, "Paragraph 2") { !strings.Contains(got, "Paragraph 2") {
t.Errorf("Missing expected text: %q", got) 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>", input: "<script>alert('x');</script><style>body{}</style><p>Keep this</p>",
wantFunc: func(t *testing.T, got string) { wantFunc: func(t *testing.T, got string) {
if strings.Contains(got, "alert") || strings.Contains(got, "body{}") { if strings.Contains(got, "alert") || strings.Contains(got, "body{}") {
t.Errorf("Expected script/style content removed, got: %q", got) t.Errorf("Expected script/style content removed, got: %q", got)
} }
if !strings.Contains(got, "Keep this") { if !strings.Contains(got, "Keep this") {
t.Errorf("Expected 'Keep this' to remain, got: %q", got) 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>", input: "<p>A</p>\n\n\n\n\n<p>B</p>",
wantFunc: func(t *testing.T, got string) { wantFunc: func(t *testing.T, got string) {
if strings.Contains(got, "\n\n\n") { if strings.Contains(got, "\n\n\n") {
t.Errorf("Expected excessive blank lines collapsed, got: %q", got) t.Errorf("Expected excessive blank lines collapsed, got: %q", got)
} }
}, },
}, },
{ {
name: "collapses horizontal whitespace", name: "collapses horizontal whitespace",
input: "<p>hello world</p>", input: "<p>hello world</p>",
wantFunc: func(t *testing.T, got string) { wantFunc: func(t *testing.T, got string) {
if strings.Contains(got, " ") { if strings.Contains(got, " ") {
t.Errorf("Expected spaces collapsed, got: %q", got) t.Errorf("Expected spaces collapsed, got: %q", got)
} }
if !strings.Contains(got, "hello world") { if !strings.Contains(got, "hello world") {
t.Errorf("Expected 'hello world', got: %q", got) t.Errorf("Expected 'hello world', got: %q", got)
} }
}, },
}, },
{ {
name: "empty input", name: "empty input",
input: "", input: "",
wantFunc: func(t *testing.T, got string) { wantFunc: func(t *testing.T, got string) {
if got != "" { if got != "" {
t.Errorf("Expected empty string, got: %q", got) t.Errorf("Expected empty string, got: %q", got)
@ -318,15 +402,19 @@ func TestWebFetchTool_extractText(t *testing.T) {
for _, tt := range tests { for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) { t.Run(tt.name, func(t *testing.T) {
got := tool.extractText(tt.input) got := tool.extractText(tt.input)
tt.wantFunc(t, got) tt.wantFunc(t, got)
}) })
} }
} }
// TestWebTool_WebFetch_MissingDomain verifies error handling for URL without domain // TestWebTool_WebFetch_MissingDomain verifies error handling for URL without domain
func TestWebTool_WebFetch_MissingDomain(t *testing.T) { func TestWebTool_WebFetch_MissingDomain(t *testing.T) {
tool := NewWebFetchTool(50000) tool := NewWebFetchTool(50000)
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"url": "https://", "url": "https://",
} }
@ -334,11 +422,13 @@ func TestWebTool_WebFetch_MissingDomain(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Should return error result // Should return error result
if !result.IsError { if !result.IsError {
t.Errorf("Expected error for URL without domain") t.Errorf("Expected error for URL without domain")
} }
// Should mention missing domain // Should mention missing domain
if !strings.Contains(result.ForLLM, "domain") && !strings.Contains(result.ForUser, "domain") { if !strings.Contains(result.ForLLM, "domain") && !strings.Contains(result.ForUser, "domain") {
t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM) t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM)
} }
@ -349,14 +439,17 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("createHTTPClient() error: %v", err) t.Fatalf("createHTTPClient() error: %v", err)
} }
if client.Timeout != 12*time.Second { if client.Timeout != 12*time.Second {
t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second) t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second)
} }
tr, ok := client.Transport.(*http.Transport) tr, ok := client.Transport.(*http.Transport)
if !ok { if !ok {
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
} }
if tr.Proxy == nil { if tr.Proxy == nil {
t.Fatal("transport.Proxy is nil, want non-nil") t.Fatal("transport.Proxy is nil, want non-nil")
} }
@ -365,10 +458,12 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("http.NewRequest() error: %v", err) t.Fatalf("http.NewRequest() error: %v", err)
} }
proxyURL, err := tr.Proxy(req) proxyURL, err := tr.Proxy(req)
if err != nil { if err != nil {
t.Fatalf("transport.Proxy(req) error: %v", err) t.Fatalf("transport.Proxy(req) error: %v", err)
} }
if proxyURL == nil || proxyURL.String() != "http://127.0.0.1:7890" { 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") 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) { func TestCreateHTTPClient_InvalidProxy(t *testing.T) {
_, err := createHTTPClient("://bad-proxy", 10*time.Second) _, err := createHTTPClient("://bad-proxy", 10*time.Second)
if err == nil { if err == nil {
t.Fatal("createHTTPClient() expected error for invalid proxy URL, got 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) tr, ok := client.Transport.(*http.Transport)
if !ok { if !ok {
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
} }
req, err := http.NewRequest("GET", "https://example.com", nil) req, err := http.NewRequest("GET", "https://example.com", nil)
if err != nil { if err != nil {
t.Fatalf("http.NewRequest() error: %v", err) t.Fatalf("http.NewRequest() error: %v", err)
} }
proxyURL, err := tr.Proxy(req) proxyURL, err := tr.Proxy(req)
if err != nil { if err != nil {
t.Fatalf("transport.Proxy(req) error: %v", err) t.Fatalf("transport.Proxy(req) error: %v", err)
} }
if proxyURL == nil || proxyURL.String() != "socks5://127.0.0.1:1080" { 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") 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) { func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) {
_, err := createHTTPClient("ftp://127.0.0.1:21", 10*time.Second) _, err := createHTTPClient("ftp://127.0.0.1:21", 10*time.Second)
if err == nil { if err == nil {
t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil") t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil")
} }
if !strings.Contains(err.Error(), "unsupported proxy scheme") { if !strings.Contains(err.Error(), "unsupported proxy scheme") {
t.Fatalf("error = %q, want to contain %q", 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) { 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("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("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("all_proxy", "") t.Setenv("all_proxy", "")
t.Setenv("NO_PROXY", "") t.Setenv("NO_PROXY", "")
t.Setenv("no_proxy", "") t.Setenv("no_proxy", "")
client, err := createHTTPClient("", 10*time.Second) client, err := createHTTPClient("", 10*time.Second)
@ -430,9 +539,11 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) {
} }
tr, ok := client.Transport.(*http.Transport) tr, ok := client.Transport.(*http.Transport)
if !ok { if !ok {
t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport)
} }
if tr.Proxy == nil { if tr.Proxy == nil {
t.Fatal("transport.Proxy is nil, want proxy function from environment") t.Fatal("transport.Proxy is nil, want proxy function from environment")
} }
@ -441,6 +552,7 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("http.NewRequest() error: %v", err) t.Fatalf("http.NewRequest() error: %v", err)
} }
if _, err := tr.Proxy(req); err != nil { if _, err := tr.Proxy(req); err != nil {
t.Fatalf("transport.Proxy(req) error: %v", err) t.Fatalf("transport.Proxy(req) error: %v", err)
} }
@ -451,9 +563,11 @@ func TestNewWebFetchToolWithProxy(t *testing.T) {
if err != nil { if err != nil {
t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) t.Fatalf("NewWebFetchToolWithProxy() error: %v", err)
} }
if tool.maxChars != 1024 { if tool.maxChars != 1024 {
t.Fatalf("maxChars = %d, want %d", tool.maxChars, 1024) t.Fatalf("maxChars = %d, want %d", tool.maxChars, 1024)
} }
if tool.proxy != "http://127.0.0.1:7890" { if tool.proxy != "http://127.0.0.1:7890" {
t.Fatalf("proxy = %q, want %q", 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 { if err != nil {
t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) t.Fatalf("NewWebFetchToolWithProxy() error: %v", err)
} }
if tool.maxChars != 50000 { if tool.maxChars != 50000 {
t.Fatalf("default maxChars = %d, want %d", tool.maxChars, 50000) t.Fatalf("default maxChars = %d, want %d", tool.maxChars, 50000)
} }
@ -471,17 +586,23 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
t.Run("perplexity", func(t *testing.T) { t.Run("perplexity", func(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
PerplexityEnabled: true, PerplexityEnabled: true,
PerplexityAPIKey: "k", PerplexityAPIKey: "k",
PerplexityMaxResults: 3, PerplexityMaxResults: 3,
Proxy: "http://127.0.0.1:7890", Proxy: "http://127.0.0.1:7890",
}) })
if err != nil { if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err) t.Fatalf("NewWebSearchTool() error: %v", err)
} }
p, ok := tool.provider.(*PerplexitySearchProvider) p, ok := tool.provider.(*PerplexitySearchProvider)
if !ok { if !ok {
t.Fatalf("provider type = %T, want *PerplexitySearchProvider", tool.provider) t.Fatalf("provider type = %T, want *PerplexitySearchProvider", tool.provider)
} }
if p.proxy != "http://127.0.0.1:7890" { if p.proxy != "http://127.0.0.1:7890" {
t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890")
} }
@ -490,17 +611,23 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
t.Run("brave", func(t *testing.T) { t.Run("brave", func(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
BraveEnabled: true, BraveEnabled: true,
BraveAPIKey: "k", BraveAPIKey: "k",
BraveMaxResults: 3, BraveMaxResults: 3,
Proxy: "http://127.0.0.1:7890", Proxy: "http://127.0.0.1:7890",
}) })
if err != nil { if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err) t.Fatalf("NewWebSearchTool() error: %v", err)
} }
p, ok := tool.provider.(*BraveSearchProvider) p, ok := tool.provider.(*BraveSearchProvider)
if !ok { if !ok {
t.Fatalf("provider type = %T, want *BraveSearchProvider", tool.provider) t.Fatalf("provider type = %T, want *BraveSearchProvider", tool.provider)
} }
if p.proxy != "http://127.0.0.1:7890" { if p.proxy != "http://127.0.0.1:7890" {
t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890")
} }
@ -509,16 +636,21 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) {
t.Run("duckduckgo", func(t *testing.T) { t.Run("duckduckgo", func(t *testing.T) {
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
DuckDuckGoEnabled: true, DuckDuckGoEnabled: true,
DuckDuckGoMaxResults: 3, DuckDuckGoMaxResults: 3,
Proxy: "http://127.0.0.1:7890", Proxy: "http://127.0.0.1:7890",
}) })
if err != nil { if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err) t.Fatalf("NewWebSearchTool() error: %v", err)
} }
p, ok := tool.provider.(*DuckDuckGoSearchProvider) p, ok := tool.provider.(*DuckDuckGoSearchProvider)
if !ok { if !ok {
t.Fatalf("provider type = %T, want *DuckDuckGoSearchProvider", tool.provider) t.Fatalf("provider type = %T, want *DuckDuckGoSearchProvider", tool.provider)
} }
if p.proxy != "http://127.0.0.1:7890" { if p.proxy != "http://127.0.0.1:7890" {
t.Fatalf("provider proxy = %q, want %q", 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 // TestWebTool_TavilySearch_Success verifies successful Tavily search
func TestWebTool_TavilySearch_Success(t *testing.T) { func TestWebTool_TavilySearch_Success(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" { if r.Method != "POST" {
t.Errorf("Expected POST request, got %s", r.Method) t.Errorf("Expected POST request, got %s", r.Method)
} }
if r.Header.Get("Content-Type") != "application/json" { if r.Header.Get("Content-Type") != "application/json" {
t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type"))
} }
// Verify payload // Verify payload
var payload map[string]any var payload map[string]any
json.NewDecoder(r.Body).Decode(&payload) json.NewDecoder(r.Body).Decode(&payload)
if payload["api_key"] != "test-key" { if payload["api_key"] != "test-key" {
t.Errorf("Expected api_key test-key, got %v", payload["api_key"]) t.Errorf("Expected api_key test-key, got %v", payload["api_key"])
} }
if payload["query"] != "test query" { if payload["query"] != "test query" {
t.Errorf("Expected query 'test query', got %v", payload["query"]) t.Errorf("Expected query 'test query', got %v", payload["query"])
} }
// Return mock response // Return mock response
response := map[string]any{ response := map[string]any{
"results": []map[string]any{ "results": []map[string]any{
{ {
"title": "Test Result 1", "title": "Test Result 1",
"url": "https://example.com/1", "url": "https://example.com/1",
"content": "Content for result 1", "content": "Content for result 1",
}, },
{ {
"title": "Test Result 2", "title": "Test Result 2",
"url": "https://example.com/2", "url": "https://example.com/2",
"content": "Content for result 2", "content": "Content for result 2",
}, },
}, },
} }
w.Header().Set("Content-Type", "application/json") w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK) w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(response) json.NewEncoder(w).Encode(response)
})) }))
defer server.Close() defer server.Close()
tool, err := NewWebSearchTool(WebSearchToolOptions{ tool, err := NewWebSearchTool(WebSearchToolOptions{
TavilyEnabled: true, TavilyEnabled: true,
TavilyAPIKey: "test-key", TavilyAPIKey: "test-key",
TavilyBaseURL: server.URL, TavilyBaseURL: server.URL,
TavilyMaxResults: 5, TavilyMaxResults: 5,
}) })
if err != nil { if err != nil {
@ -577,6 +728,7 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
} }
ctx := context.Background() ctx := context.Background()
args := map[string]any{ args := map[string]any{
"query": "test query", "query": "test query",
} }
@ -584,17 +736,21 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
result := tool.Execute(ctx, args) result := tool.Execute(ctx, args)
// Success should not be an error // Success should not be an error
if result.IsError { if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
} }
// ForUser should contain result titles and URLs // ForUser should contain result titles and URLs
if !strings.Contains(result.ForUser, "Test Result 1") || if !strings.Contains(result.ForUser, "Test Result 1") ||
!strings.Contains(result.ForUser, "https://example.com/1") { !strings.Contains(result.ForUser, "https://example.com/1") {
t.Errorf("Expected results in output, got: %s", result.ForUser) t.Errorf("Expected results in output, got: %s", result.ForUser)
} }
// Should mention via Tavily // Should mention via Tavily
if !strings.Contains(result.ForUser, "via Tavily") { if !strings.Contains(result.ForUser, "via Tavily") {
t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser)
} }

View file

@ -8,52 +8,72 @@ import (
type ( type (
workspaceOverrideKey struct{} workspaceOverrideKey struct{}
overrideFsKey struct{} overrideFsKey struct{}
) )
// WithWorkspaceOverride returns a context carrying a workspace override path // WithWorkspaceOverride returns a context carrying a workspace override path
// and a pre-built sandboxFs for that workspace. Tools will resolve file // and a pre-built sandboxFs for that workspace. Tools will resolve file
// operations against this path instead of the original workspace. // operations against this path instead of the original workspace.
// The cached sandboxFs is reused across all resolveFS calls on the same context, // The cached sandboxFs is reused across all resolveFS calls on the same context,
// avoiding per-operation allocation. // avoiding per-operation allocation.
func WithWorkspaceOverride(ctx context.Context, workspace string) context.Context { func WithWorkspaceOverride(ctx context.Context, workspace string) context.Context {
ctx = context.WithValue(ctx, workspaceOverrideKey{}, workspace) ctx = context.WithValue(ctx, workspaceOverrideKey{}, workspace)
ctx = context.WithValue(ctx, overrideFsKey{}, &sandboxFs{workspace: workspace}) ctx = context.WithValue(ctx, overrideFsKey{}, &sandboxFs{workspace: workspace})
return ctx return ctx
} }
// WorkspaceOverrideFromCtx extracts the workspace override from context, or "". // WorkspaceOverrideFromCtx extracts the workspace override from context, or "".
func WorkspaceOverrideFromCtx(ctx context.Context) string { func WorkspaceOverrideFromCtx(ctx context.Context) string {
if v, ok := ctx.Value(workspaceOverrideKey{}).(string); ok { if v, ok := ctx.Value(workspaceOverrideKey{}).(string); ok {
return v return v
} }
return "" return ""
} }
// resolveFS returns a fileSystem applying workspace override from context. // resolveFS returns a fileSystem applying workspace override from context.
// Paths under "memory/" are excluded (always use original workspace). // Paths under "memory/" are excluded (always use original workspace).
// For sandboxFs: returns the cached override instance from context. // For sandboxFs: returns the cached override instance from context.
// For hostFs (unrestricted): returns as-is. // For hostFs (unrestricted): returns as-is.
func resolveFS(ctx context.Context, fs fileSystem, path string) fileSystem { func resolveFS(ctx context.Context, fs fileSystem, path string) fileSystem {
override := WorkspaceOverrideFromCtx(ctx) override := WorkspaceOverrideFromCtx(ctx)
if override == "" { if override == "" {
return fs return fs
} }
// memory/ paths always use original workspace // memory/ paths always use original workspace
if isMemoryPath(path) { if isMemoryPath(path) {
return fs return fs
} }
// Only sandboxFs supports workspace override // Only sandboxFs supports workspace override
if sfs, ok := fs.(*sandboxFs); ok { if sfs, ok := fs.(*sandboxFs); ok {
if sfs.workspace == override { if sfs.workspace == override {
return fs return fs
} }
// Use cached sandboxFs from context // Use cached sandboxFs from context
if cached, ok := ctx.Value(overrideFsKey{}).(*sandboxFs); ok { if cached, ok := ctx.Value(overrideFsKey{}).(*sandboxFs); ok {
return cached return cached
} }
return &sandboxFs{workspace: override} 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. // isMemoryPath returns true for paths under the memory/ directory.
// Matches: "memory/MEMORY.md", "memory", "/workspace/memory/notes.md" // Matches: "memory/MEMORY.md", "memory", "/workspace/memory/notes.md"
func isMemoryPath(path string) bool { func isMemoryPath(path string) bool {
p := filepath.ToSlash(filepath.Clean(path)) p := filepath.ToSlash(filepath.Clean(path))
// Relative path starting with memory/ // Relative path starting with memory/
if strings.HasPrefix(p, "memory/") || p == "memory" { if strings.HasPrefix(p, "memory/") || p == "memory" {
return true return true
} }
// Absolute path containing /memory/ or ending with /memory // Absolute path containing /memory/ or ending with /memory
if strings.Contains(p, "/memory/") || strings.HasSuffix(p, "/memory") { if strings.Contains(p, "/memory/") || strings.HasSuffix(p, "/memory") {
return true return true
} }