From 1c12de76306836f1f97962ad69df0242e0522eca Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Fri, 13 Mar 2026 03:16:04 +0900 Subject: [PATCH] refactor: replace test files with upstream versions, extract fork-only tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace 43 test files with their upstream/main versions to eliminate test file merge conflicts entirely (43 files, 604 markers → 0). Fork-only test functions are extracted to *_ext_test.go files (19 files) which have no upstream counterpart and thus never conflict. Source-level upstream alignment: - config/defaults: AllowRemote defaults to true - config/migration: model names match upstream (gpt-5.4) - session/manager: sanitizeFilename replaces / and \ - state: log.Printf instead of log.Fatalf on mkdir failure - wecom: verifySignature returns false on empty token (fail-closed) - openclaw migration: preserves AllowRemote default Co-Authored-By: Claude Opus 4.6 --- cmd/picoclaw/internal/helpers_ext_test.go | 80 + cmd/picoclaw/internal/helpers_test.go | 73 - cmd/picoclaw/main_test.go | 3 +- pkg/agent/context_cache_test.go | 262 +- pkg/agent/context_test.go | 62 - pkg/agent/instance_ext_test.go | 52 + pkg/agent/instance_test.go | 171 +- pkg/agent/loop_ext_test.go | 2918 ++++++++++++ pkg/agent/loop_test.go | 4212 ++--------------- pkg/agent/mock_provider_test.go | 7 +- pkg/agent/registry_test.go | 63 +- pkg/channels/manager_ext_test.go | 881 ++++ pkg/channels/manager_test.go | 1085 +---- pkg/channels/matrix/matrix_test.go | 96 + pkg/channels/telegram/telegram_ext_test.go | 52 + pkg/channels/telegram/telegram_test.go | 494 +- pkg/channels/wecom/app_test.go | 6 +- pkg/channels/wecom/bot_test.go | 7 +- pkg/channels/wecom/common.go | 2 +- pkg/config/config_ext_test.go | 123 + pkg/config/config_test.go | 264 +- pkg/config/defaults.go | 1 + pkg/config/migration.go | 6 +- pkg/config/migration_test.go | 21 +- pkg/heartbeat/service_ext_test.go | 103 + pkg/heartbeat/service_test.go | 95 - pkg/logger/logger_ext_test.go | 289 ++ pkg/logger/logger_test.go | 287 -- pkg/memory/migration_test.go | 52 + .../sources/openclaw/openclaw_config.go | 1 + .../sources/openclaw/openclaw_config_test.go | 14 + pkg/providers/anthropic/provider_test.go | 6 +- pkg/providers/antigravity_provider_test.go | 2 +- pkg/providers/claude_cli_provider_ext_test.go | 283 ++ pkg/providers/claude_cli_provider_test.go | 289 +- pkg/providers/codex_cli_provider_test.go | 12 +- pkg/providers/codex_provider_test.go | 18 +- pkg/providers/factory_ext_test.go | 73 + pkg/providers/factory_provider_test.go | 24 + pkg/providers/factory_test.go | 86 +- .../openai_compat/provider_ext_test.go | 456 ++ pkg/providers/openai_compat/provider_test.go | 864 ++-- pkg/session/manager.go | 5 +- pkg/session/manager_ext_test.go | 121 + pkg/session/manager_test.go | 165 +- pkg/skills/loader_test.go | 75 + pkg/state/state.go | 4 +- pkg/state/state_ext_test.go | 38 + pkg/state/state_test.go | 44 +- pkg/tools/edit_test.go | 236 +- pkg/tools/filesystem_ext_test.go | 183 + pkg/tools/filesystem_test.go | 371 +- pkg/tools/message_test.go | 81 +- pkg/tools/registry_ext_test.go | 269 ++ pkg/tools/registry_test.go | 382 +- pkg/tools/result_test.go | 62 +- pkg/tools/shell_ext_test.go | 896 ++++ pkg/tools/shell_test.go | 1224 +---- pkg/tools/shell_timeout_unix_ext_test.go | 35 + pkg/tools/shell_timeout_unix_test.go | 31 +- pkg/tools/skills_install_test.go | 39 +- pkg/tools/skills_search_test.go | 40 +- pkg/tools/spawn_test.go | 34 +- pkg/tools/subagent_tool_ext_test.go | 102 + pkg/tools/subagent_tool_test.go | 263 +- pkg/tools/web_test.go | 789 ++- pkg/utils/string_ext_test.go | 179 + pkg/utils/string_test.go | 187 +- 68 files changed, 10498 insertions(+), 9252 deletions(-) create mode 100644 cmd/picoclaw/internal/helpers_ext_test.go create mode 100644 pkg/agent/instance_ext_test.go create mode 100644 pkg/agent/loop_ext_test.go create mode 100644 pkg/channels/manager_ext_test.go create mode 100644 pkg/channels/telegram/telegram_ext_test.go create mode 100644 pkg/config/config_ext_test.go create mode 100644 pkg/heartbeat/service_ext_test.go create mode 100644 pkg/logger/logger_ext_test.go create mode 100644 pkg/providers/claude_cli_provider_ext_test.go create mode 100644 pkg/providers/factory_ext_test.go create mode 100644 pkg/providers/openai_compat/provider_ext_test.go create mode 100644 pkg/session/manager_ext_test.go create mode 100644 pkg/state/state_ext_test.go create mode 100644 pkg/tools/filesystem_ext_test.go create mode 100644 pkg/tools/registry_ext_test.go create mode 100644 pkg/tools/shell_ext_test.go create mode 100644 pkg/tools/shell_timeout_unix_ext_test.go create mode 100644 pkg/tools/subagent_tool_ext_test.go create mode 100644 pkg/utils/string_ext_test.go diff --git a/cmd/picoclaw/internal/helpers_ext_test.go b/cmd/picoclaw/internal/helpers_ext_test.go new file mode 100644 index 000000000..7e9ca5097 --- /dev/null +++ b/cmd/picoclaw/internal/helpers_ext_test.go @@ -0,0 +1,80 @@ +package internal + +import ( + "github.com/stretchr/testify/assert" + "runtime" + "testing" +) + +func TestFormatVersion_NoGitCommit(t *testing.T) { + oldVersion, oldGit := version, gitCommit + t.Cleanup(func() { version, gitCommit = oldVersion, oldGit }) + + version = "1.2.3" + gitCommit = "" + + assert.Equal(t, "1.2.3", FormatVersion()) +} + +func TestFormatVersion_WithGitCommit(t *testing.T) { + oldVersion, oldGit := version, gitCommit + t.Cleanup(func() { version, gitCommit = oldVersion, oldGit }) + + version = "1.2.3" + gitCommit = "abc123" + + assert.Equal(t, "1.2.3 (git: abc123)", FormatVersion()) +} + +func TestFormatBuildInfo_UsesBuildTimeAndGoVersion_WhenSet(t *testing.T) { + oldBuildTime, oldGoVersion := buildTime, goVersion + t.Cleanup(func() { buildTime, goVersion = oldBuildTime, oldGoVersion }) + + buildTime = "2026-02-20T00:00:00Z" + goVersion = "go1.23.0" + + build, goVer := FormatBuildInfo() + + assert.Equal(t, buildTime, build) + assert.Equal(t, goVersion, goVer) +} + +func TestFormatBuildInfo_EmptyBuildTime_ReturnsEmptyBuild(t *testing.T) { + oldBuildTime, oldGoVersion := buildTime, goVersion + t.Cleanup(func() { buildTime, goVersion = oldBuildTime, oldGoVersion }) + + buildTime = "" + goVersion = "go1.23.0" + + build, goVer := FormatBuildInfo() + + assert.Empty(t, build) + assert.Equal(t, goVersion, goVer) +} + +func TestFormatBuildInfo_EmptyGoVersion_FallsBackToRuntimeVersion(t *testing.T) { + oldBuildTime, oldGoVersion := buildTime, goVersion + t.Cleanup(func() { buildTime, goVersion = oldBuildTime, oldGoVersion }) + + buildTime = "x" + goVersion = "" + + build, goVer := FormatBuildInfo() + + assert.Equal(t, "x", build) + assert.Equal(t, runtime.Version(), goVer) +} + +func TestGetVersion(t *testing.T) { + assert.Equal(t, "dev", GetVersion()) +} + +func TestGetConfigPath_WithEnv(t *testing.T) { + t.Setenv("PICOCLAW_CONFIG", "/tmp/custom/config.json") + t.Setenv("HOME", "/tmp/home") + + got := GetConfigPath() + want := "/tmp/custom/config.json" + + assert.Equal(t, want, got) +} diff --git a/cmd/picoclaw/internal/helpers_test.go b/cmd/picoclaw/internal/helpers_test.go index 646be1ba1..583751781 100644 --- a/cmd/picoclaw/internal/helpers_test.go +++ b/cmd/picoclaw/internal/helpers_test.go @@ -40,65 +40,6 @@ func TestGetConfigPath_WithPICOCLAW_CONFIG(t *testing.T) { assert.Equal(t, want, got) } -func TestFormatVersion_NoGitCommit(t *testing.T) { - oldVersion, oldGit := version, gitCommit - t.Cleanup(func() { version, gitCommit = oldVersion, oldGit }) - - version = "1.2.3" - gitCommit = "" - - assert.Equal(t, "1.2.3", FormatVersion()) -} - -func TestFormatVersion_WithGitCommit(t *testing.T) { - oldVersion, oldGit := version, gitCommit - t.Cleanup(func() { version, gitCommit = oldVersion, oldGit }) - - version = "1.2.3" - gitCommit = "abc123" - - assert.Equal(t, "1.2.3 (git: abc123)", FormatVersion()) -} - -func TestFormatBuildInfo_UsesBuildTimeAndGoVersion_WhenSet(t *testing.T) { - oldBuildTime, oldGoVersion := buildTime, goVersion - t.Cleanup(func() { buildTime, goVersion = oldBuildTime, oldGoVersion }) - - buildTime = "2026-02-20T00:00:00Z" - goVersion = "go1.23.0" - - build, goVer := FormatBuildInfo() - - assert.Equal(t, buildTime, build) - assert.Equal(t, goVersion, goVer) -} - -func TestFormatBuildInfo_EmptyBuildTime_ReturnsEmptyBuild(t *testing.T) { - oldBuildTime, oldGoVersion := buildTime, goVersion - t.Cleanup(func() { buildTime, goVersion = oldBuildTime, oldGoVersion }) - - buildTime = "" - goVersion = "go1.23.0" - - build, goVer := FormatBuildInfo() - - assert.Empty(t, build) - assert.Equal(t, goVersion, goVer) -} - -func TestFormatBuildInfo_EmptyGoVersion_FallsBackToRuntimeVersion(t *testing.T) { - oldBuildTime, oldGoVersion := buildTime, goVersion - t.Cleanup(func() { buildTime, goVersion = oldBuildTime, oldGoVersion }) - - buildTime = "x" - goVersion = "" - - build, goVer := FormatBuildInfo() - - assert.Equal(t, "x", build) - assert.Equal(t, runtime.Version(), goVer) -} - func TestGetConfigPath_Windows(t *testing.T) { if runtime.GOOS != "windows" { t.Skip("windows-specific HOME behavior varies; run on windows") @@ -112,17 +53,3 @@ func TestGetConfigPath_Windows(t *testing.T) { require.True(t, strings.EqualFold(got, want), "GetConfigPath() = %q, want %q", got, want) } - -func TestGetVersion(t *testing.T) { - assert.Equal(t, "dev", GetVersion()) -} - -func TestGetConfigPath_WithEnv(t *testing.T) { - t.Setenv("PICOCLAW_CONFIG", "/tmp/custom/config.json") - t.Setenv("HOME", "/tmp/home") // Also set home to ensure env is preferred - - got := GetConfigPath() - want := "/tmp/custom/config.json" - - assert.Equal(t, want, got) -} diff --git a/cmd/picoclaw/main_test.go b/cmd/picoclaw/main_test.go index 3740ba358..e622675ee 100644 --- a/cmd/picoclaw/main_test.go +++ b/cmd/picoclaw/main_test.go @@ -9,6 +9,7 @@ import ( "github.com/stretchr/testify/require" "github.com/sipeed/picoclaw/cmd/picoclaw/internal" + "github.com/sipeed/picoclaw/pkg/config" ) func TestNewPicoclawCommand(t *testing.T) { @@ -16,7 +17,7 @@ func TestNewPicoclawCommand(t *testing.T) { require.NotNil(t, cmd) - short := fmt.Sprintf("%s picoclaw - Personal AI Assistant v%s\n\n", internal.Logo, internal.GetVersion()) + short := fmt.Sprintf("%s picoclaw - Personal AI Assistant v%s\n\n", internal.Logo, config.GetVersion()) assert.Equal(t, "picoclaw", cmd.Use) assert.Equal(t, short, cmd.Short) diff --git a/pkg/agent/context_cache_test.go b/pkg/agent/context_cache_test.go index aa252ee94..707510820 100644 --- a/pkg/agent/context_cache_test.go +++ b/pkg/agent/context_cache_test.go @@ -12,103 +12,70 @@ import ( ) // setupWorkspace creates a temporary workspace with standard directories and optional files. - // Returns the tmpDir path; caller should defer os.RemoveAll(tmpDir). - func setupWorkspace(t *testing.T, files map[string]string) string { t.Helper() - tmpDir, err := os.MkdirTemp("", "picoclaw-test-*") if err != nil { t.Fatal(err) } - os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755) - os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755) - for name, content := range files { dir := filepath.Dir(filepath.Join(tmpDir, name)) - os.MkdirAll(dir, 0o755) - if err := os.WriteFile(filepath.Join(tmpDir, name), []byte(content), 0o644); err != nil { t.Fatal(err) } } - return tmpDir } // TestSingleSystemMessage verifies that BuildMessages always produces exactly one - // system message regardless of summary/history variations. - // Fix: multiple system messages break Anthropic (top-level system param) and - // Codex (only reads last system message as instructions). - func TestSingleSystemMessage(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ "IDENTITY.md": "# Identity\nTest agent.", }) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) tests := []struct { - name string - + name string history []providers.Message - summary string - message string }{ { - name: "no summary, no history", - + name: "no summary, no history", summary: "", - message: "hello", }, - { - name: "with summary", - + name: "with summary", summary: "Previous conversation discussed X", - message: "hello", }, - { name: "with history and summary", - history: []providers.Message{ {Role: "user", Content: "hi"}, - {Role: "assistant", Content: "hello"}, }, - summary: strings.Repeat("Long summary text. ", 50), - message: "new message", }, - { name: "system message in history is filtered", - history: []providers.Message{ {Role: "system", Content: "stale system prompt from previous session"}, - {Role: "user", Content: "hi"}, - {Role: "assistant", Content: "hello"}, }, - summary: "", - message: "new message", }, } @@ -118,44 +85,35 @@ func TestSingleSystemMessage(t *testing.T) { msgs := cb.BuildMessages(tt.history, tt.summary, tt.message, nil, "test", "chat1") systemCount := 0 - for _, m := range msgs { if m.Role == "system" { systemCount++ } } - if systemCount != 1 { t.Errorf("expected exactly 1 system message, got %d", systemCount) } - if msgs[0].Role != "system" { t.Errorf("first message should be system, got %s", msgs[0].Role) } - if msgs[len(msgs)-1].Role != "user" { t.Errorf("last message should be user, got %s", msgs[len(msgs)-1].Role) } // System message must contain identity (static) and time (dynamic) - sys := msgs[0].Content - if !strings.Contains(sys, "picoclaw") { t.Error("system message missing identity") } - if !strings.Contains(sys, "Current Time") { t.Error("system message missing dynamic time context") } // Summary handling - if tt.summary != "" { if !strings.Contains(sys, "CONTEXT_SUMMARY:") { t.Error("summary present but CONTEXT_SUMMARY prefix missing") } - if !strings.Contains(sys, tt.summary[:20]) { t.Error("summary content not found in system message") } @@ -169,46 +127,29 @@ func TestSingleSystemMessage(t *testing.T) { } // TestMtimeAutoInvalidation verifies that the cache detects source file changes - // via mtime without requiring explicit InvalidateCache(). - // Fix: original implementation had no auto-invalidation — edits to bootstrap files, - // memory, or skills were invisible until process restart. - func TestMtimeAutoInvalidation(t *testing.T) { tests := []struct { - name string - - file string // relative path inside workspace - - contentV1 string - - contentV2 string - + name string + file string // relative path inside workspace + contentV1 string + contentV2 string checkField string // substring to verify in rebuilt prompt }{ { - name: "bootstrap file change", - - file: "IDENTITY.md", - - contentV1: "# Original Identity", - - contentV2: "# Updated Identity", - + name: "bootstrap file change", + file: "IDENTITY.md", + contentV1: "# Original Identity", + contentV2: "# Updated Identity", checkField: "Updated Identity", }, - { - name: "memory file change", - - file: "memory/MEMORY.md", - - contentV1: "# Memory\nUser likes Go.", - - contentV2: "# Memory\nUser likes Rust.", - + name: "memory file change", + file: "memory/MEMORY.md", + contentV1: "# Memory\nUser likes Go.", + contentV2: "# Memory\nUser likes Rust.", checkField: "User likes Rust", }, } @@ -216,7 +157,6 @@ func TestMtimeAutoInvalidation(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{tt.file: tt.contentV1}) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) @@ -224,39 +164,26 @@ func TestMtimeAutoInvalidation(t *testing.T) { sp1 := cb.BuildSystemPromptWithCache() // Overwrite file and set future mtime to ensure detection. - // Use 2s offset for filesystem mtime resolution safety (some FS - // have 1s or coarser granularity, especially in CI containers). - fullPath := filepath.Join(tmpDir, tt.file) - os.WriteFile(fullPath, []byte(tt.contentV2), 0o644) - future := time.Now().Add(2 * time.Second) - os.Chtimes(fullPath, future, future) // Verify sourceFilesChangedLocked detects the mtime change - cb.systemPromptMutex.RLock() - changed := cb.sourceFilesChangedLocked() - cb.systemPromptMutex.RUnlock() - if !changed { t.Fatalf("sourceFilesChangedLocked() should detect %s change", tt.file) } // Should auto-rebuild without explicit InvalidateCache() - sp2 := cb.BuildSystemPromptWithCache() - if sp1 == sp2 { t.Errorf("cache not rebuilt after %s change", tt.file) } - if !strings.Contains(sp2, tt.checkField) { t.Errorf("rebuilt prompt missing expected content %q", tt.checkField) } @@ -264,34 +191,23 @@ func TestMtimeAutoInvalidation(t *testing.T) { } // Skills directory mtime change - t.Run("skills dir change", func(t *testing.T) { tmpDir := setupWorkspace(t, nil) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) - _ = cb.BuildSystemPromptWithCache() // populate cache // Touch skills directory (simulate new skill installed) - skillsDir := filepath.Join(tmpDir, "skills") - future := time.Now().Add(2 * time.Second) - os.Chtimes(skillsDir, future, future) // Verify sourceFilesChangedLocked detects it (cache is rebuilt) - // We confirm by checking internal state: a second call should rebuild. - cb.systemPromptMutex.RLock() - changed := cb.sourceFilesChangedLocked() - cb.systemPromptMutex.RUnlock() - if !changed { t.Error("sourceFilesChangedLocked() should detect skills dir mtime change") } @@ -299,22 +215,17 @@ func TestMtimeAutoInvalidation(t *testing.T) { } // TestExplicitInvalidateCache verifies that InvalidateCache() forces a rebuild - // even when source files haven't changed (useful for tests and reload commands). - func TestExplicitInvalidateCache(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ "IDENTITY.md": "# Test Identity", }) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) sp1 := cb.BuildSystemPromptWithCache() - cb.InvalidateCache() - sp2 := cb.BuildSystemPromptWithCache() if sp1 != sp2 { @@ -322,39 +233,29 @@ func TestExplicitInvalidateCache(t *testing.T) { } // Verify cachedAt was reset - cb.InvalidateCache() - cb.systemPromptMutex.RLock() - if !cb.cachedAt.IsZero() { t.Error("cachedAt should be zero after InvalidateCache()") } - cb.systemPromptMutex.RUnlock() } // TestCacheStability verifies that the static prompt is stable across repeated calls - // when no files change (regression test for issue #607). - func TestCacheStability(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ "IDENTITY.md": "# Identity\nContent", - - "SOUL.md": "# Soul\nContent", + "SOUL.md": "# Soul\nContent", }) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) results := make([]string, 5) - for i := range results { results[i] = cb.BuildSystemPromptWithCache() } - for i := 1; i < len(results); i++ { if results[i] != results[0] { t.Errorf("cached prompt changed between call 0 and %d", i) @@ -362,47 +263,32 @@ func TestCacheStability(t *testing.T) { } // Static prompt must NOT contain per-request data - if strings.Contains(results[0], "Current Time") { t.Error("static cached prompt should not contain time (added dynamically)") } } // TestNewFileCreationInvalidatesCache verifies that creating a source file that - // did not exist when the cache was built triggers a cache rebuild. - // This catches the "from nothing to something" edge case that the old - // modifiedSince (return false on stat error) would miss. - func TestNewFileCreationInvalidatesCache(t *testing.T) { tests := []struct { - name string - - file string // relative path inside workspace - - content string - + name string + file string // relative path inside workspace + content string checkField string // substring to verify in rebuilt prompt }{ { - name: "new bootstrap file", - - file: "SOUL.md", - - content: "# Soul\nBe kind and helpful.", - + name: "new bootstrap file", + file: "SOUL.md", + content: "# Soul\nBe kind and helpful.", checkField: "Be kind and helpful", }, - { - name: "new memory file", - - file: "memory/MEMORY.md", - - content: "# Memory\nUser prefers dark mode.", - + name: "new memory file", + file: "memory/MEMORY.md", + content: "# Memory\nUser prefers dark mode.", checkField: "User prefers dark mode", }, } @@ -410,41 +296,29 @@ func TestNewFileCreationInvalidatesCache(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Start with an empty workspace (no bootstrap/memory files) - tmpDir := setupWorkspace(t, nil) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) // Populate cache — file does not exist yet - sp1 := cb.BuildSystemPromptWithCache() - if strings.Contains(sp1, tt.checkField) { t.Fatalf("prompt should not contain %q before file is created", tt.checkField) } // Create the file after cache was built - fullPath := filepath.Join(tmpDir, tt.file) - os.MkdirAll(filepath.Dir(fullPath), 0o755) - if err := os.WriteFile(fullPath, []byte(tt.content), 0o644); err != nil { t.Fatal(err) } - // Set future mtime to guarantee detection - future := time.Now().Add(2 * time.Second) - os.Chtimes(fullPath, future, future) // Cache should auto-invalidate because file went from absent -> present - sp2 := cb.BuildSystemPromptWithCache() - if !strings.Contains(sp2, tt.checkField) { t.Errorf("cache not invalidated on new file creation: expected %q in prompt", tt.checkField) } @@ -453,89 +327,58 @@ func TestNewFileCreationInvalidatesCache(t *testing.T) { } // TestSkillFileContentChange verifies that modifying a skill file's content - // (not just the directory structure) invalidates the cache. - // This is the scenario where directory mtime alone is insufficient — on most - // filesystems, editing a file inside a directory does NOT update the parent - // directory's mtime. - func TestSkillFileContentChange(t *testing.T) { skillMD := `--- - name: test-skill - description: "A test skill" - --- - # Test Skill v1 - Original content.` tmpDir := setupWorkspace(t, map[string]string{ "skills/test-skill/SKILL.md": skillMD, }) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) // Populate cache - sp1 := cb.BuildSystemPromptWithCache() - _ = sp1 // cache is warm // Modify the skill file content (without touching the skills/ directory) - updatedSkillMD := `--- - name: test-skill - description: "An updated test skill" - --- - # Test Skill v2 - Updated content.` skillPath := filepath.Join(tmpDir, "skills", "test-skill", "SKILL.md") - if err := os.WriteFile(skillPath, []byte(updatedSkillMD), 0o644); err != nil { t.Fatal(err) } - // Set future mtime on the skill file only (NOT the directory) - future := time.Now().Add(2 * time.Second) - os.Chtimes(skillPath, future, future) // Verify that sourceFilesChangedLocked detects the content change - cb.systemPromptMutex.RLock() - changed := cb.sourceFilesChangedLocked() - cb.systemPromptMutex.RUnlock() - if !changed { t.Error("sourceFilesChangedLocked() should detect skill file content change") } // Verify cache is actually rebuilt with new content - sp2 := cb.BuildSystemPromptWithCache() - if sp1 == sp2 && strings.Contains(sp1, "test-skill") { // If the skill appeared in the prompt and the prompt didn't change, - // the cache was not invalidated. - t.Error("cache should be invalidated when skill file content changes") } } @@ -697,75 +540,53 @@ description: delete-me-v1 } // TestConcurrentBuildSystemPromptWithCache verifies that multiple goroutines - // can safely call BuildSystemPromptWithCache concurrently without producing - // empty results, panics, or data races. - // Run with: go test -race ./pkg/agent/ -run TestConcurrentBuildSystemPromptWithCache - func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { tmpDir := setupWorkspace(t, map[string]string{ - "IDENTITY.md": "# Identity\nConcurrency test agent.", - - "SOUL.md": "# Soul\nBe helpful.", - - "memory/MEMORY.md": "# Memory\nUser prefers Go.", - + "IDENTITY.md": "# Identity\nConcurrency test agent.", + "SOUL.md": "# Soul\nBe helpful.", + "memory/MEMORY.md": "# Memory\nUser prefers Go.", "skills/demo/SKILL.md": "---\nname: demo\ndescription: \"demo skill\"\n---\n# Demo", }) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) const goroutines = 20 - const iterations = 50 var wg sync.WaitGroup - errs := make(chan string, goroutines*iterations) for g := range goroutines { wg.Add(1) - go func(id int) { defer wg.Done() - for i := range iterations { result := cb.BuildSystemPromptWithCache() - if result == "" { errs <- "empty prompt returned" - return } - if !strings.Contains(result, "picoclaw") { errs <- "prompt missing identity" - return } // Also exercise BuildMessages concurrently - msgs := cb.BuildMessages(nil, "", "hello", nil, "test", "chat") - if len(msgs) < 2 { errs <- "BuildMessages returned fewer than 2 messages" - return } - if msgs[0].Role != "system" { errs <- "first message not system" - return } // Occasionally invalidate to exercise the write path - if i%10 == 0 { cb.InvalidateCache() } @@ -774,7 +595,6 @@ func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { } wg.Wait() - close(errs) for errMsg := range errs { @@ -785,90 +605,64 @@ func TestConcurrentBuildSystemPromptWithCache(t *testing.T) { // BenchmarkBuildMessagesWithCache measures caching performance. // TestEmptyWorkspaceBaselineDetectsNewFiles verifies that when the cache is - // built on an empty workspace (no tracked files exist), creating a file - // afterwards still triggers cache invalidation. This validates the - // time.Unix(1, 0) fallback for maxMtime: any real file's mtime is after epoch, - // so fileChangedSince correctly detects the absent -> present transition AND - // the mtime comparison succeeds even without artificially inflated Chtimes. - func TestEmptyWorkspaceBaselineDetectsNewFiles(t *testing.T) { // Empty workspace: no bootstrap files, no memory, no skills content. - tmpDir := setupWorkspace(t, nil) - defer os.RemoveAll(tmpDir) cb := NewContextBuilder(tmpDir) // Build cache — all tracked files are absent, maxMtime falls back to epoch. - sp1 := cb.BuildSystemPromptWithCache() // Create a bootstrap file with natural mtime (no Chtimes manipulation). - // The file's mtime should be the current wall-clock time, which is - // strictly after time.Unix(1, 0). - soulPath := filepath.Join(tmpDir, "SOUL.md") - if err := os.WriteFile(soulPath, []byte("# Soul\nNewly created."), 0o644); err != nil { t.Fatal(err) } // Cache should detect the new file via existedAtCache (absent -> present). - cb.systemPromptMutex.RLock() - changed := cb.sourceFilesChangedLocked() - cb.systemPromptMutex.RUnlock() - if !changed { t.Fatal("sourceFilesChangedLocked should detect newly created file on empty workspace") } sp2 := cb.BuildSystemPromptWithCache() - if !strings.Contains(sp2, "Newly created") { t.Error("rebuilt prompt should contain new file content") } - if sp1 == sp2 { t.Error("cache should have been invalidated after file creation") } } // BenchmarkBuildMessagesWithCache measures caching performance. - func BenchmarkBuildMessagesWithCache(b *testing.B) { tmpDir, _ := os.MkdirTemp("", "picoclaw-bench-*") - defer os.RemoveAll(tmpDir) os.MkdirAll(filepath.Join(tmpDir, "memory"), 0o755) - os.MkdirAll(filepath.Join(tmpDir, "skills"), 0o755) - for _, name := range []string{"IDENTITY.md", "SOUL.md", "USER.md"} { os.WriteFile(filepath.Join(tmpDir, name), []byte(strings.Repeat("Content.\n", 10)), 0o644) } cb := NewContextBuilder(tmpDir) - history := []providers.Message{ {Role: "user", Content: "previous message"}, - {Role: "assistant", Content: "previous response"}, } b.ResetTimer() - for i := 0; i < b.N; i++ { _ = cb.BuildMessages(history, "summary", "new message", nil, "cli", "test") } diff --git a/pkg/agent/context_test.go b/pkg/agent/context_test.go index 6f6884bbe..5756ed911 100644 --- a/pkg/agent/context_test.go +++ b/pkg/agent/context_test.go @@ -12,11 +12,9 @@ func msg(role, content string) providers.Message { func assistantWithTools(toolIDs ...string) providers.Message { calls := make([]providers.ToolCall, len(toolIDs)) - for i, id := range toolIDs { calls[i] = providers.ToolCall{ID: id, Type: "function"} } - return providers.Message{Role: "assistant", ToolCalls: calls} } @@ -26,13 +24,11 @@ func toolResult(id string) providers.Message { func TestSanitizeHistoryForProvider_EmptyHistory(t *testing.T) { result := sanitizeHistoryForProvider(nil) - if len(result) != 0 { t.Fatalf("expected empty, got %d messages", len(result)) } result = sanitizeHistoryForProvider([]providers.Message{}) - if len(result) != 0 { t.Fatalf("expected empty, got %d messages", len(result)) } @@ -41,228 +37,170 @@ func TestSanitizeHistoryForProvider_EmptyHistory(t *testing.T) { func TestSanitizeHistoryForProvider_SingleToolCall(t *testing.T) { history := []providers.Message{ msg("user", "hello"), - assistantWithTools("A"), - toolResult("A"), - msg("assistant", "done"), } result := sanitizeHistoryForProvider(history) - if len(result) != 4 { t.Fatalf("expected 4 messages, got %d", len(result)) } - assertRoles(t, result, "user", "assistant", "tool", "assistant") } func TestSanitizeHistoryForProvider_MultiToolCalls(t *testing.T) { history := []providers.Message{ msg("user", "do two things"), - assistantWithTools("A", "B"), - toolResult("A"), - toolResult("B"), - msg("assistant", "both done"), } result := sanitizeHistoryForProvider(history) - if len(result) != 5 { t.Fatalf("expected 5 messages, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant") } func TestSanitizeHistoryForProvider_AssistantToolCallAfterPlainAssistant(t *testing.T) { history := []providers.Message{ msg("user", "hi"), - msg("assistant", "thinking"), - assistantWithTools("A"), - toolResult("A"), } result := sanitizeHistoryForProvider(history) - if len(result) != 2 { t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user", "assistant") } func TestSanitizeHistoryForProvider_OrphanedLeadingTool(t *testing.T) { history := []providers.Message{ toolResult("A"), - msg("user", "hello"), } result := sanitizeHistoryForProvider(history) - if len(result) != 1 { t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user") } func TestSanitizeHistoryForProvider_ToolAfterUserDropped(t *testing.T) { history := []providers.Message{ msg("user", "hello"), - toolResult("A"), } result := sanitizeHistoryForProvider(history) - if len(result) != 1 { t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user") } func TestSanitizeHistoryForProvider_ToolAfterAssistantNoToolCalls(t *testing.T) { history := []providers.Message{ msg("user", "hello"), - msg("assistant", "hi"), - toolResult("A"), } result := sanitizeHistoryForProvider(history) - if len(result) != 2 { t.Fatalf("expected 2 messages, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user", "assistant") } func TestSanitizeHistoryForProvider_AssistantToolCallAtStart(t *testing.T) { history := []providers.Message{ assistantWithTools("A"), - toolResult("A"), - msg("user", "hello"), } result := sanitizeHistoryForProvider(history) - if len(result) != 1 { t.Fatalf("expected 1 message, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user") } func TestSanitizeHistoryForProvider_MultiToolCallsThenNewRound(t *testing.T) { history := []providers.Message{ msg("user", "do two things"), - assistantWithTools("A", "B"), - toolResult("A"), - toolResult("B"), - msg("assistant", "done"), - msg("user", "hi"), - assistantWithTools("C"), - toolResult("C"), - msg("assistant", "done again"), } result := sanitizeHistoryForProvider(history) - if len(result) != 9 { t.Fatalf("expected 9 messages, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "user", "assistant", "tool", "assistant") } func TestSanitizeHistoryForProvider_ConsecutiveMultiToolRounds(t *testing.T) { history := []providers.Message{ msg("user", "start"), - assistantWithTools("A", "B"), - toolResult("A"), - toolResult("B"), - assistantWithTools("C", "D"), - toolResult("C"), - toolResult("D"), - msg("assistant", "all done"), } result := sanitizeHistoryForProvider(history) - if len(result) != 8 { t.Fatalf("expected 8 messages, got %d: %+v", len(result), roles(result)) } - assertRoles(t, result, "user", "assistant", "tool", "tool", "assistant", "tool", "tool", "assistant") } func TestSanitizeHistoryForProvider_PlainConversation(t *testing.T) { history := []providers.Message{ msg("user", "hello"), - msg("assistant", "hi"), - msg("user", "how are you"), - msg("assistant", "fine"), } result := sanitizeHistoryForProvider(history) - if len(result) != 4 { t.Fatalf("expected 4 messages, got %d", len(result)) } - assertRoles(t, result, "user", "assistant", "user", "assistant") } func roles(msgs []providers.Message) []string { r := make([]string, len(msgs)) - for i, m := range msgs { r[i] = m.Role } - return r } func assertRoles(t *testing.T, msgs []providers.Message, expected ...string) { t.Helper() - if len(msgs) != len(expected) { t.Fatalf("role count mismatch: got %v, want %v", roles(msgs), expected) } - for i, exp := range expected { if msgs[i].Role != exp { t.Errorf("message[%d]: got role %q, want %q", i, msgs[i].Role, exp) diff --git a/pkg/agent/instance_ext_test.go b/pkg/agent/instance_ext_test.go new file mode 100644 index 000000000..27799eda2 --- /dev/null +++ b/pkg/agent/instance_ext_test.go @@ -0,0 +1,52 @@ +package agent + +import ( + "github.com/sipeed/picoclaw/pkg/config" + "os" + "testing" +) + +func TestNewAgentInstance_ResolveCandidatesFromModelListAliasWithoutProtocol(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-instance-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "glm-5", + }, + }, + + ModelList: []config.ModelConfig{ + { + ModelName: "glm-5", + + Model: "glm-5", + + APIBase: "https://api.z.ai/api/coding/paas/v4", + }, + }, + } + + provider := &mockProvider{} + + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) + + if len(agent.Candidates) != 1 { + t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) + } + + if agent.Candidates[0].Provider != "openai" { + t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openai") + } + + if agent.Candidates[0].Model != "glm-5" { + t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "glm-5") + } +} diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index 9762d0a3c..4f41ecd1c 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -12,35 +12,28 @@ func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 1234, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 1234, MaxToolIterations: 5, }, }, } configuredTemp := 1.0 - cfg.Agents.Defaults.Temperature = &configuredTemp provider := &mockProvider{} - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if agent.MaxTokens != 1234 { t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234) } - if agent.Temperature != 1.0 { t.Fatalf("Temperature = %f, want %f", agent.Temperature, 1.0) } @@ -51,29 +44,23 @@ func TestNewAgentInstance_DefaultsTemperatureWhenZero(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 1234, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 1234, MaxToolIterations: 5, }, }, } configuredTemp := 0.0 - cfg.Agents.Defaults.Temperature = &configuredTemp provider := &mockProvider{} - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if agent.Temperature != 0.0 { @@ -86,25 +73,20 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 1234, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 1234, MaxToolIterations: 5, }, }, } provider := &mockProvider{} - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) if agent.Temperature != 0.7 { @@ -113,91 +95,68 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) { } func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-instance-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "step-3.5-flash", - }, + tests := []struct { + name string + aliasName string + modelName string + apiBase string + wantProvider string + wantModel string + }{ + { + name: "alias with provider prefix", + aliasName: "step-3.5-flash", + modelName: "openrouter/stepfun/step-3.5-flash:free", + apiBase: "https://openrouter.ai/api/v1", + wantProvider: "openrouter", + wantModel: "stepfun/step-3.5-flash:free", }, - - ModelList: []config.ModelConfig{ - { - ModelName: "step-3.5-flash", - - Model: "openrouter/stepfun/step-3.5-flash:free", - - APIBase: "https://openrouter.ai/api/v1", - }, + { + name: "alias without provider prefix", + aliasName: "glm-5", + modelName: "glm-5", + apiBase: "https://api.z.ai/api/coding/paas/v4", + wantProvider: "openai", + wantModel: "glm-5", }, } - provider := &mockProvider{} + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-instance-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: tt.aliasName, + }, + }, + ModelList: []config.ModelConfig{ + { + ModelName: tt.aliasName, + Model: tt.modelName, + APIBase: tt.apiBase, + }, + }, + } - if len(agent.Candidates) != 1 { - t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) - } + provider := &mockProvider{} + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) - if agent.Candidates[0].Provider != "openrouter" { - t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openrouter") - } - - if agent.Candidates[0].Model != "stepfun/step-3.5-flash:free" { - t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "stepfun/step-3.5-flash:free") - } -} - -func TestNewAgentInstance_ResolveCandidatesFromModelListAliasWithoutProtocol(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-instance-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "glm-5", - }, - }, - - ModelList: []config.ModelConfig{ - { - ModelName: "glm-5", - - Model: "glm-5", - - APIBase: "https://api.z.ai/api/coding/paas/v4", - }, - }, - } - - provider := &mockProvider{} - - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) - - if len(agent.Candidates) != 1 { - t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) - } - - if agent.Candidates[0].Provider != "openai" { - t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openai") - } - - if agent.Candidates[0].Model != "glm-5" { - t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "glm-5") + if len(agent.Candidates) != 1 { + t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) + } + if agent.Candidates[0].Provider != tt.wantProvider { + t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, tt.wantProvider) + } + if agent.Candidates[0].Model != tt.wantModel { + t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, tt.wantModel) + } + }) } } diff --git a/pkg/agent/loop_ext_test.go b/pkg/agent/loop_ext_test.go new file mode 100644 index 000000000..a26a3fa24 --- /dev/null +++ b/pkg/agent/loop_ext_test.go @@ -0,0 +1,2918 @@ +package agent + +import ( + "context" + "fmt" + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" + "github.com/sipeed/picoclaw/pkg/tools" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" +) + +func TestRecordLastHeartbeatTarget(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + + provider := &mockProvider{} + + al := NewAgentLoop(cfg, msgBus, provider) + + target := "telegram:-100123/42" + + if err := al.RecordLastHeartbeatTarget(target); err != nil { + t.Fatalf("RecordLastHeartbeatTarget failed: %v", err) + } + + if got := al.state.GetLastHeartbeatTarget(); got != target { + t.Fatalf("GetLastHeartbeatTarget = %q, want %q", got, target) + } +} + +type mockContextualTool struct { + lastChannel string + + lastChatID string +} + +func (m *mockContextualTool) SetContext(channel, chatID string) { + m.lastChannel = channel + + m.lastChatID = chatID +} + +func TestShouldInjectReminder(t *testing.T) { + tests := []struct { + name string + + iteration int + + interval int + + want bool + }{ + {"first iteration skipped", 1, 5, false}, + + {"iteration 5 interval 5", 5, 5, true}, + + {"iteration 10 interval 5", 10, 5, true}, + + {"iteration 3 interval 5", 3, 5, false}, + + {"interval zero disabled", 5, 0, false}, + + {"interval negative disabled", 5, -1, false}, + + {"iteration 2 interval 1", 2, 1, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := shouldInjectReminder(tt.iteration, tt.interval) + + if got != tt.want { + t.Errorf("shouldInjectReminder(%d, %d) = %v, want %v", tt.iteration, tt.interval, got, tt.want) + } + }) + } +} + +func TestBuildTaskReminder_WithoutBlocker(t *testing.T) { + msg := buildTaskReminder("implement feature X", "") + + if msg.Role != "user" { + t.Errorf("expected role 'user', got %q", msg.Role) + } + + if !strings.Contains(msg.Content, "[TASK REMINDER]") { + t.Error("expected content to contain '[TASK REMINDER]'") + } + + if !strings.Contains(msg.Content, "implement feature X") { + t.Error("expected content to contain original message") + } + + if strings.Contains(msg.Content, "blocker") { + t.Error("expected content NOT to contain 'blocker' when no blocker provided") + } + + if !strings.Contains(msg.Content, "move on") { + t.Error("expected content to contain completion prompt") + } +} + +func TestBuildTaskReminder_WithBlocker(t *testing.T) { + msg := buildTaskReminder("implement feature X", "ModuleNotFoundError: No module named 'foo'") + + if msg.Role != "user" { + t.Errorf("expected role 'user', got %q", msg.Role) + } + + if !strings.Contains(msg.Content, "[TASK REMINDER]") { + t.Error("expected content to contain '[TASK REMINDER]'") + } + + if !strings.Contains(msg.Content, "implement feature X") { + t.Error("expected content to contain original message") + } + + if !strings.Contains(msg.Content, "Last blocker") { + t.Error("expected content to contain 'Last blocker'") + } + + if !strings.Contains(msg.Content, "ModuleNotFoundError") { + t.Error("expected content to contain blocker text") + } +} + +func TestResolveProvider_CachesProviders(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + Provider: "vllm", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + + Providers: config.ProvidersConfig{ + VLLM: config.ProviderConfig{ + APIKey: "test-key", + + APIBase: "https://example.com/v1", + }, + }, + } + + msgBus := bus.NewMessageBus() + + primary := &mockProvider{} + + al := NewAgentLoop(cfg, msgBus, primary) + + p1 := al.resolveProvider("vllm", "test-model", primary) + + if p1 == primary { + t.Fatal("expected a new provider from legacy providers config, not the fallback") + } + + p2 := al.resolveProvider("vllm", "test-model", primary) + + if p1 != p2 { + t.Fatal("expected same cached instance on second call") + } +} + +func TestResolveProvider_FallsBackOnError(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + Provider: "vllm", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + + primary := &mockProvider{} + + al := NewAgentLoop(cfg, msgBus, primary) + + p := al.resolveProvider("nonexistent", "unknown-model", primary) + + if p != primary { + t.Fatal("expected fallback to primary provider on creation error") + } + + if _, ok := al.providerCache["nonexistent"]; ok { + t.Fatal("failed provider should not be cached") + } +} + +func TestResolveProvider_EmptyNameReturnsFallback(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + + primary := &mockProvider{} + + al := NewAgentLoop(cfg, msgBus, primary) + + p := al.resolveProvider("", "", primary) + + if p != primary { + t.Fatal("expected fallback provider for empty name") + } +} + +func TestSlashCommandResponseSkipsPlaceholder(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + + provider := &mockProvider{} + + al := NewAgentLoop(cfg, msgBus, provider) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + + defer cancel() + + go func() { + _ = al.Run(ctx) + }() + + msgBus.PublishInbound(context.Background(), bus.InboundMessage{ + Channel: "telegram", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "/skills", + }) + + outMsg, ok := msgBus.SubscribeOutbound(ctx) + + if !ok { + t.Fatal("expected outbound message from slash command") + } + + if !outMsg.SkipPlaceholder { + t.Errorf("expected SkipPlaceholder=true for slash command response, got false") + } +} + +func TestBuildTaskReminder_Truncation(t *testing.T) { + + longMsg := strings.Repeat("あ", 1000) + + longBlocker := strings.Repeat("X", 500) + + msg := buildTaskReminder(longMsg, longBlocker) + + runeCount := strings.Count(msg.Content, "あ") + + if runeCount >= 1000 { + t.Errorf("expected task message to be truncated, got %d 'あ' runes", runeCount) + } + + if runeCount > taskReminderMaxChars { + t.Errorf("expected at most %d task runes, got %d", taskReminderMaxChars, runeCount) + } + + xCount := strings.Count(msg.Content, "X") + + if xCount >= 500 { + t.Errorf("expected blocker to be truncated, got %d 'X' chars", xCount) + } + + if xCount > blockerMaxChars { + t.Errorf("expected at most %d blocker chars, got %d", blockerMaxChars, xCount) + } +} + +func TestBuildPlanReminder(t *testing.T) { + tests := []struct { + name string + + status string + + wantOK bool + + wantSubstr string + }{ + {"interviewing", "interviewing", true, "interviewing the user"}, + + {"review", "review", true, "under review"}, + + {"executing returns false", "executing", false, ""}, + + {"empty returns false", "", false, ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + msg, ok := buildPlanReminder(tt.status) + + if ok != tt.wantOK { + t.Fatalf("buildPlanReminder(%q) ok = %v, want %v", tt.status, ok, tt.wantOK) + } + + if !ok { + return + } + + if msg.Role != "user" { + t.Errorf("expected role 'user', got %q", msg.Role) + } + + if !strings.Contains(msg.Content, tt.wantSubstr) { + t.Errorf("expected content to contain %q, got %q", tt.wantSubstr, msg.Content) + } + }) + } +} + +func TestPlanCommand_ShowNoPlan(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + response, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan"}) + + if !handled { + t.Fatal("expected /plan to be handled") + } + + if !strings.Contains(response, "No active plan") { + t.Errorf("expected 'No active plan', got %q", response) + } +} + +func TestSplitChatAndThread(t *testing.T) { + tests := []struct { + name string + + chatID string + + wantChatID string + + wantThread int + }{ + {name: "plain chat", chatID: "-100123", wantChatID: "-100123", wantThread: 0}, + + {name: "chat with thread", chatID: "-100123/77", wantChatID: "-100123", wantThread: 77}, + + {name: "invalid thread", chatID: "-100123/abc", wantChatID: "-100123", wantThread: 0}, + + {name: "empty", chatID: "", wantChatID: "", wantThread: 0}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gotChatID, gotThread := splitChatAndThread(tt.chatID) + + if gotChatID != tt.wantChatID || gotThread != tt.wantThread { + t.Fatalf( + + "splitChatAndThread(%q) = (%q, %d), want (%q, %d)", + + tt.chatID, + + gotChatID, + + gotThread, + + tt.wantChatID, + + tt.wantThread, + ) + } + }) + } +} + +func TestHeartbeatCommandThreadHerePersistsConfig(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + var saved bool + + var updatedThread int + + al.SetConfigSaver(func(cfg *config.Config) error { + saved = true + + if cfg.Channels.Telegram.HeartbeatThreadID != 42 { + t.Fatalf("HeartbeatThreadID in saver = %d, want 42", cfg.Channels.Telegram.HeartbeatThreadID) + } + + return nil + }) + + al.SetHeartbeatThreadUpdater(func(threadID int) { updatedThread = threadID }) + + msg := bus.InboundMessage{ + Content: "/heartbeat thread here", + + Channel: "telegram", + + ChatID: "-100500/42", + } + + resp, handled := al.handleCommand(context.Background(), msg) + + if !handled { + t.Fatal("expected /heartbeat command to be handled") + } + + if !strings.Contains(resp, "Heartbeat thread set to 42") { + t.Fatalf("unexpected response: %q", resp) + } + + if !saved { + t.Fatal("expected config saver to be called") + } + + if updatedThread != 42 { + t.Fatalf("updatedThread = %d, want 42", updatedThread) + } + + if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 42 { + t.Fatalf("cfg heartbeat thread = %d, want 42", got) + } + + if got := al.state.GetHeartbeatTarget(); got != "telegram:-100500" { + t.Fatalf("state heartbeat target = %q, want %q", got, "telegram:-100500") + } +} + +func TestHeartbeatCommandThreadOff(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + al.cfg.Channels.Telegram.HeartbeatThreadID = 99 + + resp, handled := al.handleCommand(context.Background(), bus.InboundMessage{ + Content: "/heartbeat thread off", + + Channel: "telegram", + + ChatID: "-100500/42", + }) + + if !handled { + t.Fatal("expected /heartbeat command to be handled") + } + + if !strings.Contains(resp, "disabled") { + t.Fatalf("unexpected response: %q", resp) + } + + if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 0 { + t.Fatalf("cfg heartbeat thread = %d, want 0", got) + } +} + +func TestPlanCommand_StartNewPlan(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + _, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan Set up monitoring"}) + + if handled { + t.Fatal("expected /plan NOT to be handled (should fall through to LLM)") + } + + msg := bus.InboundMessage{Content: "/plan Set up monitoring"} + + expanded, compact, ok := al.expandPlanCommand(msg) + + if !ok { + t.Fatal("expected expandPlanCommand to succeed") + } + + if expanded != "Set up monitoring" { + t.Errorf("expected expanded = 'Set up monitoring', got %q", expanded) + } + + if !strings.Contains(compact, "Set up monitoring") { + t.Errorf("expected compact to contain task, got %q", compact) + } + + agent := al.registry.GetDefaultAgent() + + if !agent.ContextBuilder.HasActivePlan() { + t.Error("expected active plan after expandPlanCommand") + } + + if status := agent.ContextBuilder.GetPlanStatus(); status != "interviewing" { + t.Errorf("expected 'interviewing', got %q", status) + } +} + +func TestPlanCommand_StartBlockedByExisting(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + al.expandPlanCommand(bus.InboundMessage{Content: "/plan First task"}) + + response, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan Second task"}) + + if !handled { + t.Fatal("expected second /plan to be handled (blocked)") + } + + if !strings.Contains(response, "already active") { + t.Errorf("expected 'already active', got %q", response) + } +} + +func TestPlanCommand_Clear(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan clear"}) + + if !strings.Contains(response, "Plan cleared") { + t.Errorf("expected 'Plan cleared', got %q", response) + } + + agent := al.registry.GetDefaultAgent() + + if agent.ContextBuilder.HasActivePlan() { + t.Error("expected no plan after clear") + } +} + +func TestPlanCommand_ClearNoPlan(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan clear"}) + + if !strings.Contains(response, "No active plan") { + t.Errorf("expected 'No active plan', got %q", response) + } +} + +func TestPlanCommand_Start(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := "# Active Plan\n\n> Task: Test task\n> Status: interviewing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + + _ = agent.ContextBuilder.WriteMemory(plan) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + + if !strings.Contains(response, "approved") { + t.Errorf("expected 'approved', got %q", response) + } + + if status := agent.ContextBuilder.GetPlanStatus(); status != "executing" { + t.Errorf("expected 'executing', got %q", status) + } + + if !al.planStartPending { + t.Error("expected planStartPending to be true after /plan start") + } +} + +func TestPlanCommand_StartFromReview(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + + _ = agent.ContextBuilder.WriteMemory(plan) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + + if !strings.Contains(response, "approved") { + t.Errorf("expected 'approved', got %q", response) + } + + if status := agent.ContextBuilder.GetPlanStatus(); status != "executing" { + t.Errorf("expected 'executing', got %q", status) + } + + if !al.planStartPending { + t.Error("expected planStartPending to be true after /plan start from review") + } +} + +func TestPlanCommand_StartNoPhases(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + + if !strings.Contains(response, "no phases") { + t.Errorf("expected 'no phases' error, got %q", response) + } + + agent := al.registry.GetDefaultAgent() + + if status := agent.ContextBuilder.GetPlanStatus(); status != "interviewing" { + t.Errorf("expected status to remain 'interviewing', got %q", status) + } + + if al.planStartPending { + t.Error("planStartPending must not be set when start is rejected (no phases)") + } +} + +func TestPlanCommand_StartAlreadyExecuting(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := "# Active Plan\n\n> Task: Test task\n> Status: interviewing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + + _ = agent.ContextBuilder.WriteMemory(plan) + + al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + + al.planStartPending = false + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) + + if !strings.Contains(response, "already executing") { + t.Errorf("expected 'already executing', got %q", response) + } + + if al.planStartPending { + t.Error("planStartPending must not be set when plan is already executing") + } +} + +func TestPlanCommand_Done(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := `# Active Plan + + + +> Task: Test task + +> Status: executing + +> Phase: 1 + + + +## Phase 1: Setup + +- [ ] Step one + +- [ ] Step two + + + +## Context + +Test context + +` + + agent.ContextBuilder.WriteMemory(plan) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan done 1"}) + + if !strings.Contains(response, "Marked step 1") { + t.Errorf("expected confirmation, got %q", response) + } +} + +func TestPlanCommand_DoneInvalidStep(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan done abc"}) + + if !strings.Contains(response, "positive integer") { + t.Errorf("expected step validation error, got %q", response) + } +} + +func TestPlanCommand_Add(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := `# Active Plan + + + +> Task: Test task + +> Status: executing + +> Phase: 1 + + + +## Phase 1: Setup + +- [ ] Step one + + + +## Context + +Test context + +` + + agent.ContextBuilder.WriteMemory(plan) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan add New step here"}) + + if !strings.Contains(response, "Added step") { + t.Errorf("expected 'Added step', got %q", response) + } + + content := agent.ContextBuilder.ReadMemory() + + if !strings.Contains(content, "New step here") { + t.Error("expected new step in plan content") + } +} + +func TestPlanCommand_Next(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := `# Active Plan + + + +> Task: Test task + +> Status: executing + +> Phase: 1 + + + +## Phase 1: Setup + +- [x] Step one + + + +## Phase 2: Deploy + +- [ ] Step two + + + +## Context + +Test + +` + + agent.ContextBuilder.WriteMemory(plan) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan next"}) + + if !strings.Contains(response, "phase 2") { + t.Errorf("expected 'phase 2', got %q", response) + } + + if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 2 { + t.Errorf("expected phase 2, got %d", phase) + } +} + +func TestPlanCommand_ShowActivePlan(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := `# Active Plan + + + +> Task: Deploy app + +> Status: executing + +> Phase: 1 + + + +## Phase 1: Build + +- [x] Compile code + +- [ ] Run tests + + + +## Context + +Production server + +` + + agent.ContextBuilder.WriteMemory(plan) + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan"}) + + if !strings.Contains(response, "Deploy app") { + t.Errorf("expected task name in display, got %q", response) + } + + if !strings.Contains(response, "Phase 1") { + t.Errorf("expected phase info in display, got %q", response) + } +} + +func TestAutoPhaseAdvance(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-auto-advance-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + + provider := &simpleMockProvider{response: "OK"} + + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + + if agent == nil { + t.Fatal("No default agent") + } + + plan := `# Active Plan + + + +> Task: Test auto advance + +> Status: executing + +> Phase: 1 + + + +## Phase 1: Setup + +- [x] Step one + +- [x] Step two + + + +## Phase 2: Deploy + +- [ ] Step three + + + +## Context + +Test + +` + + agent.ContextBuilder.WriteMemory(plan) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + _, err = al.ProcessDirectWithChannel(ctx, "continue", "auto-advance-test", "test", "chat1") + if err != nil { + t.Fatalf("ProcessDirectWithChannel failed: %v", err) + } + + if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 2 { + t.Errorf("expected phase auto-advanced to 2, got %d", phase) + } +} + +func TestAutoCompleteClears(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-auto-complete-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + + provider := &simpleMockProvider{response: "All done"} + + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + + if agent == nil { + t.Fatal("No default agent") + } + + plan := `# Active Plan + + + +> Task: Test auto complete + +> Status: executing + +> Phase: 1 + + + +## Phase 1: Setup + +- [x] Step one + +- [x] Step two + + + +## Context + +Test + +` + + agent.ContextBuilder.WriteMemory(plan) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + _, err = al.ProcessDirectWithChannel(ctx, "finish up", "auto-complete-test", "test", "chat1") + if err != nil { + t.Fatalf("ProcessDirectWithChannel failed: %v", err) + } + + if !agent.ContextBuilder.HasActivePlan() { + t.Error("expected plan to be retained after completion") + } + + if status := agent.ContextBuilder.GetPlanStatus(); status != "completed" { + t.Errorf("expected plan status 'completed', got %q", status) + } + + if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 1 { + t.Errorf("expected phase 1 (total phases), got %d", phase) + } +} + +func TestIsToolAllowedDuringInterview_FuzzyNames(t *testing.T) { + tests := []struct { + name string + + args map[string]any + + want bool + }{ + + {"read_file", nil, true}, + + {"list_dir", nil, true}, + + {"web_search", nil, true}, + + {"web_fetch", nil, true}, + + {"readfile", nil, true}, + + {"ReadFile", nil, true}, + + {"listdir", nil, true}, + + {"websearch", nil, true}, + + {"webfetch", nil, true}, + + {"message", nil, true}, + + {"Message", nil, true}, + + {"edit_file", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, + + {"editfile", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, + + {"EditFile", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, + + {"edit_file", map[string]any{"path": "/ws/main.go"}, false}, + + {"editfile", map[string]any{"path": "/ws/main.go"}, false}, + + {"exec", map[string]any{"command": "find . -name '*.py'"}, true}, + + {"exec", map[string]any{"command": "ls -la"}, true}, + + {"exec", map[string]any{"command": "grep -r TODO ."}, true}, + + {"exec", map[string]any{"command": "cat README.md"}, true}, + + {"exec", map[string]any{"command": "cd /home/user/project && find . -type f"}, true}, + + {"exec", map[string]any{"command": "cd /tmp && rm -rf *"}, false}, + + {"exec", map[string]any{"command": "find . > output.txt"}, false}, + + {"exec", map[string]any{"command": "ls -la >> log.txt"}, false}, + + {"exec", map[string]any{"command": "cat foo | tee bar.txt"}, false}, + + {"exec", map[string]any{"command": "cat ../../etc/passwd"}, false}, + + {"exec", map[string]any{"command": "find ../../"}, false}, + + {"exec", map[string]any{"command": "ls ../secret"}, false}, + + {"exec", map[string]any{"command": "cat /etc/passwd"}, false}, + + {"exec", map[string]any{"command": "find /etc -name '*.conf'"}, false}, + + {"exec", map[string]any{"command": "ls /root"}, false}, + + {"exec", map[string]any{"command": "rm -rf /"}, false}, + + {"exec", map[string]any{"command": "mv a b"}, false}, + + {"exec", nil, false}, + + {"exec", map[string]any{"command": ""}, false}, + + {"Exec", nil, false}, + } + + for _, tt := range tests { + got := isToolAllowedDuringInterview(tt.name, tt.args) + + if got != tt.want { + t.Errorf("isToolAllowedDuringInterview(%q, %v) = %v, want %v", tt.name, tt.args, got, tt.want) + } + } +} + +func TestBuildArgsSnippet_ExecStripsCD(t *testing.T) { + tests := []struct { + name string + + tool string + + args map[string]any + + workspace string + + wantSnip string + }{ + { + name: "exec strips cd prefix", + + tool: "exec", + + args: map[string]any{ + "command": "cd /home/user/workspace/project/my-projects && pytest tests/test_integration.py", + }, + + workspace: "/home/user/workspace", + + wantSnip: "pytest tests/test_integration.py", + }, + + { + name: "exec no cd prefix, flags stripped", + + tool: "exec", + + args: map[string]any{"command": "ls -la"}, + + workspace: "/ws", + + wantSnip: "ls", + }, + + { + name: "exec empty command", + + tool: "exec", + + args: map[string]any{}, + + workspace: "/ws", + + wantSnip: "{}", + }, + + { + name: "read_file strips workspace", + + tool: "read_file", + + args: map[string]any{"path": "/home/user/workspace/src/main.go"}, + + workspace: "/home/user/workspace", + + wantSnip: "src/main.go", + }, + + { + name: "edit_file shows path", + + tool: "edit_file", + + args: map[string]any{"path": "/ws/config.json", "old_text": "old value here"}, + + workspace: "/ws", + + wantSnip: "config.json", + }, + + { + name: "file tool long path prioritizes filename", + + tool: "read_file", + + args: map[string]any{ + "path": "/ws/projects/terra-py-form/src/terra_py_form/hot/state/backend.py", + }, + + workspace: "/ws", + + wantSnip: "projects/terra-py-form/src/terra_py_form/hot/sta\u2026/backend.py", + }, + + { + name: "unknown tool shows raw JSON", + + tool: "web_search", + + args: map[string]any{"query": "hello"}, + + workspace: "/ws", + + wantSnip: `{"query":"hello"}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := buildArgsSnippet(tt.tool, tt.args, tt.workspace) + + if got != tt.wantSnip { + t.Errorf("buildArgsSnippet(%q) = %q, want %q", tt.tool, got, tt.wantSnip) + } + }) + } +} + +func TestFormatCompactEntry(t *testing.T) { + tests := []struct { + name string + + entry toolLogEntry + + wantSub string // must be a substring + + wantMark string // result marker must appear + + noTime bool // if true, duration should NOT appear + }{ + { + name: "exec short entry", + + entry: toolLogEntry{Name: "[1] exec", ArgsSnip: "ls", Result: "✓ 1.0s"}, + + wantSub: "exec ls", + + wantMark: "✓ 1.0s", + }, + + { + name: "exec long entry truncated from end", + + entry: toolLogEntry{ + Name: "[2] exec", + + ArgsSnip: "pytest tests/integration/test_very_long_name.py", + + Result: "✗ 3.0s", + }, + + wantMark: "✗", + }, + + { + name: "file tool omits duration, shows filename", + + entry: toolLogEntry{ + Name: "[3] edit_file", + + ArgsSnip: "projects/terra/src/deep/nested/backend.py", + + Result: "✓ 0.0s", + }, + + wantSub: "backend.py", + + wantMark: "✓", + + noTime: true, + }, + + { + name: "file tool path truncates from start", + + entry: toolLogEntry{ + Name: "[4] read_file", + + ArgsSnip: "projects/terra-py-form/src/terra_py_form/hot/state/backend.py", + + Result: "✓ 0.1s", + }, + + wantSub: "backend.py", + + wantMark: "✓", + + noTime: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := formatCompactEntry(tt.entry) + + if tt.wantSub != "" && !strings.Contains(got, tt.wantSub) { + t.Errorf("expected to contain %q, got: %q", tt.wantSub, got) + } + + if !strings.Contains(got, tt.wantMark) { + t.Errorf("result marker %q missing from: %q", tt.wantMark, got) + } + + if tt.noTime && strings.Contains(got, "0s") { + t.Errorf("file tool should omit duration, got: %q", got) + } + + if runeLen := len([]rune(got)); runeLen > maxEntryLineWidth { + t.Errorf("entry too wide: %d runes (max %d): %q", runeLen, maxEntryLineWidth, got) + } + }) + } +} + +func TestBuildRichStatus(t *testing.T) { + task := &activeTask{ + Iteration: 3, + + MaxIter: 20, + + toolLog: []toolLogEntry{ + {Name: "exec", ArgsSnip: "ls -la", Result: "✓ 1.2s"}, + + {Name: "exec", ArgsSnip: "pytest tests/", Result: "✓ 5.0s"}, + + {Name: "read_file", ArgsSnip: "src/main.go", Result: "⏳"}, + }, + } + + got := buildRichStatus(task, false, "/home/user/my-projects") + + mustContain := []string{ + "Task in progress (3/20)", + + "my-projects", + + "read_file", + + "No errors", + } + + for _, s := range mustContain { + if !strings.Contains(got, s) { + t.Errorf("expected output to contain %q, got:\n%s", s, got) + } + } + + if strings.Contains(got, "Reply to intervene") { + t.Error("non-background task should not have reply prompt") + } + + bgGot := buildRichStatus(task, true, "/home/user/my-projects") + + if !strings.Contains(bgGot, "Reply to intervene") { + t.Error("background task should have reply prompt") + } +} + +func TestBuildRichStatus_ProjectDir(t *testing.T) { + + task := &activeTask{ + Iteration: 1, + + MaxIter: 10, + + projectDir: "terra-py-form", + + toolLog: []toolLogEntry{ + {Name: "exec", ArgsSnip: "ls", Result: "✓ 0.1s"}, + }, + } + + got := buildRichStatus(task, false, "/home/user/.picoclaw/workspace") + + if !strings.Contains(got, "terra-py-form") { + t.Errorf("expected projectDir in output, got:\n%s", got) + } + + task2 := &activeTask{ + Iteration: 1, + + MaxIter: 10, + + fileCommonDir: "projects/terra-py-form", + + toolLog: []toolLogEntry{ + {Name: "read_file", ArgsSnip: "src/main.py", Result: "✓ 0.1s"}, + }, + } + + got2 := buildRichStatus(task2, false, "/home/user/.picoclaw/workspace") + + if !strings.Contains(got2, "terra-py-form") { + t.Errorf("expected fileCommonDir basename in output, got:\n%s", got2) + } + + task3 := &activeTask{ + Iteration: 1, + + MaxIter: 10, + + toolLog: []toolLogEntry{ + {Name: "exec", ArgsSnip: "ls", Result: "✓ 0.1s"}, + }, + } + + for _, ws := range []string{"/home/user/my-project/", "/home/user/my-project"} { + got := buildRichStatus(task3, false, ws) + + if !strings.Contains(got, "my-project") { + t.Errorf("workspace %q: expected 'my-project' in output, got:\n%s", ws, got) + } + } +} + +func TestExtractExecProjectDir(t *testing.T) { + tests := []struct { + name string + + cmd string + + want string + }{ + {"cd deep path", "cd /ws/projects/terra-py-form && pytest", "terra-py-form"}, + + {"cd direct subdir", "cd /ws/my-app && make build", "my-app"}, + + {"cd trailing slash", "cd /ws/my-app/ && ls", "my-app"}, + + {"cd to workspace", "cd /ws && ls", "ws"}, + + {"no cd prefix", "pytest tests/", ""}, + + {"empty command", "", ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + args := map[string]any{"command": tt.cmd} + + got := extractExecProjectDir(args) + + if got != tt.want { + t.Errorf("extractExecProjectDir(%q) = %q, want %q", tt.cmd, got, tt.want) + } + }) + } +} + +func TestFileParentRelDir(t *testing.T) { + ws := "/home/user/.picoclaw/workspace" + + tests := []struct { + name string + + path string + + want string + }{ + {"deep path", ws + "/projects/terra/src/main.py", "projects/terra/src"}, + + {"direct subdir", ws + "/my-app/README.md", "my-app"}, + + {"workspace root file", ws + "/notes.txt", ""}, + + {"outside workspace", "/tmp/foo.txt", ""}, + + {"trailing slash ws", ws + "/projects/terra/src/main.py", "projects/terra/src"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := fileParentRelDir(tt.path, ws) + + if got != tt.want { + t.Errorf("fileParentRelDir(%q, ws) = %q, want %q", tt.path, got, tt.want) + } + }) + } +} + +func TestCommonDirPrefix(t *testing.T) { + tests := []struct { + name string + + a, b string + + want string + }{ + {"same dir", "projects/terra/src", "projects/terra/src", "projects/terra/src"}, + + {"converge to project", "projects/terra/src", "projects/terra/tests", "projects/terra"}, + + {"converge to top", "projects/terra/src", "projects/other/tests", "projects"}, + + {"no common", "aaa/bbb", "ccc/ddd", ""}, + + {"one is prefix", "projects/terra", "projects/terra/src", "projects/terra"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := commonDirPrefix(tt.a, tt.b) + + if got != tt.want { + t.Errorf("commonDirPrefix(%q, %q) = %q, want %q", tt.a, tt.b, got, tt.want) + } + }) + } +} + +func TestDisplayProjectDir(t *testing.T) { + + task1 := &activeTask{projectDir: "my-app", fileCommonDir: "projects/other"} + + if got := displayProjectDir(task1); got != "my-app" { + t.Errorf("expected 'my-app', got %q", got) + } + + task2 := &activeTask{fileCommonDir: "projects/terra-py-form"} + + if got := displayProjectDir(task2); got != "terra-py-form" { + t.Errorf("expected 'terra-py-form', got %q", got) + } + + task3 := &activeTask{fileCommonDir: "my-app"} + + if got := displayProjectDir(task3); got != "my-app" { + t.Errorf("expected 'my-app', got %q", got) + } + + task4 := &activeTask{} + + if got := displayProjectDir(task4); got != "" { + t.Errorf("expected empty, got %q", got) + } +} + +func TestBuildRichStatus_FixedHeight(t *testing.T) { + + countLines := func(s string) int { + return strings.Count(s, "\n") + } + + task0 := &activeTask{Iteration: 1, MaxIter: 10} + + lines0 := countLines(buildRichStatus(task0, true, "/ws/p")) + + task1 := &activeTask{ + Iteration: 1, MaxIter: 10, + + toolLog: []toolLogEntry{{Name: "exec", ArgsSnip: "ls", Result: "⏳"}}, + } + + lines1 := countLines(buildRichStatus(task1, true, "/ws/p")) + + task5 := &activeTask{Iteration: 5, MaxIter: 10} + + for i := 0; i < 5; i++ { + task5.toolLog = append(task5.toolLog, toolLogEntry{ + Name: fmt.Sprintf("[%d] exec", i), ArgsSnip: "cmd", Result: "✓ 1.0s", + }) + } + + lines5 := countLines(buildRichStatus(task5, true, "/ws/p")) + + task5err := &activeTask{Iteration: 5, MaxIter: 10} + + for i := 0; i < 5; i++ { + task5err.toolLog = append(task5err.toolLog, toolLogEntry{ + Name: fmt.Sprintf("[%d] exec", i), ArgsSnip: "cmd", Result: "✓ 1.0s", + }) + } + + errEntry := toolLogEntry{ + Name: "[3] exec", ArgsSnip: "pytest", Result: "✗ 2.0s", + + ErrDetail: "FAILED test\nExit code: 1", + } + + task5err.lastError = &errEntry + + lines5err := countLines(buildRichStatus(task5err, true, "/ws/p")) + + if lines0 != lines1 || lines1 != lines5 || lines5 != lines5err { + t.Errorf("line counts should be equal: 0=%d, 1=%d, 5=%d, 5+err=%d", + + lines0, lines1, lines5, lines5err) + } +} + +func TestBuildRichStatus_StickyError(t *testing.T) { + + errEntry := toolLogEntry{ + Name: "[2] exec", ArgsSnip: "pytest", Result: "✗ 3.2s", + + ErrDetail: "FAILED test_login\nExit code: 1", + } + + task := &activeTask{ + Iteration: 5, + + MaxIter: 10, + + toolLog: []toolLogEntry{ + {Name: "[3] read_file", ArgsSnip: "src/auth.py", Result: "✓ 0.1s"}, + + {Name: "[4] edit_file", ArgsSnip: "src/auth.py", Result: "✓ 0.2s"}, + + {Name: "[5] exec", ArgsSnip: "pytest --retry", Result: "⏳"}, + }, + + lastError: &errEntry, + } + + got := buildRichStatus(task, false, "/ws/p") + + if !strings.Contains(got, "FAILED test_login") { + t.Errorf("expected sticky error detail in error section, got:\n%s", got) + } + + if !strings.Contains(got, "\u274C") { + t.Errorf("expected ❌ error header, got:\n%s", got) + } + + if !strings.Contains(got, "pytest --retry") { + t.Errorf("expected latest entry command, got:\n%s", got) + } +} + +func TestBuildRichStatus_LatestEntryNoInlineResult(t *testing.T) { + longCmd := "uv run pytest tests/hot/test_state_backend_integration.py" + + task := &activeTask{ + Iteration: 2, + + MaxIter: 10, + + toolLog: []toolLogEntry{ + {Name: "exec", ArgsSnip: "ls -la", Result: "\u2713 0.5s"}, + + {Name: "exec", ArgsSnip: longCmd, Result: "\u23F3"}, + }, + } + + got := buildRichStatus(task, false, "/ws/my-project") + + if !strings.Contains(got, "integration.py") { + t.Errorf("latest entry should show filename, got:\n%s", got) + } + + if !strings.Contains(got, " \u23F3") { + t.Errorf("latest entry result should be on indented line, got:\n%s", got) + } + + if !strings.Contains(got, "my-project") { + t.Errorf("should show project name, got:\n%s", got) + } + + lines := strings.Split(got, "\n") + + sepCount := 0 + + for _, l := range lines { + if strings.HasPrefix(l, "\u2501") { + sepCount++ + } + } + + if sepCount != 1 { + t.Errorf("expected exactly 1 separator, got %d in:\n%s", sepCount, got) + } +} + +func TestSanitizeHistoryForProvider_MultiToolCall(t *testing.T) { + + history := []providers.Message{ + {Role: "user", Content: "hello"}, + + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ + {ID: "a", Function: &providers.FunctionCall{Name: "exec"}}, + + {ID: "b", Function: &providers.FunctionCall{Name: "read_file"}}, + }}, + + {Role: "tool", Content: "ok", ToolCallID: "a"}, + + {Role: "tool", Content: "ok", ToolCallID: "b"}, + + {Role: "assistant", Content: "done"}, + } + + got := sanitizeHistoryForProvider(history) + + if len(got) != 5 { + roles := make([]string, len(got)) + + for i, m := range got { + roles[i] = m.Role + } + + t.Fatalf("expected 5 messages, got %d: %v", len(got), roles) + } + + toolCount := 0 + + for _, m := range got { + if m.Role == "tool" { + toolCount++ + } + } + + if toolCount != 2 { + t.Errorf("expected 2 tool results, got %d", toolCount) + } +} + +func TestPlanNudge_ForegroundExecution(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + provider := &countingMockProvider{} + + msgBus := bus.NewMessageBus() + + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + + if agent == nil { + t.Fatal("no default agent") + } + + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n\n## Context\n" + + agent.ContextBuilder.WriteMemory(plan) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + msg := bus.InboundMessage{ + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "continue working", + + SessionKey: "nudge-test", + } + + _, err = al.processMessage(ctx, msg) + if err != nil { + t.Fatalf("processMessage failed: %v", err) + } + + if provider.callCount < 2 { + t.Errorf("expected at least 2 provider calls (nudge should trigger continuation), got %d", provider.callCount) + } +} + +func TestPlanNudge_NoNudgeWhenAllStepsComplete(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + provider := &countingMockProvider{} + + msgBus := bus.NewMessageBus() + + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + + if agent == nil { + t.Fatal("no default agent") + } + + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [x] Step one\n- [x] Step two\n\n## Context\n" + + agent.ContextBuilder.WriteMemory(plan) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + msg := bus.InboundMessage{ + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "all done", + + SessionKey: "nudge-test-complete", + } + + _, err = al.processMessage(ctx, msg) + if err != nil { + t.Fatalf("processMessage failed: %v", err) + } + + if provider.callCount != 1 { + t.Errorf("expected exactly 1 provider call (no nudge needed), got %d", provider.callCount) + } +} + +func TestPlanNudge_ProgressMessage(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "test-model", + + MaxTokens: 4096, + + MaxToolIterations: 10, + }, + }, + } + + var nudgeContent string + + provider := &nudgeCaptureMockProvider{onSecondCall: func(msgs []providers.Message) { + + for i := len(msgs) - 1; i >= 0; i-- { + if msgs[i].Role == "user" { + nudgeContent = msgs[i].Content + + break + } + } + }} + + msgBus := bus.NewMessageBus() + + al := NewAgentLoop(cfg, msgBus, provider) + + agent := al.registry.GetDefaultAgent() + + if agent == nil { + t.Fatal("no default agent") + } + + plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n- [ ] Step three\n\n## Context\n" + + agent.ContextBuilder.WriteMemory(plan) + + provider.onFirstCall = func() { + updated := strings.Replace(agent.ContextBuilder.ReadMemory(), "- [ ] Step one", "- [x] Step one", 1) + + agent.ContextBuilder.WriteMemory(updated) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + + defer cancel() + + msg := bus.InboundMessage{ + Channel: "test", + + SenderID: "user1", + + ChatID: "chat1", + + Content: "work on the plan", + + SessionKey: "nudge-progress-test", + } + + _, err = al.processMessage(ctx, msg) + if err != nil { + t.Fatalf("processMessage failed: %v", err) + } + + if !strings.Contains(nudgeContent, "Progress recorded") { + t.Errorf("expected 'Progress recorded' nudge, got %q", nudgeContent) + } + + if !strings.Contains(nudgeContent, "2 unchecked steps remain") { + t.Errorf("expected '2 unchecked steps remain' in nudge, got %q", nudgeContent) + } +} + +type nudgeCaptureMockProvider struct { + callCount int + + onFirstCall func() + + onSecondCall func([]providers.Message) +} + +func TestConsumeStream_NormalCompletion(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} + + ch <- protocoltypes.StreamEvent{ContentDelta: "world!"} + + ch <- protocoltypes.StreamEvent{ + FinishReason: "stop", + + Usage: &providers.UsageInfo{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7}, + } + + close(ch) + }() + + ctx, cancel := context.WithCancel(context.Background()) + + defer cancel() + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if detected { + t.Fatal("expected detected=false for normal content") + } + + if resp.Content != "Hello world!" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") + } + + if resp.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") + } + + if resp.Usage == nil || resp.Usage.TotalTokens != 7 { + t.Errorf("Usage.TotalTokens = %v, want 7", resp.Usage) + } + + _ = ctx +} + +func TestConsumeStream_DetectsRepetition(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 64) + + cancelCalled := false + + ctx, cancel := context.WithCancel(context.Background()) + + wrappedCancel := func() { + cancelCalled = true + + cancel() + } + + repeatedChunk := strings.Repeat("abcdefghij", 50) + + go func() { + + for i := 0; i < 6; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} + } + + for i := 0; i < 10; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} + } + + close(ch) + }() + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if !detected { + t.Fatal("expected repetition detection to trigger") + } + + if !cancelCalled { + t.Error("expected cancelFn to be called") + } + + if len(resp.Content) >= 3000+10*len("more data") { + t.Errorf("Content length = %d, expected less than full output", len(resp.Content)) + } + + _ = ctx +} + +func TestConsumeStream_ToolCallAccumulation(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + + go func() { + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ID: "call_1", Name: "test_fn", ArgumentsDelta: `{"ke`}, + }, + } + + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ArgumentsDelta: `y":"val"}`}, + }, + } + + ch <- protocoltypes.StreamEvent{FinishReason: "tool_calls"} + + close(ch) + }() + + _, cancel := context.WithCancel(context.Background()) + + defer cancel() + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if detected { + t.Fatal("expected no repetition detection for tool calls") + } + + if len(resp.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls)) + } + + if resp.ToolCalls[0].Name != "test_fn" { + t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_fn") + } + + if resp.ToolCalls[0].Arguments["key"] != "val" { + t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "val") + } +} + +func TestConsumeStream_StreamError(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 4) + + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "partial"} + + ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("read error")} + + close(ch) + }() + + _, cancel := context.WithCancel(context.Background()) + + defer cancel() + + _, _, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) + + if err == nil { + t.Fatal("expected error, got nil") + } + + if !strings.Contains(err.Error(), "read error") { + t.Errorf("error = %q, want to contain %q", err.Error(), "read error") + } +} + +func TestConsumeStream_OnChunkCallback(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} + + ch <- protocoltypes.StreamEvent{ContentDelta: "world"} + + ch <- protocoltypes.StreamEvent{ContentDelta: "!"} + + ch <- protocoltypes.StreamEvent{FinishReason: "stop"} + + close(ch) + }() + + _, cancel := context.WithCancel(context.Background()) + + defer cancel() + + var chunks []string + + onChunk := func(accumulated, _ string) { + chunks = append(chunks, accumulated) + } + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, onChunk) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if detected { + t.Fatal("expected detected=false") + } + + if resp.Content != "Hello world!" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") + } + + if len(chunks) != 3 { + t.Fatalf("onChunk called %d times, want 3", len(chunks)) + } + + if chunks[0] != "Hello " { + t.Errorf("chunks[0] = %q, want %q", chunks[0], "Hello ") + } + + if chunks[1] != "Hello world" { + t.Errorf("chunks[1] = %q, want %q", chunks[1], "Hello world") + } + + if chunks[2] != "Hello world!" { + t.Errorf("chunks[2] = %q, want %q", chunks[2], "Hello world!") + } +} + +func TestConsumeStream_OnChunkWithRepetitionDetection(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 64) + + cancelCalled := false + + ctx, cancel := context.WithCancel(context.Background()) + + wrappedCancel := func() { + cancelCalled = true + + cancel() + } + + repeatedChunk := strings.Repeat("abcdefghij", 50) + + go func() { + for i := 0; i < 6; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} + } + + for i := 0; i < 10; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} + } + + close(ch) + }() + + var chunkCount int + + onChunk := func(_, _ string) { + chunkCount++ + } + + _, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, onChunk) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if !detected { + t.Fatal("expected repetition detection to trigger") + } + + if !cancelCalled { + t.Error("expected cancelFn to be called") + } + + if chunkCount == 0 { + t.Error("expected onChunk to be called at least once") + } + + _ = ctx +} + +type modelCapturingMockProvider struct { + mu sync.Mutex + + models []string + + response string +} + +func TestAgentLoop_PlanModel_UsedDuringInterviewing(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-planmodel-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "normal-model", + + PlanModel: "plan-model", + + MaxTokens: 4096, + + MaxToolIterations: 2, + }, + }, + } + + msgBus := bus.NewMessageBus() + + provider := &modelCapturingMockProvider{response: "Plan interview response"} + + al := NewAgentLoop(cfg, msgBus, provider) + + defaultAgent := al.registry.GetDefaultAgent() + + if defaultAgent == nil { + t.Fatal("No default agent found") + } + + memoryDir := filepath.Join(tmpDir, "memory") + + os.MkdirAll(memoryDir, 0o755) + + memoryPath := filepath.Join(memoryDir, "MEMORY.md") + + memoryContent := "# Active Plan\n\n> Task: Test plan model\n> Status: interviewing\n> Phase: 1\n" + + if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { + t.Fatalf("Failed to write MEMORY.md: %v", wErr) + } + + _, err = al.ProcessDirectWithChannel( + + context.Background(), + + "Hello, plan model test", + + "test-plan-session", + + "test", + + "test-chat", + ) + if err != nil { + t.Fatalf("ProcessDirectWithChannel failed: %v", err) + } + + provider.mu.Lock() + + defer provider.mu.Unlock() + + if len(provider.models) == 0 { + t.Fatal("Expected at least one Chat call") + } + + if provider.models[0] != "plan-model" { + t.Errorf("Expected plan model 'plan-model' during interviewing, got %q", provider.models[0]) + } +} + +func TestAgentLoop_PlanModel_NotUsedDuringExecuting(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-planmodel-exec-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "normal-model", + + PlanModel: "plan-model", + + MaxTokens: 4096, + + MaxToolIterations: 2, + }, + }, + } + + msgBus := bus.NewMessageBus() + + provider := &modelCapturingMockProvider{response: "Executing response"} + + al := NewAgentLoop(cfg, msgBus, provider) + + defaultAgent := al.registry.GetDefaultAgent() + + if defaultAgent == nil { + t.Fatal("No default agent found") + } + + memoryDir := filepath.Join(tmpDir, "memory") + + os.MkdirAll(memoryDir, 0o755) + + memoryPath := filepath.Join(memoryDir, "MEMORY.md") + + memoryContent := `# Active Plan + + + +> Task: Test plan model + +> Status: executing + +> Phase: 1 + + + +## Phase 1: Build + +- [ ] Run build + +` + + if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { + t.Fatalf("Failed to write MEMORY.md: %v", wErr) + } + + _, err = al.ProcessDirectWithChannel( + + context.Background(), + + "Hello, executing test", + + "test-exec-session", + + "test", + + "test-chat", + ) + if err != nil { + t.Fatalf("ProcessDirectWithChannel failed: %v", err) + } + + provider.mu.Lock() + + defer provider.mu.Unlock() + + if len(provider.models) == 0 { + t.Fatal("Expected at least one Chat call") + } + + if provider.models[0] != "normal-model" { + t.Errorf("Expected normal model 'normal-model' during executing, got %q", provider.models[0]) + } +} + +func TestAgentLoop_PlanModel_ResolvesProviderForSingleCandidate(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-planmodel-resolve-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + + Model: "MiniMax-M2.5", + + PlanModel: "openai/gpt-5.2", + + MaxTokens: 4096, + + MaxToolIterations: 2, + }, + }, + } + + msgBus := bus.NewMessageBus() + + mainProvider := &modelCapturingMockProvider{response: "wrong provider response"} + + al := NewAgentLoop(cfg, msgBus, mainProvider) + + resolvedProvider := &modelCapturingMockProvider{response: "correct provider response"} + + al.providerCache["openai/gpt-5.2"] = resolvedProvider + + memoryDir := filepath.Join(tmpDir, "memory") + + os.MkdirAll(memoryDir, 0o755) + + memoryPath := filepath.Join(memoryDir, "MEMORY.md") + + memoryContent := "# Active Plan\n\n> Task: Test provider resolution\n> Status: interviewing\n> Phase: 1\n" + + if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { + t.Fatalf("Failed to write MEMORY.md: %v", wErr) + } + + _, err = al.ProcessDirectWithChannel( + + context.Background(), + + "Hello, resolve provider test", + + "test-resolve-session", + + "test", + + "test-chat", + ) + if err != nil { + t.Fatalf("ProcessDirectWithChannel failed: %v", err) + } + + resolvedProvider.mu.Lock() + + defer resolvedProvider.mu.Unlock() + + mainProvider.mu.Lock() + + defer mainProvider.mu.Unlock() + + if len(resolvedProvider.models) == 0 { + t.Fatal("Expected resolved provider to receive Chat call, but it got none") + } + + if resolvedProvider.models[0] != "gpt-5.2" { + t.Errorf("Expected resolved provider to receive model 'gpt-5.2', got %q", resolvedProvider.models[0]) + } + + if len(mainProvider.models) > 0 { + t.Errorf("Expected main provider to receive no Chat calls during plan model phase, got %d calls with models %v", + + len(mainProvider.models), mainProvider.models) + } +} + +func TestPlanCommand_StartClear(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + + _ = agent.ContextBuilder.WriteMemory(plan) + + agent.Sessions.AddMessage("test-session", "user", "hello") + + agent.Sessions.AddMessage("test-session", "assistant", "world") + + agent.Sessions.SetSummary("test-session", "some summary") + + response, handled := al.handleCommand(context.Background(), bus.InboundMessage{ + Content: "/plan start clear", + + SessionKey: "test-session", + }) + + if !handled { + t.Fatal("expected /plan start clear to be handled") + } + + if !strings.Contains(response, "clean history") { + t.Errorf("expected 'clean history' in response, got %q", response) + } + + if !al.planStartPending { + t.Error("expected planStartPending to be true") + } + + if !al.planClearHistory { + t.Error("expected planClearHistory to be true") + } + + al.planStartPending = false + + clearHistory := al.planClearHistory + + al.planClearHistory = false + + if clearHistory { + agent.Sessions.SetHistory("test-session", nil) + + agent.Sessions.SetSummary("test-session", "") + + _ = agent.Sessions.Save("test-session") + } + + history := agent.Sessions.GetHistory("test-session") + + if len(history) != 0 { + t.Errorf("expected empty history after clear, got %d messages", len(history)) + } + + summary := agent.Sessions.GetSummary("test-session") + + if summary != "" { + t.Errorf("expected empty summary after clear, got %q", summary) + } +} + +func TestPlanCommand_StartWithoutClear_PreservesHistory(t *testing.T) { + al, cleanup := newTestAgentLoop(t) + + defer cleanup() + + agent := al.registry.GetDefaultAgent() + + plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" + + _ = agent.ContextBuilder.WriteMemory(plan) + + agent.Sessions.AddMessage("test-session", "user", "hello") + + agent.Sessions.AddMessage("test-session", "assistant", "world") + + agent.Sessions.SetSummary("test-session", "some summary") + + response, _ := al.handleCommand(context.Background(), bus.InboundMessage{ + Content: "/plan start", + + SessionKey: "test-session", + }) + + if strings.Contains(response, "clean history") { + t.Errorf("did not expect 'clean history' in response, got %q", response) + } + + if al.planClearHistory { + t.Error("planClearHistory should be false for /plan start without clear") + } + + history := agent.Sessions.GetHistory("test-session") + + if len(history) != 2 { + t.Errorf("expected 2 history messages preserved, got %d", len(history)) + } + + summary := agent.Sessions.GetSummary("test-session") + + if summary != "some summary" { + t.Errorf("expected summary preserved, got %q", summary) + } +} + +func TestFilterInterviewTools(t *testing.T) { + allDefs := []providers.ToolDefinition{ + {Function: protocoltypes.ToolFunctionDefinition{Name: "read_file"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "list_dir"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "web_search"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "web_fetch"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "message"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "edit_file"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "append_file"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "write_file"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "exec"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "logs"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "spawn_subagent"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "skills_search"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "skills_install"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "bg_monitor"}}, + + {Function: protocoltypes.ToolFunctionDefinition{Name: "i2c_transfer"}}, + } + + filtered := filterInterviewTools(allDefs) + + if len(filtered) != 10 { + names := make([]string, len(filtered)) + + for i, d := range filtered { + names[i] = d.Function.Name + } + + t.Errorf("expected 10 allowed tools, got %d: %v", len(filtered), names) + } + + disallowed := map[string]bool{ + "spawnsubagent": true, "skillssearch": true, + + "skillsinstall": true, "bgmonitor": true, "ictransfer": true, + } + + for _, d := range filtered { + norm := tools.NormalizeToolName(d.Function.Name) + + if disallowed[norm] { + t.Errorf("disallowed tool %q should have been filtered out", d.Function.Name) + } + } +} + +func TestBuildStreamingDisplay_ContentOnly(t *testing.T) { + display := buildStreamingDisplay("hello world", "") + + if !strings.HasSuffix(display, " \u2589") { + t.Error("expected cursor suffix") + } + + if strings.Contains(display, "\U0001f9e0") { + t.Error("should not contain brain emoji when no reasoning") + } + + lines := strings.Count(display, "\n") + 1 + + if lines != streamingDisplayLines+1 { + t.Logf("display:\n%s", display) + } +} + +func TestBuildStreamingDisplay_ReasoningOnly(t *testing.T) { + display := buildStreamingDisplay("", "let me think about this") + + if !strings.Contains(display, "\U0001f9e0") { + t.Error("expected brain emoji for reasoning phase") + } + + if !strings.Contains(display, "Thinking...") { + t.Error("expected Thinking... header") + } + + if !strings.HasSuffix(display, " \u2589") { + t.Error("expected cursor suffix") + } +} + +func TestBuildStreamingDisplay_Both(t *testing.T) { + display := buildStreamingDisplay("the answer is 42", "first I considered...") + + if !strings.Contains(display, "\U0001f9e0") { + t.Error("expected brain emoji") + } + + if !strings.Contains(display, "responding") { + t.Error("expected responding header when both present") + } + + if !strings.Contains(display, "the answer is 42") { + t.Error("expected content in display") + } +} + +func TestFormatDurationMs(t *testing.T) { + tests := []struct { + ms int64 + + want string + }{ + {0, "0ms"}, + + {500, "500ms"}, + + {999, "999ms"}, + + {1000, "1.0s"}, + + {1200, "1.2s"}, + + {3500, "3.5s"}, + + {59900, "59.9s"}, + + {60000, "1m"}, + + {61000, "1m1s"}, + + {65000, "1m5s"}, + + {120000, "2m"}, + + {3661000, "61m1s"}, + } + + for _, tt := range tests { + t.Run(fmt.Sprintf("%dms", tt.ms), func(t *testing.T) { + got := formatDurationMs(tt.ms) + + if got != tt.want { + t.Errorf("formatDurationMs(%d) = %q, want %q", tt.ms, got, tt.want) + } + }) + } +} + +func TestFormatSubagentCompletion(t *testing.T) { + tests := []struct { + name string + + label string + + metadata map[string]string + + want string + }{ + { + "no metadata", + + "scout-1", + + nil, + + "📋 scout-1 completed.", + }, + + { + "empty metadata", + + "scout-1", + + map[string]string{}, + + "📋 scout-1 completed.", + }, + + { + "duration and tool calls", + + "scout-1", + + map[string]string{"duration_ms": "3200", "tool_calls": "5"}, + + "📋 scout-1 completed (3.2s, 5 tool calls).", + }, + + { + "single tool call", + + "coder-1", + + map[string]string{"duration_ms": "1200", "tool_calls": "1"}, + + "📋 coder-1 completed (1.2s, 1 tool call).", + }, + + { + "duration only", + + "scout-2", + + map[string]string{"duration_ms": "65000", "tool_calls": "0"}, + + "📋 scout-2 completed (1m5s).", + }, + + { + "tool calls only", + + "scout-3", + + map[string]string{"duration_ms": "0", "tool_calls": "10"}, + + "📋 scout-3 completed (10 tool calls).", + }, + + { + "zero everything", + + "scout-4", + + map[string]string{"duration_ms": "0", "tool_calls": "0"}, + + "📋 scout-4 completed.", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := formatSubagentCompletion(tt.label, tt.metadata) + + if got != tt.want { + t.Errorf("formatSubagentCompletion(%q, %v) = %q, want %q", tt.label, tt.metadata, got, tt.want) + } + }) + } +} diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index ee7f45e81..cab82e176 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -7,490 +7,290 @@ import ( "path/filepath" "slices" "strings" - "sync" "testing" "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/providers" - "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" + "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/tools" ) type fakeChannel struct{ id string } -func (f *fakeChannel) Name() string { return "fake" } - -func (f *fakeChannel) Start(ctx context.Context) error { return nil } - -func (f *fakeChannel) Stop(ctx context.Context) error { return nil } - +func (f *fakeChannel) Name() string { return "fake" } +func (f *fakeChannel) Start(ctx context.Context) error { return nil } +func (f *fakeChannel) Stop(ctx context.Context) error { return nil } func (f *fakeChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { return nil } +func (f *fakeChannel) IsRunning() bool { return true } +func (f *fakeChannel) IsAllowed(string) bool { return true } +func (f *fakeChannel) IsAllowedSender(sender bus.SenderInfo) bool { return true } +func (f *fakeChannel) ReasoningChannelID() string { return f.id } -func (f *fakeChannel) IsRunning() bool { return true } - -func (f *fakeChannel) IsAllowed(string) bool { return true } - -func (f *fakeChannel) IsAllowedSender(sender bus.SenderInfo) bool { return true } - -func (f *fakeChannel) ReasoningChannelID() string { return f.id } - -func TestRecordLastChannel(t *testing.T) { - // Create temp workspace - +func newTestAgentLoop( + t *testing.T, +) (al *AgentLoop, cfg *config.Config, msgBus *bus.MessageBus, provider *mockProvider, cleanup func()) { + t.Helper() tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - - defer os.RemoveAll(tmpDir) - - // Create test config - - cfg := &config.Config{ + cfg = &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } + msgBus = bus.NewMessageBus() + provider = &mockProvider{} + al = NewAgentLoop(cfg, msgBus, provider) + return al, cfg, msgBus, provider, func() { os.RemoveAll(tmpDir) } +} - // Create agent loop - - msgBus := bus.NewMessageBus() - - provider := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, provider) - - // Test RecordLastChannel +func TestRecordLastChannel(t *testing.T) { + al, cfg, msgBus, provider, cleanup := newTestAgentLoop(t) + defer cleanup() testChannel := "test-channel" - - err = al.RecordLastChannel(testChannel) - if err != nil { + if err := al.RecordLastChannel(testChannel); err != nil { t.Fatalf("RecordLastChannel failed: %v", err) } - - // Verify channel was saved - - lastChannel := al.state.GetLastChannel() - - if lastChannel != testChannel { - t.Errorf("Expected channel '%s', got '%s'", testChannel, lastChannel) + if got := al.state.GetLastChannel(); got != testChannel { + t.Errorf("Expected channel '%s', got '%s'", testChannel, got) } - - // Verify persistence by creating a new agent loop - al2 := NewAgentLoop(cfg, msgBus, provider) - - if al2.state.GetLastChannel() != testChannel { - t.Errorf("Expected persistent channel '%s', got '%s'", testChannel, al2.state.GetLastChannel()) + if got := al2.state.GetLastChannel(); got != testChannel { + t.Errorf("Expected persistent channel '%s', got '%s'", testChannel, got) } } func TestRecordLastChatID(t *testing.T) { - // Create temp workspace - - tmpDir, err := os.MkdirTemp("", "agent-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - // Create test config - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - // Create agent loop - - msgBus := bus.NewMessageBus() - - provider := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, provider) - - // Test RecordLastChatID + al, cfg, msgBus, provider, cleanup := newTestAgentLoop(t) + defer cleanup() testChatID := "test-chat-id-123" - - err = al.RecordLastChatID(testChatID) - if err != nil { + if err := al.RecordLastChatID(testChatID); err != nil { t.Fatalf("RecordLastChatID failed: %v", err) } - - // Verify chat ID was saved - - lastChatID := al.state.GetLastChatID() - - if lastChatID != testChatID { - t.Errorf("Expected chat ID '%s', got '%s'", testChatID, lastChatID) + if got := al.state.GetLastChatID(); got != testChatID { + t.Errorf("Expected chat ID '%s', got '%s'", testChatID, got) } - - // Verify persistence by creating a new agent loop - al2 := NewAgentLoop(cfg, msgBus, provider) - - if al2.state.GetLastChatID() != testChatID { - t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, al2.state.GetLastChatID()) - } -} - -func TestRecordLastHeartbeatTarget(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - msgBus := bus.NewMessageBus() - - provider := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, provider) - - target := "telegram:-100123/42" - - if err := al.RecordLastHeartbeatTarget(target); err != nil { - t.Fatalf("RecordLastHeartbeatTarget failed: %v", err) - } - - if got := al.state.GetLastHeartbeatTarget(); got != target { - t.Fatalf("GetLastHeartbeatTarget = %q, want %q", got, target) + if got := al2.state.GetLastChatID(); got != testChatID { + t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, got) } } func TestNewAgentLoop_StateInitialized(t *testing.T) { // Create temp workspace - tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) // Create test config - cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } // Create agent loop - msgBus := bus.NewMessageBus() - provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) // Verify state manager is initialized - if al.state == nil { t.Error("Expected state manager to be initialized") } // Verify state directory was created - stateDir := filepath.Join(tmpDir, "state") - if _, err := os.Stat(stateDir); os.IsNotExist(err) { t.Error("Expected state directory to exist") } } // TestToolRegistry_ToolRegistration verifies tools can be registered and retrieved - func TestToolRegistry_ToolRegistration(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() - provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) // Register a custom tool - customTool := &mockCustomTool{} - al.RegisterTool(customTool) // Verify tool is registered by checking it doesn't panic on GetStartupInfo - // (actual tool retrieval is tested in tools package tests) - info := al.GetStartupInfo() - toolsInfo := info["tools"].(map[string]any) - toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list - found := slices.Contains(toolsList, "mock_custom") - if !found { t.Error("Expected custom tool to be registered") } } -// TestToolContext_Updates verifies tool context is updated with channel/chatID - +// TestToolContext_Updates verifies tool context helpers work correctly func TestToolContext_Updates(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) + ctx := tools.WithToolContext(context.Background(), "telegram", "chat-42") + + if got := tools.ToolChannel(ctx); got != "telegram" { + t.Errorf("expected channel 'telegram', got %q", got) + } + if got := tools.ToolChatID(ctx); got != "chat-42" { + t.Errorf("expected chatID 'chat-42', got %q", got) } - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, + // Empty context returns empty strings + if got := tools.ToolChannel(context.Background()); got != "" { + t.Errorf("expected empty channel from bare context, got %q", got) } - - msgBus := bus.NewMessageBus() - - provider := &simpleMockProvider{response: "OK"} - - _ = NewAgentLoop(cfg, msgBus, provider) - - // Verify that ContextualTool interface is defined and can be implemented - - // This test validates the interface contract exists - - ctxTool := &mockContextualTool{} - - // Verify the tool implements the interface correctly - - var _ tools.ContextualTool = ctxTool } // TestToolRegistry_GetDefinitions verifies tool definitions can be retrieved - func TestToolRegistry_GetDefinitions(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() - provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) // Register a test tool and verify it shows up in startup info - testTool := &mockCustomTool{} - al.RegisterTool(testTool) info := al.GetStartupInfo() - toolsInfo := info["tools"].(map[string]any) - toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list - found := slices.Contains(toolsList, "mock_custom") - if !found { t.Error("Expected custom tool to be registered") } } // TestAgentLoop_GetStartupInfo verifies startup info contains tools - func TestAgentLoop_GetStartupInfo(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = tmpDir + cfg.Agents.Defaults.Model = "test-model" + cfg.Agents.Defaults.MaxTokens = 4096 + cfg.Agents.Defaults.MaxToolIterations = 10 msgBus := bus.NewMessageBus() - provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) info := al.GetStartupInfo() // Verify tools info exists - toolsInfo, ok := info["tools"] - if !ok { t.Fatal("Expected 'tools' key in startup info") } toolsMap, ok := toolsInfo.(map[string]any) - if !ok { t.Fatal("Expected 'tools' to be a map") } count, ok := toolsMap["count"] - if !ok { t.Fatal("Expected 'count' in tools info") } // Should have default tools registered - if count.(int) == 0 { t.Error("Expected at least some tools to be registered") } } // TestAgentLoop_Stop verifies Stop() sets running to false - func TestAgentLoop_Stop(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() - provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) // Note: running is only set to true when Run() is called - // We can't test that without starting the event loop - // Instead, verify the Stop method can be called safely - al.Stop() // Verify running is false (initial state or after Stop) - if al.running.Load() { t.Error("Expected agent to be stopped (or never started)") } @@ -504,18 +304,13 @@ type simpleMockProvider struct { func (m *simpleMockProvider) Chat( ctx context.Context, - messages []providers.Message, - tools []providers.ToolDefinition, - model string, - opts map[string]any, ) (*providers.LLMResponse, error) { return &providers.LLMResponse{ - Content: m.response, - + Content: m.response, ToolCalls: []providers.ToolCall{}, }, nil } @@ -524,8 +319,30 @@ func (m *simpleMockProvider) GetDefaultModel() string { return "mock-model" } -// mockCustomTool is a simple mock tool for registration testing +type countingMockProvider struct { + response string + calls int +} +func (m *countingMockProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, +) (*providers.LLMResponse, error) { + m.calls++ + return &providers.LLMResponse{ + Content: m.response, + ToolCalls: []providers.ToolCall{}, + }, nil +} + +func (m *countingMockProvider) GetDefaultModel() string { + return "counting-mock-model" +} + +// mockCustomTool is a simple mock tool for registration testing type mockCustomTool struct{} func (m *mockCustomTool) Name() string { @@ -538,8 +355,7 @@ func (m *mockCustomTool) Description() string { func (m *mockCustomTool) Parameters() map[string]any { return map[string]any{ - "type": "object", - + "type": "object", "properties": map[string]any{}, } } @@ -548,209 +364,322 @@ func (m *mockCustomTool) Execute(ctx context.Context, args map[string]any) *tool return tools.SilentResult("Custom tool executed") } -// mockContextualTool tracks context updates - -type mockContextualTool struct { - lastChannel string - - lastChatID string -} - -func (m *mockContextualTool) Name() string { - return "mock_contextual" -} - -func (m *mockContextualTool) Description() string { - return "Mock contextual tool" -} - -func (m *mockContextualTool) Parameters() map[string]any { - return map[string]any{ - "type": "object", - - "properties": map[string]any{}, - } -} - -func (m *mockContextualTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult { - return tools.SilentResult("Contextual tool executed") -} - -func (m *mockContextualTool) SetContext(channel, chatID string) { - m.lastChannel = channel - - m.lastChatID = chatID -} - // testHelper executes a message and returns the response - type testHelper struct { al *AgentLoop } func (h testHelper) executeAndGetResponse(tb testing.TB, ctx context.Context, msg bus.InboundMessage) string { // Use a short timeout to avoid hanging - timeoutCtx, cancel := context.WithTimeout(ctx, responseTimeout) - defer cancel() response, err := h.al.processMessage(timeoutCtx, msg) if err != nil { tb.Fatalf("processMessage failed: %v", err) } - return response } const responseTimeout = 3 * time.Second -// TestToolResult_SilentToolDoesNotSendUserMessage verifies silent tools don't trigger outbound - -func TestToolResult_SilentToolDoesNotSendUserMessage(t *testing.T) { +func TestProcessMessage_UsesRouteSessionKey(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() - - provider := &simpleMockProvider{response: "File operation complete"} - + provider := &simpleMockProvider{response: "ok"} al := NewAgentLoop(cfg, msgBus, provider) + msg := bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "hello", + Peer: bus.Peer{ + Kind: "direct", + ID: "user1", + }, + } + + route := al.registry.ResolveRoute(routing.RouteInput{ + Channel: msg.Channel, + Peer: extractPeer(msg), + }) + sessionKey := route.SessionKey + + defaultAgent := al.registry.GetDefaultAgent() + if defaultAgent == nil { + t.Fatal("No default agent found") + } + + helper := testHelper{al: al} + _ = helper.executeAndGetResponse(t, context.Background(), msg) + + history := defaultAgent.Sessions.GetHistory(sessionKey) + if len(history) != 2 { + t.Fatalf("expected session history len=2, got %d", len(history)) + } + if history[0].Role != "user" || history[0].Content != "hello" { + t.Fatalf("unexpected first message in session: %+v", history[0]) + } +} + +func TestProcessMessage_CommandOutcomes(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + Session: config.SessionConfig{ + DMScope: "per-channel-peer", + }, + } + + msgBus := bus.NewMessageBus() + provider := &countingMockProvider{response: "LLM reply"} + al := NewAgentLoop(cfg, msgBus, provider) + helper := testHelper{al: al} + + baseMsg := bus.InboundMessage{ + Channel: "whatsapp", + SenderID: "user1", + ChatID: "chat1", + Peer: bus.Peer{ + Kind: "direct", + ID: "user1", + }, + } + + showResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: baseMsg.Channel, + SenderID: baseMsg.SenderID, + ChatID: baseMsg.ChatID, + Content: "/show channel", + Peer: baseMsg.Peer, + }) + if showResp != "Current Channel: whatsapp" { + t.Fatalf("unexpected /show reply: %q", showResp) + } + if provider.calls != 0 { + t.Fatalf("LLM should not be called for handled command, calls=%d", provider.calls) + } + + fooResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: baseMsg.Channel, + SenderID: baseMsg.SenderID, + ChatID: baseMsg.ChatID, + Content: "/foo", + Peer: baseMsg.Peer, + }) + if fooResp != "LLM reply" { + t.Fatalf("unexpected /foo reply: %q", fooResp) + } + if provider.calls != 1 { + t.Fatalf("LLM should be called exactly once after /foo passthrough, calls=%d", provider.calls) + } + + newResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: baseMsg.Channel, + SenderID: baseMsg.SenderID, + ChatID: baseMsg.ChatID, + Content: "/new", + Peer: baseMsg.Peer, + }) + if newResp != "LLM reply" { + t.Fatalf("unexpected /new reply: %q", newResp) + } + if provider.calls != 2 { + t.Fatalf("LLM should be called for passthrough /new command, calls=%d", provider.calls) + } +} + +func TestProcessMessage_SwitchModelShowModelConsistency(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Provider: "openai", + Model: "before-switch", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &countingMockProvider{response: "LLM reply"} + al := NewAgentLoop(cfg, msgBus, provider) + helper := testHelper{al: al} + + switchResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "/switch model to after-switch", + Peer: bus.Peer{ + Kind: "direct", + ID: "user1", + }, + }) + if !strings.Contains(switchResp, "Switched model from before-switch to after-switch") { + t.Fatalf("unexpected /switch reply: %q", switchResp) + } + + showResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ + Channel: "telegram", + SenderID: "user1", + ChatID: "chat1", + Content: "/show model", + Peer: bus.Peer{ + Kind: "direct", + ID: "user1", + }, + }) + if !strings.Contains(showResp, "Current Model: after-switch (Provider: openai)") { + t.Fatalf("unexpected /show model reply after switch: %q", showResp) + } + + if provider.calls != 0 { + t.Fatalf("LLM should not be called for /switch and /show, calls=%d", provider.calls) + } +} + +// TestToolResult_SilentToolDoesNotSendUserMessage verifies silent tools don't trigger outbound +func TestToolResult_SilentToolDoesNotSendUserMessage(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &simpleMockProvider{response: "File operation complete"} + al := NewAgentLoop(cfg, msgBus, provider) helper := testHelper{al: al} // ReadFileTool returns SilentResult, which should not send user message - ctx := context.Background() - msg := bus.InboundMessage{ - Channel: "test", - - SenderID: "user1", - - ChatID: "chat1", - - Content: "read test.txt", - + Channel: "test", + SenderID: "user1", + ChatID: "chat1", + Content: "read test.txt", SessionKey: "test-session", } response := helper.executeAndGetResponse(t, ctx, msg) // Silent tool should return the LLM's response directly - if response != "File operation complete" { t.Errorf("Expected 'File operation complete', got: %s", response) } } // TestToolResult_UserFacingToolDoesSendMessage verifies user-facing tools trigger outbound - func TestToolResult_UserFacingToolDoesSendMessage(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } msgBus := bus.NewMessageBus() - provider := &simpleMockProvider{response: "Command output: hello world"} - al := NewAgentLoop(cfg, msgBus, provider) - helper := testHelper{al: al} // ExecTool returns UserResult, which should send user message - ctx := context.Background() - msg := bus.InboundMessage{ - Channel: "test", - - SenderID: "user1", - - ChatID: "chat1", - - Content: "run hello", - + Channel: "test", + SenderID: "user1", + ChatID: "chat1", + Content: "run hello", SessionKey: "test-session", } response := helper.executeAndGetResponse(t, ctx, msg) // User-facing tool should include the output in final response - if response != "Command output: hello world" { t.Errorf("Expected 'Command output: hello world', got: %s", response) } } // failFirstMockProvider fails on the first N calls with a specific error - type failFirstMockProvider struct { - failures int - + failures int currentCall int - - failError error - + failError error successResp string } func (m *failFirstMockProvider) Chat( ctx context.Context, - messages []providers.Message, - tools []providers.ToolDefinition, - model string, - opts map[string]any, ) (*providers.LLMResponse, error) { m.currentCall++ - if m.currentCall <= m.failures { return nil, m.failError } - return &providers.LLMResponse{ - Content: m.successResp, - + Content: m.successResp, ToolCalls: []providers.ToolCall{}, }, nil } @@ -760,24 +689,19 @@ func (m *failFirstMockProvider) GetDefaultModel() string { } // TestAgentLoop_ContextExhaustionRetry verify that the agent retries on context errors - func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, @@ -786,61 +710,39 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { msgBus := bus.NewMessageBus() // Create a provider that fails once with a context error - contextErr := fmt.Errorf("InvalidParameter: Total tokens of image and text exceed max message tokens") - provider := &failFirstMockProvider{ - failures: 1, - - failError: contextErr, - + failures: 1, + failError: contextErr, successResp: "Recovered from context error", } al := NewAgentLoop(cfg, msgBus, provider) // Inject some history to simulate a full context - sessionKey := "test-session-context" - // Create dummy history - history := []providers.Message{ {Role: "system", Content: "System prompt"}, - {Role: "user", Content: "Old message 1"}, - {Role: "assistant", Content: "Old response 1"}, - {Role: "user", Content: "Old message 2"}, - {Role: "assistant", Content: "Old response 2"}, - {Role: "user", Content: "Trigger message"}, } - defaultAgent := al.registry.GetDefaultAgent() - if defaultAgent == nil { t.Fatal("No default agent found") } - defaultAgent.Sessions.SetHistory(sessionKey, history) // Call ProcessDirectWithChannel - // Note: ProcessDirectWithChannel calls processMessage which will execute runLLMIteration - response, err := al.ProcessDirectWithChannel( - context.Background(), - "Trigger message", - sessionKey, - "test", - "test-chat", ) if err != nil { @@ -852,63 +754,69 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { } // We expect 2 calls: 1st failed, 2nd succeeded - if provider.currentCall != 2 { t.Errorf("Expected 2 calls (1 fail + 1 success), got %d", provider.currentCall) } // Check final history length - finalHistory := defaultAgent.Sessions.GetHistory(sessionKey) - // We verify that the history has been modified (compressed) - // Original length: 6 - // Expected behavior: compression drops ~50% of history (mid slice) - // We can assert that the length is NOT what it would be without compression. - // Without compression: 6 + 1 (new user msg) + 1 (assistant msg) = 8 - if len(finalHistory) >= 8 { t.Errorf("Expected history to be compressed (len < 8), got %d", len(finalHistory)) } } -func TestShouldInjectReminder(t *testing.T) { - tests := []struct { - name string +func TestProcessDirectWithChannel_InitializesMCPInAgentMode(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) - iteration int - - interval int - - want bool - }{ - {"first iteration skipped", 1, 5, false}, - - {"iteration 5 interval 5", 5, 5, true}, - - {"iteration 10 interval 5", 10, 5, true}, - - {"iteration 3 interval 5", 3, 5, false}, - - {"interval zero disabled", 5, 0, false}, - - {"interval negative disabled", 5, -1, false}, - - {"iteration 2 interval 1", 2, 1, true}, + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + Tools: config.ToolsConfig{ + MCP: config.MCPConfig{ + ToolConfig: config.ToolConfig{ + Enabled: true, + }, + }, + }, } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := shouldInjectReminder(tt.iteration, tt.interval) + msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + defer al.Close() - if got != tt.want { - t.Errorf("shouldInjectReminder(%d, %d) = %v, want %v", tt.iteration, tt.interval, got, tt.want) - } - }) + if al.mcp.hasManager() { + t.Fatal("expected MCP manager to be nil before first direct processing") + } + + _, err = al.ProcessDirectWithChannel( + context.Background(), + "hello", + "session-1", + "cli", + "direct", + ) + if err != nil { + t.Fatalf("ProcessDirectWithChannel failed: %v", err) + } + + if !al.mcp.hasManager() { + t.Fatal("expected MCP manager to be initialized in direct agent mode") } } @@ -917,96 +825,63 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) { if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - defer os.RemoveAll(tmpDir) cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) - chManager, err := channels.NewManager(&config.Config{}, bus.NewMessageBus(), nil) if err != nil { t.Fatalf("Failed to create channel manager: %v", err) } - for name, id := range map[string]string{ - "whatsapp": "rid-whatsapp", - - "telegram": "rid-telegram", - - "feishu": "rid-feishu", - - "discord": "rid-discord", - - "maixcam": "rid-maixcam", - - "qq": "rid-qq", - - "dingtalk": "rid-dingtalk", - - "slack": "rid-slack", - - "line": "rid-line", - - "onebot": "rid-onebot", - - "wecom": "rid-wecom", - + "whatsapp": "rid-whatsapp", + "telegram": "rid-telegram", + "feishu": "rid-feishu", + "discord": "rid-discord", + "maixcam": "rid-maixcam", + "qq": "rid-qq", + "dingtalk": "rid-dingtalk", + "slack": "rid-slack", + "line": "rid-line", + "onebot": "rid-onebot", + "wecom": "rid-wecom", "wecom_app": "rid-wecom-app", } { chManager.RegisterChannel(name, &fakeChannel{id: id}) } - al.SetChannelManager(chManager) - tests := []struct { channel string - - wantID string + wantID string }{ {channel: "whatsapp", wantID: "rid-whatsapp"}, - {channel: "telegram", wantID: "rid-telegram"}, - {channel: "feishu", wantID: "rid-feishu"}, - {channel: "discord", wantID: "rid-discord"}, - {channel: "maixcam", wantID: "rid-maixcam"}, - {channel: "qq", wantID: "rid-qq"}, - {channel: "dingtalk", wantID: "rid-dingtalk"}, - {channel: "slack", wantID: "rid-slack"}, - {channel: "line", wantID: "rid-line"}, - {channel: "onebot", wantID: "rid-onebot"}, - {channel: "wecom", wantID: "rid-wecom"}, - {channel: "wecom_app", wantID: "rid-wecom-app"}, - {channel: "unknown", wantID: ""}, } for _, tt := range tests { t.Run(tt.channel, func(t *testing.T) { got := al.targetReasoningChannelID(tt.channel) - if got != tt.wantID { t.Fatalf("targetReasoningChannelID(%q) = %q, want %q", tt.channel, got, tt.wantID) } @@ -1014,3097 +889,34 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) { } } -func TestBuildTaskReminder_WithoutBlocker(t *testing.T) { - msg := buildTaskReminder("implement feature X", "") - - if msg.Role != "user" { - t.Errorf("expected role 'user', got %q", msg.Role) - } - - if !strings.Contains(msg.Content, "[TASK REMINDER]") { - t.Error("expected content to contain '[TASK REMINDER]'") - } - - if !strings.Contains(msg.Content, "implement feature X") { - t.Error("expected content to contain original message") - } - - if strings.Contains(msg.Content, "blocker") { - t.Error("expected content NOT to contain 'blocker' when no blocker provided") - } - - if !strings.Contains(msg.Content, "move on") { - t.Error("expected content to contain completion prompt") - } -} - -func TestBuildTaskReminder_WithBlocker(t *testing.T) { - msg := buildTaskReminder("implement feature X", "ModuleNotFoundError: No module named 'foo'") - - if msg.Role != "user" { - t.Errorf("expected role 'user', got %q", msg.Role) - } - - if !strings.Contains(msg.Content, "[TASK REMINDER]") { - t.Error("expected content to contain '[TASK REMINDER]'") - } - - if !strings.Contains(msg.Content, "implement feature X") { - t.Error("expected content to contain original message") - } - - if !strings.Contains(msg.Content, "Last blocker") { - t.Error("expected content to contain 'Last blocker'") - } - - if !strings.Contains(msg.Content, "ModuleNotFoundError") { - t.Error("expected content to contain blocker text") - } -} - -func TestResolveProvider_CachesProviders(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - Provider: "vllm", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - - Providers: config.ProvidersConfig{ - VLLM: config.ProviderConfig{ - APIKey: "test-key", - - APIBase: "https://example.com/v1", - }, - }, - } - - msgBus := bus.NewMessageBus() - - primary := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, primary) - - // First call creates and caches a provider for "vllm/test-model" - - p1 := al.resolveProvider("vllm", "test-model", primary) - - if p1 == primary { - t.Fatal("expected a new provider from legacy providers config, not the fallback") - } - - // Calling again should return the same instance (cached) - - p2 := al.resolveProvider("vllm", "test-model", primary) - - if p1 != p2 { - t.Fatal("expected same cached instance on second call") - } -} - -func TestResolveProvider_FallsBackOnError(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - Provider: "vllm", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - msgBus := bus.NewMessageBus() - - primary := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, primary) - - // Request a provider that can't be created (no config for "nonexistent") - - p := al.resolveProvider("nonexistent", "unknown-model", primary) - - if p != primary { - t.Fatal("expected fallback to primary provider on creation error") - } - - // Ensure the failed provider is NOT cached - - if _, ok := al.providerCache["nonexistent"]; ok { - t.Fatal("failed provider should not be cached") - } -} - -func TestResolveProvider_EmptyNameReturnsFallback(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - msgBus := bus.NewMessageBus() - - primary := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, primary) - - p := al.resolveProvider("", "", primary) - - if p != primary { - t.Fatal("expected fallback provider for empty name") - } -} - -// TestSlashCommandResponseSkipsPlaceholder verifies that slash command responses - -// are published with SkipPlaceholder=true so they don't overwrite the ongoing task - -// status bubble. - -func TestSlashCommandResponseSkipsPlaceholder(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - msgBus := bus.NewMessageBus() - - provider := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, provider) - - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - - defer cancel() - - go func() { - _ = al.Run(ctx) - }() - - // Send a slash command - - msgBus.PublishInbound(context.Background(), bus.InboundMessage{ - Channel: "telegram", - - SenderID: "user1", - - ChatID: "chat1", - - Content: "/skills", - }) - - // Read the outbound message - - outMsg, ok := msgBus.SubscribeOutbound(ctx) - - if !ok { - t.Fatal("expected outbound message from slash command") - } - - if !outMsg.SkipPlaceholder { - t.Errorf("expected SkipPlaceholder=true for slash command response, got false") - } -} - -func TestBuildTaskReminder_Truncation(t *testing.T) { - // Build a long message (1000 runes) - - longMsg := strings.Repeat("あ", 1000) - - longBlocker := strings.Repeat("X", 500) - - msg := buildTaskReminder(longMsg, longBlocker) - - // The full message should NOT contain 1000 'あ' characters - - runeCount := strings.Count(msg.Content, "あ") - - if runeCount >= 1000 { - t.Errorf("expected task message to be truncated, got %d 'あ' runes", runeCount) - } - - // Should be at most taskReminderMaxChars (500) runes for the task part - - if runeCount > taskReminderMaxChars { - t.Errorf("expected at most %d task runes, got %d", taskReminderMaxChars, runeCount) - } - - // Blocker should be truncated too - - xCount := strings.Count(msg.Content, "X") - - if xCount >= 500 { - t.Errorf("expected blocker to be truncated, got %d 'X' chars", xCount) - } - - if xCount > blockerMaxChars { - t.Errorf("expected at most %d blocker chars, got %d", blockerMaxChars, xCount) - } -} - -func TestBuildPlanReminder(t *testing.T) { - tests := []struct { - name string - - status string - - wantOK bool - - wantSubstr string - }{ - {"interviewing", "interviewing", true, "interviewing the user"}, - - {"review", "review", true, "under review"}, - - {"executing returns false", "executing", false, ""}, - - {"empty returns false", "", false, ""}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - msg, ok := buildPlanReminder(tt.status) - - if ok != tt.wantOK { - t.Fatalf("buildPlanReminder(%q) ok = %v, want %v", tt.status, ok, tt.wantOK) - } - - if !ok { - return - } - - if msg.Role != "user" { - t.Errorf("expected role 'user', got %q", msg.Role) - } - - if !strings.Contains(msg.Content, tt.wantSubstr) { - t.Errorf("expected content to contain %q, got %q", tt.wantSubstr, msg.Content) - } - }) - } -} - -// ---------- /plan command tests ---------- - -func newTestAgentLoop(t *testing.T) (*AgentLoop, func()) { - t.Helper() - - tmpDir, err := os.MkdirTemp("", "agent-plan-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - msgBus := bus.NewMessageBus() - - provider := &mockProvider{} - - al := NewAgentLoop(cfg, msgBus, provider) - - return al, func() { os.RemoveAll(tmpDir) } -} - -func TestPlanCommand_ShowNoPlan(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - response, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan"}) - - if !handled { - t.Fatal("expected /plan to be handled") - } - - if !strings.Contains(response, "No active plan") { - t.Errorf("expected 'No active plan', got %q", response) - } -} - -func TestSplitChatAndThread(t *testing.T) { - tests := []struct { - name string - - chatID string - - wantChatID string - - wantThread int - }{ - {name: "plain chat", chatID: "-100123", wantChatID: "-100123", wantThread: 0}, - - {name: "chat with thread", chatID: "-100123/77", wantChatID: "-100123", wantThread: 77}, - - {name: "invalid thread", chatID: "-100123/abc", wantChatID: "-100123", wantThread: 0}, - - {name: "empty", chatID: "", wantChatID: "", wantThread: 0}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - gotChatID, gotThread := splitChatAndThread(tt.chatID) - - if gotChatID != tt.wantChatID || gotThread != tt.wantThread { - t.Fatalf( - - "splitChatAndThread(%q) = (%q, %d), want (%q, %d)", - - tt.chatID, - - gotChatID, - - gotThread, - - tt.wantChatID, - - tt.wantThread, - ) - } - }) - } -} - -func TestHeartbeatCommandThreadHerePersistsConfig(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - var saved bool - - var updatedThread int - - al.SetConfigSaver(func(cfg *config.Config) error { - saved = true - - if cfg.Channels.Telegram.HeartbeatThreadID != 42 { - t.Fatalf("HeartbeatThreadID in saver = %d, want 42", cfg.Channels.Telegram.HeartbeatThreadID) - } - - return nil - }) - - al.SetHeartbeatThreadUpdater(func(threadID int) { updatedThread = threadID }) - - msg := bus.InboundMessage{ - Content: "/heartbeat thread here", - - Channel: "telegram", - - ChatID: "-100500/42", - } - - resp, handled := al.handleCommand(context.Background(), msg) - - if !handled { - t.Fatal("expected /heartbeat command to be handled") - } - - if !strings.Contains(resp, "Heartbeat thread set to 42") { - t.Fatalf("unexpected response: %q", resp) - } - - if !saved { - t.Fatal("expected config saver to be called") - } - - if updatedThread != 42 { - t.Fatalf("updatedThread = %d, want 42", updatedThread) - } - - if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 42 { - t.Fatalf("cfg heartbeat thread = %d, want 42", got) - } - - if got := al.state.GetHeartbeatTarget(); got != "telegram:-100500" { - t.Fatalf("state heartbeat target = %q, want %q", got, "telegram:-100500") - } -} - -func TestHeartbeatCommandThreadOff(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - al.cfg.Channels.Telegram.HeartbeatThreadID = 99 - - resp, handled := al.handleCommand(context.Background(), bus.InboundMessage{ - Content: "/heartbeat thread off", - - Channel: "telegram", - - ChatID: "-100500/42", - }) - - if !handled { - t.Fatal("expected /heartbeat command to be handled") - } - - if !strings.Contains(resp, "disabled") { - t.Fatalf("unexpected response: %q", resp) - } - - if got := al.cfg.Channels.Telegram.HeartbeatThreadID; got != 0 { - t.Fatalf("cfg heartbeat thread = %d, want 0", got) - } -} - -func TestPlanCommand_StartNewPlan(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - // /plan should NOT be handled by handleCommand — it falls through - - // to the LLM queue via expandPlanCommand. - - _, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan Set up monitoring"}) - - if handled { - t.Fatal("expected /plan NOT to be handled (should fall through to LLM)") - } - - // expandPlanCommand writes the seed and rewrites the message - - msg := bus.InboundMessage{Content: "/plan Set up monitoring"} - - expanded, compact, ok := al.expandPlanCommand(msg) - - if !ok { - t.Fatal("expected expandPlanCommand to succeed") - } - - if expanded != "Set up monitoring" { - t.Errorf("expected expanded = 'Set up monitoring', got %q", expanded) - } - - if !strings.Contains(compact, "Set up monitoring") { - t.Errorf("expected compact to contain task, got %q", compact) - } - - // Verify plan was created - - agent := al.registry.GetDefaultAgent() - - if !agent.ContextBuilder.HasActivePlan() { - t.Error("expected active plan after expandPlanCommand") - } - - if status := agent.ContextBuilder.GetPlanStatus(); status != "interviewing" { - t.Errorf("expected 'interviewing', got %q", status) - } -} - -func TestPlanCommand_StartBlockedByExisting(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - // Start first plan via expandPlanCommand - - al.expandPlanCommand(bus.InboundMessage{Content: "/plan First task"}) - - // Try to start another — handleCommand should block it on the fast path - - response, handled := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan Second task"}) - - if !handled { - t.Fatal("expected second /plan to be handled (blocked)") - } - - if !strings.Contains(response, "already active") { - t.Errorf("expected 'already active', got %q", response) - } -} - -func TestPlanCommand_Clear(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - // Start plan then clear - - al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan clear"}) - - if !strings.Contains(response, "Plan cleared") { - t.Errorf("expected 'Plan cleared', got %q", response) - } - - agent := al.registry.GetDefaultAgent() - - if agent.ContextBuilder.HasActivePlan() { - t.Error("expected no plan after clear") - } -} - -func TestPlanCommand_ClearNoPlan(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan clear"}) - - if !strings.Contains(response, "No active plan") { - t.Errorf("expected 'No active plan', got %q", response) - } -} - -func TestPlanCommand_Start(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - // Create interviewing plan with phases (start requires phases) - - plan := "# Active Plan\n\n> Task: Test task\n> Status: interviewing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" - - _ = agent.ContextBuilder.WriteMemory(plan) - - // Transition to executing via /plan start - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) - - if !strings.Contains(response, "approved") { - t.Errorf("expected 'approved', got %q", response) - } - - if status := agent.ContextBuilder.GetPlanStatus(); status != "executing" { - t.Errorf("expected 'executing', got %q", status) - } - - // planStartPending must be set so Run() enqueues an LLM trigger - - if !al.planStartPending { - t.Error("expected planStartPending to be true after /plan start") - } -} - -func TestPlanCommand_StartFromReview(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - // Create a plan in review status - - plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" - - _ = agent.ContextBuilder.WriteMemory(plan) - - // Approve via /plan start - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) - - if !strings.Contains(response, "approved") { - t.Errorf("expected 'approved', got %q", response) - } - - if status := agent.ContextBuilder.GetPlanStatus(); status != "executing" { - t.Errorf("expected 'executing', got %q", status) - } - - if !al.planStartPending { - t.Error("expected planStartPending to be true after /plan start from review") - } -} - -func TestPlanCommand_StartNoPhases(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - // Create interviewing plan without phases - - al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) - - // Should be blocked because no phases exist - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) - - if !strings.Contains(response, "no phases") { - t.Errorf("expected 'no phases' error, got %q", response) - } - - agent := al.registry.GetDefaultAgent() - - if status := agent.ContextBuilder.GetPlanStatus(); status != "interviewing" { - t.Errorf("expected status to remain 'interviewing', got %q", status) - } - - if al.planStartPending { - t.Error("planStartPending must not be set when start is rejected (no phases)") - } -} - -func TestPlanCommand_StartAlreadyExecuting(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - // Create interviewing plan with phases, then start - - plan := "# Active Plan\n\n> Task: Test task\n> Status: interviewing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" - - _ = agent.ContextBuilder.WriteMemory(plan) - - al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) - - // Clear the flag from the first call (simulating Run() consuming it) - - al.planStartPending = false - - // Try start again — should be rejected - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan start"}) - - if !strings.Contains(response, "already executing") { - t.Errorf("expected 'already executing', got %q", response) - } - - if al.planStartPending { - t.Error("planStartPending must not be set when plan is already executing") - } -} - -func TestPlanCommand_Done(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - // Write a plan directly with phases - - plan := `# Active Plan - - - -> Task: Test task - -> Status: executing - -> Phase: 1 - - - -## Phase 1: Setup - -- [ ] Step one - -- [ ] Step two - - - -## Context - -Test context - -` - - agent.ContextBuilder.WriteMemory(plan) - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan done 1"}) - - if !strings.Contains(response, "Marked step 1") { - t.Errorf("expected confirmation, got %q", response) - } -} - -func TestPlanCommand_DoneInvalidStep(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - al.expandPlanCommand(bus.InboundMessage{Content: "/plan Test task"}) - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan done abc"}) - - if !strings.Contains(response, "positive integer") { - t.Errorf("expected step validation error, got %q", response) - } -} - -func TestPlanCommand_Add(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - plan := `# Active Plan - - - -> Task: Test task - -> Status: executing - -> Phase: 1 - - - -## Phase 1: Setup - -- [ ] Step one - - - -## Context - -Test context - -` - - agent.ContextBuilder.WriteMemory(plan) - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan add New step here"}) - - if !strings.Contains(response, "Added step") { - t.Errorf("expected 'Added step', got %q", response) - } - - content := agent.ContextBuilder.ReadMemory() - - if !strings.Contains(content, "New step here") { - t.Error("expected new step in plan content") - } -} - -func TestPlanCommand_Next(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - plan := `# Active Plan - - - -> Task: Test task - -> Status: executing - -> Phase: 1 - - - -## Phase 1: Setup - -- [x] Step one - - - -## Phase 2: Deploy - -- [ ] Step two - - - -## Context - -Test - -` - - agent.ContextBuilder.WriteMemory(plan) - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan next"}) - - if !strings.Contains(response, "phase 2") { - t.Errorf("expected 'phase 2', got %q", response) - } - - if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 2 { - t.Errorf("expected phase 2, got %d", phase) - } -} - -func TestPlanCommand_ShowActivePlan(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - plan := `# Active Plan - - - -> Task: Deploy app - -> Status: executing - -> Phase: 1 - - - -## Phase 1: Build - -- [x] Compile code - -- [ ] Run tests - - - -## Context - -Production server - -` - - agent.ContextBuilder.WriteMemory(plan) - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{Content: "/plan"}) - - if !strings.Contains(response, "Deploy app") { - t.Errorf("expected task name in display, got %q", response) - } - - if !strings.Contains(response, "Phase 1") { - t.Errorf("expected phase info in display, got %q", response) - } -} - -// TestAutoPhaseAdvance verifies that auto-advance sends notification after LLM iteration - -// when current phase is complete. - -func TestAutoPhaseAdvance(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-auto-advance-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - msgBus := bus.NewMessageBus() - - provider := &simpleMockProvider{response: "OK"} - - al := NewAgentLoop(cfg, msgBus, provider) - - agent := al.registry.GetDefaultAgent() - - if agent == nil { - t.Fatal("No default agent") - } - - // Write plan with phase 1 complete - - plan := `# Active Plan - - - -> Task: Test auto advance - -> Status: executing - -> Phase: 1 - - - -## Phase 1: Setup - -- [x] Step one - -- [x] Step two - - - -## Phase 2: Deploy - -- [ ] Step three - - - -## Context - -Test - -` - - agent.ContextBuilder.WriteMemory(plan) - - // Process a message which triggers runAgentLoop - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - - defer cancel() - - _, err = al.ProcessDirectWithChannel(ctx, "continue", "auto-advance-test", "test", "chat1") - if err != nil { - t.Fatalf("ProcessDirectWithChannel failed: %v", err) - } - - // After processing, phase should be auto-advanced - - if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 2 { - t.Errorf("expected phase auto-advanced to 2, got %d", phase) - } -} - -// TestAutoCompleteClears verifies that plan is marked completed with correct phase when all phases are complete. - -func TestAutoCompleteClears(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-auto-complete-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - msgBus := bus.NewMessageBus() - - provider := &simpleMockProvider{response: "All done"} - - al := NewAgentLoop(cfg, msgBus, provider) - - agent := al.registry.GetDefaultAgent() - - if agent == nil { - t.Fatal("No default agent") - } - - // Write fully complete plan - - plan := `# Active Plan - - - -> Task: Test auto complete - -> Status: executing - -> Phase: 1 - - - -## Phase 1: Setup - -- [x] Step one - -- [x] Step two - - - -## Context - -Test - -` - - agent.ContextBuilder.WriteMemory(plan) - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - - defer cancel() - - _, err = al.ProcessDirectWithChannel(ctx, "finish up", "auto-complete-test", "test", "chat1") - if err != nil { - t.Fatalf("ProcessDirectWithChannel failed: %v", err) - } - - // Plan should be kept with status "completed" and phase set to total - - if !agent.ContextBuilder.HasActivePlan() { - t.Error("expected plan to be retained after completion") - } - - if status := agent.ContextBuilder.GetPlanStatus(); status != "completed" { - t.Errorf("expected plan status 'completed', got %q", status) - } - - if phase := agent.ContextBuilder.GetCurrentPhase(); phase != 1 { - t.Errorf("expected phase 1 (total phases), got %d", phase) - } -} - -func TestIsToolAllowedDuringInterview_FuzzyNames(t *testing.T) { - tests := []struct { - name string - - args map[string]any - - want bool - }{ - // Exact names — read tools allowed - - {"read_file", nil, true}, - - {"list_dir", nil, true}, - - {"web_search", nil, true}, - - {"web_fetch", nil, true}, - - // Fuzzy variants — should also be allowed - - {"readfile", nil, true}, - - {"ReadFile", nil, true}, - - {"listdir", nil, true}, - - {"websearch", nil, true}, - - {"webfetch", nil, true}, - - // Message tool — allowed (needed for interview questions) - - {"message", nil, true}, - - {"Message", nil, true}, - - // Write to MEMORY.md — allowed - - {"edit_file", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, - - {"editfile", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, - - {"EditFile", map[string]any{"path": "/ws/memory/MEMORY.md"}, true}, - - // Write to non-MEMORY.md — blocked - - {"edit_file", map[string]any{"path": "/ws/main.go"}, false}, - - {"editfile", map[string]any{"path": "/ws/main.go"}, false}, - - // exec — read-only commands allowed - - {"exec", map[string]any{"command": "find . -name '*.py'"}, true}, - - {"exec", map[string]any{"command": "ls -la"}, true}, - - {"exec", map[string]any{"command": "grep -r TODO ."}, true}, - - {"exec", map[string]any{"command": "cat README.md"}, true}, - - // exec — cd prefix stripped - - {"exec", map[string]any{"command": "cd /home/user/project && find . -type f"}, true}, - - {"exec", map[string]any{"command": "cd /tmp && rm -rf *"}, false}, - - // exec — write operators blocked - - {"exec", map[string]any{"command": "find . > output.txt"}, false}, - - {"exec", map[string]any{"command": "ls -la >> log.txt"}, false}, - - {"exec", map[string]any{"command": "cat foo | tee bar.txt"}, false}, - - // exec — path traversal blocked - - {"exec", map[string]any{"command": "cat ../../etc/passwd"}, false}, - - {"exec", map[string]any{"command": "find ../../"}, false}, - - {"exec", map[string]any{"command": "ls ../secret"}, false}, - - // exec — absolute paths blocked - - {"exec", map[string]any{"command": "cat /etc/passwd"}, false}, - - {"exec", map[string]any{"command": "find /etc -name '*.conf'"}, false}, - - {"exec", map[string]any{"command": "ls /root"}, false}, - - // exec — write commands blocked - - {"exec", map[string]any{"command": "rm -rf /"}, false}, - - {"exec", map[string]any{"command": "mv a b"}, false}, - - // exec — no args / empty command blocked - - {"exec", nil, false}, - - {"exec", map[string]any{"command": ""}, false}, - - {"Exec", nil, false}, - } - - for _, tt := range tests { - got := isToolAllowedDuringInterview(tt.name, tt.args) - - if got != tt.want { - t.Errorf("isToolAllowedDuringInterview(%q, %v) = %v, want %v", tt.name, tt.args, got, tt.want) - } - } -} - -func TestBuildArgsSnippet_ExecStripsCD(t *testing.T) { - tests := []struct { - name string - - tool string - - args map[string]any - - workspace string - - wantSnip string - }{ - { - name: "exec strips cd prefix", - - tool: "exec", - - args: map[string]any{ - "command": "cd /home/user/workspace/project/my-projects && pytest tests/test_integration.py", - }, - - workspace: "/home/user/workspace", - - wantSnip: "pytest tests/test_integration.py", - }, - - { - name: "exec no cd prefix, flags stripped", - - tool: "exec", - - args: map[string]any{"command": "ls -la"}, - - workspace: "/ws", - - wantSnip: "ls", - }, - - { - name: "exec empty command", - - tool: "exec", - - args: map[string]any{}, - - workspace: "/ws", - - wantSnip: "{}", - }, - - { - name: "read_file strips workspace", - - tool: "read_file", - - args: map[string]any{"path": "/home/user/workspace/src/main.go"}, - - workspace: "/home/user/workspace", - - wantSnip: "src/main.go", - }, - - { - name: "edit_file shows path", - - tool: "edit_file", - - args: map[string]any{"path": "/ws/config.json", "old_text": "old value here"}, - - workspace: "/ws", - - wantSnip: "config.json", - }, - - { - name: "file tool long path prioritizes filename", - - tool: "read_file", - - args: map[string]any{ - "path": "/ws/projects/terra-py-form/src/terra_py_form/hot/state/backend.py", - }, - - workspace: "/ws", - - wantSnip: "projects/terra-py-form/src/terra_py_form/hot/sta\u2026/backend.py", - }, - - { - name: "unknown tool shows raw JSON", - - tool: "web_search", - - args: map[string]any{"query": "hello"}, - - workspace: "/ws", - - wantSnip: `{"query":"hello"}`, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := buildArgsSnippet(tt.tool, tt.args, tt.workspace) - - if got != tt.wantSnip { - t.Errorf("buildArgsSnippet(%q) = %q, want %q", tt.tool, got, tt.wantSnip) - } - }) - } -} - -func TestFormatCompactEntry(t *testing.T) { - tests := []struct { - name string - - entry toolLogEntry - - wantSub string // must be a substring - - wantMark string // result marker must appear - - noTime bool // if true, duration should NOT appear - }{ - { - name: "exec short entry", - - entry: toolLogEntry{Name: "[1] exec", ArgsSnip: "ls", Result: "✓ 1.0s"}, - - wantSub: "exec ls", - - wantMark: "✓ 1.0s", // exec keeps duration - - }, - - { - name: "exec long entry truncated from end", - - entry: toolLogEntry{ - Name: "[2] exec", - - ArgsSnip: "pytest tests/integration/test_very_long_name.py", - - Result: "✗ 3.0s", - }, - - wantMark: "✗", - }, - - { - name: "file tool omits duration, shows filename", - - entry: toolLogEntry{ - Name: "[3] edit_file", - - ArgsSnip: "projects/terra/src/deep/nested/backend.py", - - Result: "✓ 0.0s", - }, - - wantSub: "backend.py", - - wantMark: "✓", - - noTime: true, - }, - - { - name: "file tool path truncates from start", - - entry: toolLogEntry{ - Name: "[4] read_file", - - ArgsSnip: "projects/terra-py-form/src/terra_py_form/hot/state/backend.py", - - Result: "✓ 0.1s", - }, - - wantSub: "backend.py", - - wantMark: "✓", - - noTime: true, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := formatCompactEntry(tt.entry) - - if tt.wantSub != "" && !strings.Contains(got, tt.wantSub) { - t.Errorf("expected to contain %q, got: %q", tt.wantSub, got) - } - - if !strings.Contains(got, tt.wantMark) { - t.Errorf("result marker %q missing from: %q", tt.wantMark, got) - } - - if tt.noTime && strings.Contains(got, "0s") { - t.Errorf("file tool should omit duration, got: %q", got) - } - - // Must not exceed maxEntryLineWidth - - if runeLen := len([]rune(got)); runeLen > maxEntryLineWidth { - t.Errorf("entry too wide: %d runes (max %d): %q", runeLen, maxEntryLineWidth, got) - } - }) - } -} - -func TestBuildRichStatus(t *testing.T) { - task := &activeTask{ - Iteration: 3, - - MaxIter: 20, - - toolLog: []toolLogEntry{ - {Name: "exec", ArgsSnip: "ls -la", Result: "✓ 1.2s"}, - - {Name: "exec", ArgsSnip: "pytest tests/", Result: "✓ 5.0s"}, - - {Name: "read_file", ArgsSnip: "src/main.go", Result: "⏳"}, - }, - } - - got := buildRichStatus(task, false, "/home/user/my-projects") - - mustContain := []string{ - "Task in progress (3/20)", - - "my-projects", - - "read_file", // latest entry - - "No errors", // no error yet - - } - - for _, s := range mustContain { - if !strings.Contains(got, s) { - t.Errorf("expected output to contain %q, got:\n%s", s, got) - } - } - - // Non-background: should NOT have reply prompt - - if strings.Contains(got, "Reply to intervene") { - t.Error("non-background task should not have reply prompt") - } - - // Background: should have reply prompt - - bgGot := buildRichStatus(task, true, "/home/user/my-projects") - - if !strings.Contains(bgGot, "Reply to intervene") { - t.Error("background task should have reply prompt") - } -} - -func TestBuildRichStatus_ProjectDir(t *testing.T) { - // exec-based projectDir takes priority - - task := &activeTask{ - Iteration: 1, - - MaxIter: 10, - - projectDir: "terra-py-form", - - toolLog: []toolLogEntry{ - {Name: "exec", ArgsSnip: "ls", Result: "✓ 0.1s"}, - }, - } - - got := buildRichStatus(task, false, "/home/user/.picoclaw/workspace") - - if !strings.Contains(got, "terra-py-form") { - t.Errorf("expected projectDir in output, got:\n%s", got) - } - - // fileCommonDir fallback - - task2 := &activeTask{ - Iteration: 1, - - MaxIter: 10, - - fileCommonDir: "projects/terra-py-form", - - toolLog: []toolLogEntry{ - {Name: "read_file", ArgsSnip: "src/main.py", Result: "✓ 0.1s"}, - }, - } - - got2 := buildRichStatus(task2, false, "/home/user/.picoclaw/workspace") - - if !strings.Contains(got2, "terra-py-form") { - t.Errorf("expected fileCommonDir basename in output, got:\n%s", got2) - } - - // workspace basename fallback with trailing slash - - task3 := &activeTask{ - Iteration: 1, - - MaxIter: 10, - - toolLog: []toolLogEntry{ - {Name: "exec", ArgsSnip: "ls", Result: "✓ 0.1s"}, - }, - } - - for _, ws := range []string{"/home/user/my-project/", "/home/user/my-project"} { - got := buildRichStatus(task3, false, ws) - - if !strings.Contains(got, "my-project") { - t.Errorf("workspace %q: expected 'my-project' in output, got:\n%s", ws, got) - } - } -} - -func TestExtractExecProjectDir(t *testing.T) { - tests := []struct { - name string - - cmd string - - want string - }{ - {"cd deep path", "cd /ws/projects/terra-py-form && pytest", "terra-py-form"}, - - {"cd direct subdir", "cd /ws/my-app && make build", "my-app"}, - - {"cd trailing slash", "cd /ws/my-app/ && ls", "my-app"}, - - {"cd to workspace", "cd /ws && ls", "ws"}, - - {"no cd prefix", "pytest tests/", ""}, - - {"empty command", "", ""}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - args := map[string]any{"command": tt.cmd} - - got := extractExecProjectDir(args) - - if got != tt.want { - t.Errorf("extractExecProjectDir(%q) = %q, want %q", tt.cmd, got, tt.want) - } - }) - } -} - -func TestFileParentRelDir(t *testing.T) { - ws := "/home/user/.picoclaw/workspace" - - tests := []struct { - name string - - path string - - want string - }{ - {"deep path", ws + "/projects/terra/src/main.py", "projects/terra/src"}, - - {"direct subdir", ws + "/my-app/README.md", "my-app"}, - - {"workspace root file", ws + "/notes.txt", ""}, - - {"outside workspace", "/tmp/foo.txt", ""}, - - {"trailing slash ws", ws + "/projects/terra/src/main.py", "projects/terra/src"}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := fileParentRelDir(tt.path, ws) - - if got != tt.want { - t.Errorf("fileParentRelDir(%q, ws) = %q, want %q", tt.path, got, tt.want) - } - }) - } -} - -func TestCommonDirPrefix(t *testing.T) { - tests := []struct { - name string - - a, b string - - want string - }{ - {"same dir", "projects/terra/src", "projects/terra/src", "projects/terra/src"}, - - {"converge to project", "projects/terra/src", "projects/terra/tests", "projects/terra"}, - - {"converge to top", "projects/terra/src", "projects/other/tests", "projects"}, - - {"no common", "aaa/bbb", "ccc/ddd", ""}, - - {"one is prefix", "projects/terra", "projects/terra/src", "projects/terra"}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := commonDirPrefix(tt.a, tt.b) - - if got != tt.want { - t.Errorf("commonDirPrefix(%q, %q) = %q, want %q", tt.a, tt.b, got, tt.want) - } - }) - } -} - -func TestDisplayProjectDir(t *testing.T) { - // exec projectDir wins - - task1 := &activeTask{projectDir: "my-app", fileCommonDir: "projects/other"} - - if got := displayProjectDir(task1); got != "my-app" { - t.Errorf("expected 'my-app', got %q", got) - } - - // fileCommonDir fallback: basename - - task2 := &activeTask{fileCommonDir: "projects/terra-py-form"} - - if got := displayProjectDir(task2); got != "terra-py-form" { - t.Errorf("expected 'terra-py-form', got %q", got) - } - - // single component - - task3 := &activeTask{fileCommonDir: "my-app"} - - if got := displayProjectDir(task3); got != "my-app" { - t.Errorf("expected 'my-app', got %q", got) - } - - // empty - - task4 := &activeTask{} - - if got := displayProjectDir(task4); got != "" { - t.Errorf("expected empty, got %q", got) - } -} - -func TestBuildRichStatus_FixedHeight(t *testing.T) { - // Test that output has the same number of lines regardless of entry count - - countLines := func(s string) int { - return strings.Count(s, "\n") - } - - // 0 entries - - task0 := &activeTask{Iteration: 1, MaxIter: 10} - - lines0 := countLines(buildRichStatus(task0, true, "/ws/p")) - - // 1 entry - - task1 := &activeTask{ - Iteration: 1, MaxIter: 10, - - toolLog: []toolLogEntry{{Name: "exec", ArgsSnip: "ls", Result: "⏳"}}, - } - - lines1 := countLines(buildRichStatus(task1, true, "/ws/p")) - - // 5 entries - - task5 := &activeTask{Iteration: 5, MaxIter: 10} - - for i := 0; i < 5; i++ { - task5.toolLog = append(task5.toolLog, toolLogEntry{ - Name: fmt.Sprintf("[%d] exec", i), ArgsSnip: "cmd", Result: "✓ 1.0s", - }) - } - - lines5 := countLines(buildRichStatus(task5, true, "/ws/p")) - - // 5 entries + sticky error - - task5err := &activeTask{Iteration: 5, MaxIter: 10} - - for i := 0; i < 5; i++ { - task5err.toolLog = append(task5err.toolLog, toolLogEntry{ - Name: fmt.Sprintf("[%d] exec", i), ArgsSnip: "cmd", Result: "✓ 1.0s", - }) - } - - errEntry := toolLogEntry{ - Name: "[3] exec", ArgsSnip: "pytest", Result: "✗ 2.0s", - - ErrDetail: "FAILED test\nExit code: 1", - } - - task5err.lastError = &errEntry - - lines5err := countLines(buildRichStatus(task5err, true, "/ws/p")) - - if lines0 != lines1 || lines1 != lines5 || lines5 != lines5err { - t.Errorf("line counts should be equal: 0=%d, 1=%d, 5=%d, 5+err=%d", - - lines0, lines1, lines5, lines5err) - } -} - -func TestBuildRichStatus_StickyError(t *testing.T) { - // Error from a past entry sticks in the error section - - errEntry := toolLogEntry{ - Name: "[2] exec", ArgsSnip: "pytest", Result: "✗ 3.2s", - - ErrDetail: "FAILED test_login\nExit code: 1", - } - - task := &activeTask{ - Iteration: 5, - - MaxIter: 10, - - toolLog: []toolLogEntry{ - {Name: "[3] read_file", ArgsSnip: "src/auth.py", Result: "✓ 0.1s"}, - - {Name: "[4] edit_file", ArgsSnip: "src/auth.py", Result: "✓ 0.2s"}, - - {Name: "[5] exec", ArgsSnip: "pytest --retry", Result: "⏳"}, - }, - - lastError: &errEntry, - } - - got := buildRichStatus(task, false, "/ws/p") - - // Error section should show the sticky error in code block - - if !strings.Contains(got, "FAILED test_login") { - t.Errorf("expected sticky error detail in error section, got:\n%s", got) - } - - if !strings.Contains(got, "\u274C") { // ❌ - t.Errorf("expected ❌ error header, got:\n%s", got) - } - - // Latest entry is NOT the error - - if !strings.Contains(got, "pytest --retry") { - t.Errorf("expected latest entry command, got:\n%s", got) - } -} - -func TestBuildRichStatus_LatestEntryNoInlineResult(t *testing.T) { - longCmd := "uv run pytest tests/hot/test_state_backend_integration.py" - - task := &activeTask{ - Iteration: 2, - - MaxIter: 10, - - toolLog: []toolLogEntry{ - {Name: "exec", ArgsSnip: "ls -la", Result: "\u2713 0.5s"}, - - {Name: "exec", ArgsSnip: longCmd, Result: "\u23F3"}, - }, - } - - got := buildRichStatus(task, false, "/ws/my-project") - - // Latest entry shows command (possibly truncated) with filename visible - - if !strings.Contains(got, "integration.py") { - t.Errorf("latest entry should show filename, got:\n%s", got) - } - - // Result on separate indented line - - if !strings.Contains(got, " \u23F3") { - t.Errorf("latest entry result should be on indented line, got:\n%s", got) - } - - // Project name shown - - if !strings.Contains(got, "my-project") { - t.Errorf("should show project name, got:\n%s", got) - } - - // No second separator before error section - - lines := strings.Split(got, "\n") - - sepCount := 0 - - for _, l := range lines { - if strings.HasPrefix(l, "\u2501") { - sepCount++ - } - } - - if sepCount != 1 { - t.Errorf("expected exactly 1 separator, got %d in:\n%s", sepCount, got) - } -} - -func TestSanitizeHistoryForProvider_MultiToolCall(t *testing.T) { - // Regression: assistant with 2+ tool_calls had 2nd+ tool results dropped - - // because the check only allowed tool after assistant, not after sibling tool. - - history := []providers.Message{ - {Role: "user", Content: "hello"}, - - {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ - {ID: "a", Function: &providers.FunctionCall{Name: "exec"}}, - - {ID: "b", Function: &providers.FunctionCall{Name: "read_file"}}, - }}, - - {Role: "tool", Content: "ok", ToolCallID: "a"}, - - {Role: "tool", Content: "ok", ToolCallID: "b"}, - - {Role: "assistant", Content: "done"}, - } - - got := sanitizeHistoryForProvider(history) - - // All 5 messages must survive - - if len(got) != 5 { - roles := make([]string, len(got)) - - for i, m := range got { - roles[i] = m.Role - } - - t.Fatalf("expected 5 messages, got %d: %v", len(got), roles) - } - - // Verify both tool results present - - toolCount := 0 - - for _, m := range got { - if m.Role == "tool" { - toolCount++ - } - } - - if toolCount != 2 { - t.Errorf("expected 2 tool results, got %d", toolCount) - } -} - -// ---------- plan nudge tests ---------- - -// countingMockProvider counts Chat calls and always returns text-only responses. - -type countingMockProvider struct { - callCount int -} - -func (m *countingMockProvider) Chat( - ctx context.Context, - - messages []providers.Message, - - tools []providers.ToolDefinition, - - model string, - - opts map[string]any, -) (*providers.LLMResponse, error) { - m.callCount++ - - return &providers.LLMResponse{ - Content: fmt.Sprintf("Response %d", m.callCount), - - ToolCalls: []providers.ToolCall{}, - }, nil -} - -func (m *countingMockProvider) GetDefaultModel() string { - return "mock-counting-model" -} - -func TestPlanNudge_ForegroundExecution(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - provider := &countingMockProvider{} - - msgBus := bus.NewMessageBus() - - al := NewAgentLoop(cfg, msgBus, provider) - - agent := al.registry.GetDefaultAgent() - - if agent == nil { - t.Fatal("no default agent") - } - - // Write a plan in executing status with unchecked steps - - plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n\n## Context\n" - - agent.ContextBuilder.WriteMemory(plan) - - // Process a foreground message (no background metadata) - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - - defer cancel() - - msg := bus.InboundMessage{ - Channel: "test", - - SenderID: "user1", - - ChatID: "chat1", - - Content: "continue working", - - SessionKey: "nudge-test", - } - - _, err = al.processMessage(ctx, msg) - if err != nil { - t.Fatalf("processMessage failed: %v", err) - } - - // The provider should have been called at least 2 times: - - // 1st call: returns text-only → nudge fires (unchecked steps remain) - - // 2nd call: returns text-only → nudge already fired, loop exits - - if provider.callCount < 2 { - t.Errorf("expected at least 2 provider calls (nudge should trigger continuation), got %d", provider.callCount) - } -} - -func TestPlanNudge_NoNudgeWhenAllStepsComplete(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - provider := &countingMockProvider{} - - msgBus := bus.NewMessageBus() - - al := NewAgentLoop(cfg, msgBus, provider) - - agent := al.registry.GetDefaultAgent() - - if agent == nil { - t.Fatal("no default agent") - } - - // Write a plan where all steps are already checked - - plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [x] Step one\n- [x] Step two\n\n## Context\n" - - agent.ContextBuilder.WriteMemory(plan) - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - - defer cancel() - - msg := bus.InboundMessage{ - Channel: "test", - - SenderID: "user1", - - ChatID: "chat1", - - Content: "all done", - - SessionKey: "nudge-test-complete", - } - - _, err = al.processMessage(ctx, msg) - if err != nil { - t.Fatalf("processMessage failed: %v", err) - } - - // No unchecked steps → preUnchecked=0 → no nudge → only 1 provider call - - if provider.callCount != 1 { - t.Errorf("expected exactly 1 provider call (no nudge needed), got %d", provider.callCount) - } -} - -func TestPlanNudge_ProgressMessage(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-nudge-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - - MaxToolIterations: 10, - }, - }, - } - - // Provider that checks the nudge message content on the 2nd call - - var nudgeContent string - - provider := &nudgeCaptureMockProvider{onSecondCall: func(msgs []providers.Message) { - // The last user message should be the nudge - - for i := len(msgs) - 1; i >= 0; i-- { - if msgs[i].Role == "user" { - nudgeContent = msgs[i].Content - - break - } - } - }} - - msgBus := bus.NewMessageBus() - - al := NewAgentLoop(cfg, msgBus, provider) - - agent := al.registry.GetDefaultAgent() - - if agent == nil { - t.Fatal("no default agent") - } - - // Write a plan with 3 unchecked steps; the provider edits memory to - - // mark one step between calls (simulated by the first-call hook). - - plan := "# Active Plan\n\n> Task: Test\n> Status: executing\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n- [ ] Step two\n- [ ] Step three\n\n## Context\n" - - agent.ContextBuilder.WriteMemory(plan) - - // After the first LLM response (no tool calls), simulate that - - // one step was marked [x] externally (as if the AI did it via tool). - - // We do this by hooking the provider's first call to mutate memory. - - provider.onFirstCall = func() { - updated := strings.Replace(agent.ContextBuilder.ReadMemory(), "- [ ] Step one", "- [x] Step one", 1) - - agent.ContextBuilder.WriteMemory(updated) - } - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - - defer cancel() - - msg := bus.InboundMessage{ - Channel: "test", - - SenderID: "user1", - - ChatID: "chat1", - - Content: "work on the plan", - - SessionKey: "nudge-progress-test", - } - - _, err = al.processMessage(ctx, msg) - if err != nil { - t.Fatalf("processMessage failed: %v", err) - } - - // Should have gotten the "Progress recorded" nudge (not the "none were marked" one) - - if !strings.Contains(nudgeContent, "Progress recorded") { - t.Errorf("expected 'Progress recorded' nudge, got %q", nudgeContent) - } - - if !strings.Contains(nudgeContent, "2 unchecked steps remain") { - t.Errorf("expected '2 unchecked steps remain' in nudge, got %q", nudgeContent) - } -} - -// nudgeCaptureMockProvider calls hooks on 1st and 2nd Chat invocations. - -type nudgeCaptureMockProvider struct { - callCount int - - onFirstCall func() - - onSecondCall func([]providers.Message) -} - -func (m *nudgeCaptureMockProvider) Chat( - ctx context.Context, - - messages []providers.Message, - - tools []providers.ToolDefinition, - - model string, - - opts map[string]any, -) (*providers.LLMResponse, error) { - m.callCount++ - - if m.callCount == 1 && m.onFirstCall != nil { - m.onFirstCall() - } - - if m.callCount == 2 && m.onSecondCall != nil { - m.onSecondCall(messages) - } - - return &providers.LLMResponse{ - Content: fmt.Sprintf("Response %d", m.callCount), - - ToolCalls: []providers.ToolCall{}, - }, nil -} - -func (m *nudgeCaptureMockProvider) GetDefaultModel() string { - return "mock-nudge-model" -} - -// --- consumeStreamWithRepetitionDetection tests --- - -func TestConsumeStream_NormalCompletion(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 8) - - go func() { - ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} - - ch <- protocoltypes.StreamEvent{ContentDelta: "world!"} - - ch <- protocoltypes.StreamEvent{ - FinishReason: "stop", - - Usage: &providers.UsageInfo{PromptTokens: 5, CompletionTokens: 2, TotalTokens: 7}, - } - - close(ch) - }() - - ctx, cancel := context.WithCancel(context.Background()) - - defer cancel() - - resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if detected { - t.Fatal("expected detected=false for normal content") - } - - if resp.Content != "Hello world!" { - t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") - } - - if resp.FinishReason != "stop" { - t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") - } - - if resp.Usage == nil || resp.Usage.TotalTokens != 7 { - t.Errorf("Usage.TotalTokens = %v, want 7", resp.Usage) - } - - _ = ctx // keep linter happy -} - -func TestConsumeStream_DetectsRepetition(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 64) - - cancelCalled := false - - ctx, cancel := context.WithCancel(context.Background()) - - wrappedCancel := func() { - cancelCalled = true - - cancel() - } - - // Send enough repetitive content to trigger detection. - - // The pattern "abcdefghij" repeated many times will have very low n-gram uniqueness. - - repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk - - go func() { - // Send 6 chunks of repetitive content = 3000 chars total, - - // each with 500 runes. The check triggers after every 1000 runes - - // when content > 2000 chars. - - for i := 0; i < 6; i++ { - ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} - } - - // Send more data that should be ignored after detection. - - for i := 0; i < 10; i++ { - ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} - } - - close(ch) - }() - - resp, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, nil) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if !detected { - t.Fatal("expected repetition detection to trigger") - } - - if !cancelCalled { - t.Error("expected cancelFn to be called") - } - - // The response should be shorter than the full 3000+ chars - - // because detection triggers early. - - if len(resp.Content) >= 3000+10*len("more data") { - t.Errorf("Content length = %d, expected less than full output", len(resp.Content)) - } - - _ = ctx -} - -func TestConsumeStream_ToolCallAccumulation(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 8) - - go func() { - ch <- protocoltypes.StreamEvent{ - ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ - {Index: 0, ID: "call_1", Name: "test_fn", ArgumentsDelta: `{"ke`}, - }, - } - - ch <- protocoltypes.StreamEvent{ - ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ - {Index: 0, ArgumentsDelta: `y":"val"}`}, - }, - } - - ch <- protocoltypes.StreamEvent{FinishReason: "tool_calls"} - - close(ch) - }() - - _, cancel := context.WithCancel(context.Background()) - - defer cancel() - - resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if detected { - t.Fatal("expected no repetition detection for tool calls") - } - - if len(resp.ToolCalls) != 1 { - t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls)) - } - - if resp.ToolCalls[0].Name != "test_fn" { - t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_fn") - } - - if resp.ToolCalls[0].Arguments["key"] != "val" { - t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "val") - } -} - -func TestConsumeStream_StreamError(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 4) - - go func() { - ch <- protocoltypes.StreamEvent{ContentDelta: "partial"} - - ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("read error")} - - close(ch) - }() - - _, cancel := context.WithCancel(context.Background()) - - defer cancel() - - _, _, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) - - if err == nil { - t.Fatal("expected error, got nil") - } - - if !strings.Contains(err.Error(), "read error") { - t.Errorf("error = %q, want to contain %q", err.Error(), "read error") - } -} - -func TestConsumeStream_OnChunkCallback(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 8) - - go func() { - ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} - - ch <- protocoltypes.StreamEvent{ContentDelta: "world"} - - ch <- protocoltypes.StreamEvent{ContentDelta: "!"} - - ch <- protocoltypes.StreamEvent{FinishReason: "stop"} - - close(ch) - }() - - _, cancel := context.WithCancel(context.Background()) - - defer cancel() - - var chunks []string - - onChunk := func(accumulated, _ string) { - chunks = append(chunks, accumulated) - } - - resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, onChunk) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if detected { - t.Fatal("expected detected=false") - } - - if resp.Content != "Hello world!" { - t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") - } - - // onChunk should be called once per content delta (3 times) - - if len(chunks) != 3 { - t.Fatalf("onChunk called %d times, want 3", len(chunks)) - } - - if chunks[0] != "Hello " { - t.Errorf("chunks[0] = %q, want %q", chunks[0], "Hello ") - } - - if chunks[1] != "Hello world" { - t.Errorf("chunks[1] = %q, want %q", chunks[1], "Hello world") - } - - if chunks[2] != "Hello world!" { - t.Errorf("chunks[2] = %q, want %q", chunks[2], "Hello world!") - } -} - -func TestConsumeStream_OnChunkWithRepetitionDetection(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 64) - - cancelCalled := false - - ctx, cancel := context.WithCancel(context.Background()) - - wrappedCancel := func() { - cancelCalled = true - - cancel() - } - - repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk - - go func() { - for i := 0; i < 6; i++ { - ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} - } - - for i := 0; i < 10; i++ { - ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} - } - - close(ch) - }() - - var chunkCount int - - onChunk := func(_, _ string) { - chunkCount++ - } - - _, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, onChunk) - if err != nil { - t.Fatalf("unexpected error: %v", err) - } - - if !detected { - t.Fatal("expected repetition detection to trigger") - } - - if !cancelCalled { - t.Error("expected cancelFn to be called") - } - - // onChunk should have been called at least once before detection - - if chunkCount == 0 { - t.Error("expected onChunk to be called at least once") - } - - _ = ctx -} - -// modelCapturingMockProvider records which model was passed to Chat. - -type modelCapturingMockProvider struct { - mu sync.Mutex - - models []string - - response string -} - -func (m *modelCapturingMockProvider) Chat( - ctx context.Context, - - messages []providers.Message, - - tools_ []providers.ToolDefinition, - - model string, - - opts map[string]any, -) (*providers.LLMResponse, error) { - m.mu.Lock() - - m.models = append(m.models, model) - - m.mu.Unlock() - - return &providers.LLMResponse{ - Content: m.response, - - ToolCalls: []providers.ToolCall{}, - }, nil -} - -func (m *modelCapturingMockProvider) GetDefaultModel() string { - return "mock-capture-model" -} - -func TestAgentLoop_PlanModel_UsedDuringInterviewing(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-planmodel-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "normal-model", - - PlanModel: "plan-model", - - MaxTokens: 4096, - - MaxToolIterations: 2, - }, - }, - } - - msgBus := bus.NewMessageBus() - - provider := &modelCapturingMockProvider{response: "Plan interview response"} - - al := NewAgentLoop(cfg, msgBus, provider) - - defaultAgent := al.registry.GetDefaultAgent() - - if defaultAgent == nil { - t.Fatal("No default agent found") - } - - // Write MEMORY.md with interviewing status to activate plan model - - memoryDir := filepath.Join(tmpDir, "memory") - - os.MkdirAll(memoryDir, 0o755) - - memoryPath := filepath.Join(memoryDir, "MEMORY.md") - - memoryContent := "# Active Plan\n\n> Task: Test plan model\n> Status: interviewing\n> Phase: 1\n" - - if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { - t.Fatalf("Failed to write MEMORY.md: %v", wErr) - } - - _, err = al.ProcessDirectWithChannel( - - context.Background(), - - "Hello, plan model test", - - "test-plan-session", - - "test", - - "test-chat", - ) - if err != nil { - t.Fatalf("ProcessDirectWithChannel failed: %v", err) - } - - provider.mu.Lock() - - defer provider.mu.Unlock() - - if len(provider.models) == 0 { - t.Fatal("Expected at least one Chat call") - } - - // The first call should use the plan model since we're in interviewing state - - if provider.models[0] != "plan-model" { - t.Errorf("Expected plan model 'plan-model' during interviewing, got %q", provider.models[0]) - } -} - -func TestAgentLoop_PlanModel_NotUsedDuringExecuting(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-planmodel-exec-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "normal-model", - - PlanModel: "plan-model", - - MaxTokens: 4096, - - MaxToolIterations: 2, - }, - }, - } - - msgBus := bus.NewMessageBus() - - provider := &modelCapturingMockProvider{response: "Executing response"} - - al := NewAgentLoop(cfg, msgBus, provider) - - defaultAgent := al.registry.GetDefaultAgent() - - if defaultAgent == nil { - t.Fatal("No default agent found") - } - - // Write MEMORY.md with executing status - should use normal model - - memoryDir := filepath.Join(tmpDir, "memory") - - os.MkdirAll(memoryDir, 0o755) - - memoryPath := filepath.Join(memoryDir, "MEMORY.md") - - memoryContent := `# Active Plan - - - -> Task: Test plan model - -> Status: executing - -> Phase: 1 - - - -## Phase 1: Build - -- [ ] Run build - -` - - if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { - t.Fatalf("Failed to write MEMORY.md: %v", wErr) - } - - _, err = al.ProcessDirectWithChannel( - - context.Background(), - - "Hello, executing test", - - "test-exec-session", - - "test", - - "test-chat", - ) - if err != nil { - t.Fatalf("ProcessDirectWithChannel failed: %v", err) - } - - provider.mu.Lock() - - defer provider.mu.Unlock() - - if len(provider.models) == 0 { - t.Fatal("Expected at least one Chat call") - } - - // During executing phase, should use normal model, not plan model - - if provider.models[0] != "normal-model" { - t.Errorf("Expected normal model 'normal-model' during executing, got %q", provider.models[0]) - } -} - -func TestAgentLoop_PlanModel_ResolvesProviderForSingleCandidate(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-test-planmodel-resolve-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - - defer os.RemoveAll(tmpDir) - - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "MiniMax-M2.5", - - PlanModel: "openai/gpt-5.2", - - MaxTokens: 4096, - - MaxToolIterations: 2, - }, - }, - } - - msgBus := bus.NewMessageBus() - - // The main provider simulates the wrong provider (e.g. MiniMax). - - mainProvider := &modelCapturingMockProvider{response: "wrong provider response"} - - al := NewAgentLoop(cfg, msgBus, mainProvider) - - // Inject a mock provider into the cache so resolveProvider returns it - - // for the "openai/gpt-5.2" candidate (provider="openai", model="gpt-5.2"). - - resolvedProvider := &modelCapturingMockProvider{response: "correct provider response"} - - al.providerCache["openai/gpt-5.2"] = resolvedProvider - - // Write MEMORY.md with interviewing status to activate plan model - - memoryDir := filepath.Join(tmpDir, "memory") - - os.MkdirAll(memoryDir, 0o755) - - memoryPath := filepath.Join(memoryDir, "MEMORY.md") - - memoryContent := "# Active Plan\n\n> Task: Test provider resolution\n> Status: interviewing\n> Phase: 1\n" - - if wErr := os.WriteFile(memoryPath, []byte(memoryContent), 0o644); wErr != nil { - t.Fatalf("Failed to write MEMORY.md: %v", wErr) - } - - _, err = al.ProcessDirectWithChannel( - - context.Background(), - - "Hello, resolve provider test", - - "test-resolve-session", - - "test", - - "test-chat", - ) - if err != nil { - t.Fatalf("ProcessDirectWithChannel failed: %v", err) - } - - resolvedProvider.mu.Lock() - - defer resolvedProvider.mu.Unlock() - - mainProvider.mu.Lock() - - defer mainProvider.mu.Unlock() - - // The resolved provider should have been called with the stripped model name - - if len(resolvedProvider.models) == 0 { - t.Fatal("Expected resolved provider to receive Chat call, but it got none") - } - - if resolvedProvider.models[0] != "gpt-5.2" { - t.Errorf("Expected resolved provider to receive model 'gpt-5.2', got %q", resolvedProvider.models[0]) - } - - // The main provider should NOT have been called for the LLM request - - if len(mainProvider.models) > 0 { - t.Errorf("Expected main provider to receive no Chat calls during plan model phase, got %d calls with models %v", - - len(mainProvider.models), mainProvider.models) - } -} - -func TestPlanCommand_StartClear(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - // Create a plan in review status with phases - - plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" - - _ = agent.ContextBuilder.WriteMemory(plan) - - // Seed session history so we can verify it gets cleared - - agent.Sessions.AddMessage("test-session", "user", "hello") - - agent.Sessions.AddMessage("test-session", "assistant", "world") - - agent.Sessions.SetSummary("test-session", "some summary") - - // Approve with clear - - response, handled := al.handleCommand(context.Background(), bus.InboundMessage{ - Content: "/plan start clear", - - SessionKey: "test-session", - }) - - if !handled { - t.Fatal("expected /plan start clear to be handled") - } - - if !strings.Contains(response, "clean history") { - t.Errorf("expected 'clean history' in response, got %q", response) - } - - if !al.planStartPending { - t.Error("expected planStartPending to be true") - } - - if !al.planClearHistory { - t.Error("expected planClearHistory to be true") - } - - // Simulate what Run() does when planStartPending is set - - al.planStartPending = false - - clearHistory := al.planClearHistory - - al.planClearHistory = false - - if clearHistory { - agent.Sessions.SetHistory("test-session", nil) - - agent.Sessions.SetSummary("test-session", "") - - _ = agent.Sessions.Save("test-session") - } - - // Verify history and summary are cleared - - history := agent.Sessions.GetHistory("test-session") - - if len(history) != 0 { - t.Errorf("expected empty history after clear, got %d messages", len(history)) - } - - summary := agent.Sessions.GetSummary("test-session") - - if summary != "" { - t.Errorf("expected empty summary after clear, got %q", summary) - } -} - -func TestPlanCommand_StartWithoutClear_PreservesHistory(t *testing.T) { - al, cleanup := newTestAgentLoop(t) - - defer cleanup() - - agent := al.registry.GetDefaultAgent() - - // Create a plan in review status with phases - - plan := "# Active Plan\n\n> Task: Test task\n> Status: review\n> Phase: 1\n\n## Phase 1: Setup\n- [ ] Step one\n\n## Context\n" - - _ = agent.ContextBuilder.WriteMemory(plan) - - // Seed session history - - agent.Sessions.AddMessage("test-session", "user", "hello") - - agent.Sessions.AddMessage("test-session", "assistant", "world") - - agent.Sessions.SetSummary("test-session", "some summary") - - // Approve without clear - - response, _ := al.handleCommand(context.Background(), bus.InboundMessage{ - Content: "/plan start", - - SessionKey: "test-session", - }) - - if strings.Contains(response, "clean history") { - t.Errorf("did not expect 'clean history' in response, got %q", response) - } - - if al.planClearHistory { - t.Error("planClearHistory should be false for /plan start without clear") - } - - // Verify history is preserved - - history := agent.Sessions.GetHistory("test-session") - - if len(history) != 2 { - t.Errorf("expected 2 history messages preserved, got %d", len(history)) - } - - summary := agent.Sessions.GetSummary("test-session") - - if summary != "some summary" { - t.Errorf("expected summary preserved, got %q", summary) - } -} - -func TestFilterInterviewTools(t *testing.T) { - allDefs := []providers.ToolDefinition{ - {Function: protocoltypes.ToolFunctionDefinition{Name: "read_file"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "list_dir"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "web_search"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "web_fetch"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "message"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "edit_file"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "append_file"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "write_file"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "exec"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "logs"}}, - - // These should be filtered out: - - {Function: protocoltypes.ToolFunctionDefinition{Name: "spawn_subagent"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "skills_search"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "skills_install"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "bg_monitor"}}, - - {Function: protocoltypes.ToolFunctionDefinition{Name: "i2c_transfer"}}, - } - - filtered := filterInterviewTools(allDefs) - - // Should keep exactly the 10 allowed tools - - if len(filtered) != 10 { - names := make([]string, len(filtered)) - - for i, d := range filtered { - names[i] = d.Function.Name - } - - t.Errorf("expected 10 allowed tools, got %d: %v", len(filtered), names) - } - - // Verify none of the disallowed tools slipped through - - disallowed := map[string]bool{ - "spawnsubagent": true, "skillssearch": true, - - "skillsinstall": true, "bgmonitor": true, "ictransfer": true, - } - - for _, d := range filtered { - norm := tools.NormalizeToolName(d.Function.Name) - - if disallowed[norm] { - t.Errorf("disallowed tool %q should have been filtered out", d.Function.Name) - } - } -} - -func TestBuildStreamingDisplay_ContentOnly(t *testing.T) { - display := buildStreamingDisplay("hello world", "") - - if !strings.HasSuffix(display, " \u2589") { - t.Error("expected cursor suffix") - } - - if strings.Contains(display, "\U0001f9e0") { - t.Error("should not contain brain emoji when no reasoning") - } - - lines := strings.Count(display, "\n") + 1 - - if lines != streamingDisplayLines+1 { // TailPad lines + cursor on last line - t.Logf("display:\n%s", display) - } -} - -func TestBuildStreamingDisplay_ReasoningOnly(t *testing.T) { - display := buildStreamingDisplay("", "let me think about this") - - if !strings.Contains(display, "\U0001f9e0") { - t.Error("expected brain emoji for reasoning phase") - } - - if !strings.Contains(display, "Thinking...") { - t.Error("expected Thinking... header") - } - - if !strings.HasSuffix(display, " \u2589") { - t.Error("expected cursor suffix") - } -} - -func TestBuildStreamingDisplay_Both(t *testing.T) { - display := buildStreamingDisplay("the answer is 42", "first I considered...") - - if !strings.Contains(display, "\U0001f9e0") { - t.Error("expected brain emoji") - } - - if !strings.Contains(display, "responding") { - t.Error("expected responding header when both present") - } - - if !strings.Contains(display, "the answer is 42") { - t.Error("expected content in display") - } -} - func TestHandleReasoning(t *testing.T) { newLoop := func(t *testing.T) (*AgentLoop, *bus.MessageBus) { t.Helper() - tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - t.Cleanup(func() { _ = os.RemoveAll(tmpDir) }) - cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: tmpDir, - - Model: "test-model", - - MaxTokens: 4096, - + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, MaxToolIterations: 10, }, }, } - msgBus := bus.NewMessageBus() - return NewAgentLoop(cfg, msgBus, &mockProvider{}), msgBus } t.Run("skips when any required field is empty", func(t *testing.T) { al, msgBus := newLoop(t) - al.handleReasoning(context.Background(), "reasoning", "telegram", "") ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) - defer cancel() - if msg, ok := msgBus.SubscribeOutbound(ctx); ok { t.Fatalf("expected no outbound message, got %+v", msg) } @@ -4112,19 +924,14 @@ func TestHandleReasoning(t *testing.T) { t.Run("publishes one message for non telegram", func(t *testing.T) { al, msgBus := newLoop(t) - al.handleReasoning(context.Background(), "hello reasoning", "slack", "channel-1") ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() - msg, ok := msgBus.SubscribeOutbound(ctx) - if !ok { t.Fatal("expected an outbound message") } - if msg.Channel != "slack" || msg.ChatID != "channel-1" || msg.Content != "hello reasoning" { t.Fatalf("unexpected outbound message: %+v", msg) } @@ -4132,17 +939,12 @@ func TestHandleReasoning(t *testing.T) { t.Run("publishes one message for telegram", func(t *testing.T) { al, msgBus := newLoop(t) - reasoning := "hello telegram reasoning" - al.handleReasoning(context.Background(), reasoning, "telegram", "tg-chat") ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() - msg, ok := msgBus.SubscribeOutbound(ctx) - if !ok { t.Fatal("expected outbound message") } @@ -4150,33 +952,23 @@ func TestHandleReasoning(t *testing.T) { if msg.Channel != "telegram" { t.Fatalf("expected telegram channel message, got %+v", msg) } - if msg.ChatID != "tg-chat" { t.Fatalf("expected chatID tg-chat, got %+v", msg) } - if msg.Content != reasoning { t.Fatalf("content mismatch: got %q want %q", msg.Content, reasoning) } }) - t.Run("expired ctx", func(t *testing.T) { al, msgBus := newLoop(t) - reasoning := "hello telegram reasoning" - ctx, cancel := context.WithCancel(context.Background()) - cancel() - al.handleReasoning(ctx, reasoning, "telegram", "tg-chat") ctx, cancel = context.WithTimeout(context.Background(), 200*time.Millisecond) - defer cancel() - msg, ok := msgBus.SubscribeOutbound(ctx) - if ok { t.Fatalf("expected no outbound message, got %+v", msg) } @@ -4186,209 +978,191 @@ func TestHandleReasoning(t *testing.T) { al, msgBus := newLoop(t) // Fill the outbound bus buffer until a publish would block. - // Use a short timeout to detect when the buffer is full, - // rather than hardcoding the buffer size. - for i := 0; ; i++ { fillCtx, fillCancel := context.WithTimeout(context.Background(), 50*time.Millisecond) - err := msgBus.PublishOutbound(fillCtx, bus.OutboundMessage{ Channel: "filler", - - ChatID: "filler", - + ChatID: "filler", Content: fmt.Sprintf("filler-%d", i), }) - fillCancel() - if err != nil { // Buffer is full (timed out trying to send). - break } } // Use a short-deadline parent context to bound the test. - ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) - defer cancel() start := time.Now() - al.handleReasoning(ctx, "should timeout", "slack", "channel-full") - elapsed := time.Since(start) // handleReasoning uses a 5s internal timeout, but the parent ctx - // expires in 500ms. It should return within ~500ms, not 5s. - if elapsed > 2*time.Second { t.Fatalf("handleReasoning blocked too long (%v); expected prompt return", elapsed) } // Drain the bus and verify the reasoning message was NOT published - // (it should have been dropped due to timeout). - drainCtx, drainCancel := context.WithTimeout(context.Background(), 100*time.Millisecond) - defer drainCancel() - foundReasoning := false - for { msg, ok := msgBus.SubscribeOutbound(drainCtx) - if !ok { break } - if msg.Content == "should timeout" { foundReasoning = true } } - if foundReasoning { t.Fatal("expected reasoning message to be dropped when bus is full, but it was published") } }) } -func TestFormatDurationMs(t *testing.T) { - tests := []struct { - ms int64 +func TestResolveMediaRefs_ResolvesToBase64(t *testing.T) { + store := media.NewFileMediaStore() + dir := t.TempDir() - want string - }{ - {0, "0ms"}, - - {500, "500ms"}, - - {999, "999ms"}, - - {1000, "1.0s"}, - - {1200, "1.2s"}, - - {3500, "3.5s"}, - - {59900, "59.9s"}, - - {60000, "1m"}, - - {61000, "1m1s"}, - - {65000, "1m5s"}, - - {120000, "2m"}, - - {3661000, "61m1s"}, + // Create a minimal valid PNG (8-byte header is enough for filetype detection) + pngPath := filepath.Join(dir, "test.png") + // PNG magic: 0x89 P N G \r \n 0x1A \n + minimal IHDR + pngHeader := []byte{ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature + 0x00, 0x00, 0x00, 0x0D, // IHDR length + 0x49, 0x48, 0x44, 0x52, // "IHDR" + 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, // 1x1 RGB + 0x00, 0x00, 0x00, // no interlace + 0x90, 0x77, 0x53, 0xDE, // CRC + } + if err := os.WriteFile(pngPath, pngHeader, 0o644); err != nil { + t.Fatal(err) + } + ref, err := store.Store(pngPath, media.MediaMeta{}, "test") + if err != nil { + t.Fatal(err) } - for _, tt := range tests { - t.Run(fmt.Sprintf("%dms", tt.ms), func(t *testing.T) { - got := formatDurationMs(tt.ms) + messages := []providers.Message{ + {Role: "user", Content: "describe this", Media: []string{ref}}, + } + result := resolveMediaRefs(messages, store, config.DefaultMaxMediaSize) - if got != tt.want { - t.Errorf("formatDurationMs(%d) = %q, want %q", tt.ms, got, tt.want) - } - }) + if len(result[0].Media) != 1 { + t.Fatalf("expected 1 resolved media, got %d", len(result[0].Media)) + } + if !strings.HasPrefix(result[0].Media[0], "data:image/png;base64,") { + t.Fatalf("expected data:image/png;base64, prefix, got %q", result[0].Media[0][:40]) } } -func TestFormatSubagentCompletion(t *testing.T) { - tests := []struct { - name string +func TestResolveMediaRefs_SkipsOversizedFile(t *testing.T) { + store := media.NewFileMediaStore() + dir := t.TempDir() - label string - - metadata map[string]string - - want string - }{ - { - "no metadata", - - "scout-1", - - nil, - - "📋 scout-1 completed.", - }, - - { - "empty metadata", - - "scout-1", - - map[string]string{}, - - "📋 scout-1 completed.", - }, - - { - "duration and tool calls", - - "scout-1", - - map[string]string{"duration_ms": "3200", "tool_calls": "5"}, - - "📋 scout-1 completed (3.2s, 5 tool calls).", - }, - - { - "single tool call", - - "coder-1", - - map[string]string{"duration_ms": "1200", "tool_calls": "1"}, - - "📋 coder-1 completed (1.2s, 1 tool call).", - }, - - { - "duration only", - - "scout-2", - - map[string]string{"duration_ms": "65000", "tool_calls": "0"}, - - "📋 scout-2 completed (1m5s).", - }, - - { - "tool calls only", - - "scout-3", - - map[string]string{"duration_ms": "0", "tool_calls": "10"}, - - "📋 scout-3 completed (10 tool calls).", - }, - - { - "zero everything", - - "scout-4", - - map[string]string{"duration_ms": "0", "tool_calls": "0"}, - - "📋 scout-4 completed.", - }, + bigPath := filepath.Join(dir, "big.png") + // Write PNG header + padding to exceed limit + data := make([]byte, 1024+1) // 1KB + 1 byte + copy(data, []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}) + if err := os.WriteFile(bigPath, data, 0o644); err != nil { + t.Fatal(err) } + ref, _ := store.Store(bigPath, media.MediaMeta{}, "test") - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := formatSubagentCompletion(tt.label, tt.metadata) + messages := []providers.Message{ + {Role: "user", Content: "hi", Media: []string{ref}}, + } + // Use a tiny limit (1KB) so the file is oversized + result := resolveMediaRefs(messages, store, 1024) - if got != tt.want { - t.Errorf("formatSubagentCompletion(%q, %v) = %q, want %q", tt.label, tt.metadata, got, tt.want) - } - }) + if len(result[0].Media) != 0 { + t.Fatalf("expected 0 media (oversized), got %d", len(result[0].Media)) + } +} + +func TestResolveMediaRefs_SkipsUnknownType(t *testing.T) { + store := media.NewFileMediaStore() + dir := t.TempDir() + + txtPath := filepath.Join(dir, "readme.txt") + if err := os.WriteFile(txtPath, []byte("hello world"), 0o644); err != nil { + t.Fatal(err) + } + ref, _ := store.Store(txtPath, media.MediaMeta{}, "test") + + messages := []providers.Message{ + {Role: "user", Content: "hi", Media: []string{ref}}, + } + result := resolveMediaRefs(messages, store, config.DefaultMaxMediaSize) + + if len(result[0].Media) != 0 { + t.Fatalf("expected 0 media (unknown type), got %d", len(result[0].Media)) + } +} + +func TestResolveMediaRefs_PassesThroughNonMediaRefs(t *testing.T) { + messages := []providers.Message{ + {Role: "user", Content: "hi", Media: []string{"https://example.com/img.png"}}, + } + result := resolveMediaRefs(messages, nil, config.DefaultMaxMediaSize) + + if len(result[0].Media) != 1 || result[0].Media[0] != "https://example.com/img.png" { + t.Fatalf("expected passthrough of non-media:// URL, got %v", result[0].Media) + } +} + +func TestResolveMediaRefs_DoesNotMutateOriginal(t *testing.T) { + store := media.NewFileMediaStore() + dir := t.TempDir() + pngPath := filepath.Join(dir, "test.png") + pngHeader := []byte{ + 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, + 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, + 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, + 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xDE, + } + os.WriteFile(pngPath, pngHeader, 0o644) + ref, _ := store.Store(pngPath, media.MediaMeta{}, "test") + + original := []providers.Message{ + {Role: "user", Content: "hi", Media: []string{ref}}, + } + originalRef := original[0].Media[0] + + resolveMediaRefs(original, store, config.DefaultMaxMediaSize) + + if original[0].Media[0] != originalRef { + t.Fatal("resolveMediaRefs mutated original message slice") + } +} + +func TestResolveMediaRefs_UsesMetaContentType(t *testing.T) { + store := media.NewFileMediaStore() + dir := t.TempDir() + + // File with JPEG content but stored with explicit content type + jpegPath := filepath.Join(dir, "photo") + jpegHeader := []byte{0xFF, 0xD8, 0xFF, 0xE0} // JPEG magic bytes + os.WriteFile(jpegPath, jpegHeader, 0o644) + ref, _ := store.Store(jpegPath, media.MediaMeta{ContentType: "image/jpeg"}, "test") + + messages := []providers.Message{ + {Role: "user", Content: "hi", Media: []string{ref}}, + } + result := resolveMediaRefs(messages, store, config.DefaultMaxMediaSize) + + if len(result[0].Media) != 1 { + t.Fatalf("expected 1 media, got %d", len(result[0].Media)) + } + if !strings.HasPrefix(result[0].Media[0], "data:image/jpeg;base64,") { + t.Fatalf("expected jpeg prefix, got %q", result[0].Media[0][:30]) } } diff --git a/pkg/agent/mock_provider_test.go b/pkg/agent/mock_provider_test.go index f4042fd01..4962810dc 100644 --- a/pkg/agent/mock_provider_test.go +++ b/pkg/agent/mock_provider_test.go @@ -10,18 +10,13 @@ type mockProvider struct{} func (m *mockProvider) Chat( ctx context.Context, - messages []providers.Message, - tools []providers.ToolDefinition, - model string, - opts map[string]any, ) (*providers.LLMResponse, error) { return &providers.LLMResponse{ - Content: "Mock response", - + Content: "Mock response", ToolCalls: []providers.ToolCall{}, }, nil } diff --git a/pkg/agent/registry_test.go b/pkg/agent/registry_test.go index 5a53f92e6..518bb441f 100644 --- a/pkg/agent/registry_test.go +++ b/pkg/agent/registry_test.go @@ -12,13 +12,9 @@ type mockRegistryProvider struct{} func (m *mockRegistryProvider) Chat( ctx context.Context, - messages []providers.Message, - tools []providers.ToolDefinition, - model string, - options map[string]any, ) (*providers.LLMResponse, error) { return &providers.LLMResponse{Content: "mock", FinishReason: "stop"}, nil @@ -32,15 +28,11 @@ func testCfg(agents []config.AgentConfig) *config.Config { return &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ - Workspace: "/tmp/picoclaw-test-registry", - - Model: "gpt-4", - - MaxTokens: 8192, - + Workspace: "/tmp/picoclaw-test-registry", + Model: "gpt-4", + MaxTokens: 8192, MaxToolIterations: 10, }, - List: agents, }, } @@ -48,21 +40,17 @@ func testCfg(agents []config.AgentConfig) *config.Config { func TestNewAgentRegistry_ImplicitMain(t *testing.T) { cfg := testCfg(nil) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) ids := registry.ListAgentIDs() - if len(ids) != 1 || ids[0] != "main" { t.Errorf("expected implicit main agent, got %v", ids) } agent, ok := registry.GetAgent("main") - if !ok || agent == nil { t.Fatal("expected to find 'main' agent") } - if agent.ID != "main" { t.Errorf("agent.ID = %q, want 'main'", agent.ID) } @@ -71,30 +59,24 @@ func TestNewAgentRegistry_ImplicitMain(t *testing.T) { func TestNewAgentRegistry_ExplicitAgents(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "sales", Default: true, Name: "Sales Bot"}, - {ID: "support", Name: "Support Bot"}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) ids := registry.ListAgentIDs() - if len(ids) != 2 { t.Fatalf("expected 2 agents, got %d: %v", len(ids), ids) } sales, ok := registry.GetAgent("sales") - if !ok || sales == nil { t.Fatal("expected to find 'sales' agent") } - if sales.Name != "Sales Bot" { t.Errorf("sales.Name = %q, want 'Sales Bot'", sales.Name) } support, ok := registry.GetAgent("support") - if !ok || support == nil { t.Fatal("expected to find 'support' agent") } @@ -104,15 +86,12 @@ func TestAgentRegistry_GetAgent_Normalize(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "my-agent", Default: true}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, ok := registry.GetAgent("My-Agent") - if !ok || agent == nil { t.Fatal("expected to find agent with normalized ID") } - if agent.ID != "my-agent" { t.Errorf("agent.ID = %q, want 'my-agent'", agent.ID) } @@ -121,16 +100,12 @@ func TestAgentRegistry_GetAgent_Normalize(t *testing.T) { func TestAgentRegistry_GetDefaultAgent(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "alpha"}, - {ID: "beta", Default: true}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) // GetDefaultAgent first checks for "main", then returns any - agent := registry.GetDefaultAgent() - if agent == nil { t.Fatal("expected a default agent") } @@ -139,36 +114,27 @@ func TestAgentRegistry_GetDefaultAgent(t *testing.T) { func TestAgentRegistry_CanSpawnSubagent(t *testing.T) { cfg := testCfg([]config.AgentConfig{ { - ID: "parent", - + ID: "parent", Default: true, - Subagents: &config.SubagentsConfig{ AllowAgents: []string{"child1", "child2"}, }, }, - {ID: "child1"}, - {ID: "child2"}, - {ID: "restricted"}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) if !registry.CanSpawnSubagent("parent", "child1") { t.Error("expected parent to be allowed to spawn child1") } - if !registry.CanSpawnSubagent("parent", "child2") { t.Error("expected parent to be allowed to spawn child2") } - if registry.CanSpawnSubagent("parent", "restricted") { t.Error("expected parent to NOT be allowed to spawn restricted") } - if registry.CanSpawnSubagent("child1", "child2") { t.Error("expected child1 to NOT be allowed to spawn (no subagents config)") } @@ -177,24 +143,19 @@ func TestAgentRegistry_CanSpawnSubagent(t *testing.T) { func TestAgentRegistry_CanSpawnSubagent_Wildcard(t *testing.T) { cfg := testCfg([]config.AgentConfig{ { - ID: "admin", - + ID: "admin", Default: true, - Subagents: &config.SubagentsConfig{ AllowAgents: []string{"*"}, }, }, - {ID: "any-agent"}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) if !registry.CanSpawnSubagent("admin", "any-agent") { t.Error("expected wildcard to allow spawning any agent") } - if !registry.CanSpawnSubagent("admin", "nonexistent") { t.Error("expected wildcard to allow spawning even nonexistent agents") } @@ -202,15 +163,12 @@ func TestAgentRegistry_CanSpawnSubagent_Wildcard(t *testing.T) { func TestAgentInstance_Model(t *testing.T) { model := &config.AgentModelConfig{Primary: "claude-opus"} - cfg := testCfg([]config.AgentConfig{ {ID: "custom", Default: true, Model: model}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, _ := registry.GetAgent("custom") - if agent.Model != "claude-opus" { t.Errorf("agent.Model = %q, want 'claude-opus'", agent.Model) } @@ -220,13 +178,10 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "inherit", Default: true}, }) - cfg.Agents.Defaults.ModelFallbacks = []string{"openai/gpt-4o-mini", "anthropic/haiku"} - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, _ := registry.GetAgent("inherit") - if len(agent.Fallbacks) != 2 { t.Errorf("expected 2 fallbacks inherited from defaults, got %d", len(agent.Fallbacks)) } @@ -234,22 +189,16 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) { func TestAgentInstance_FallbackExplicitEmpty(t *testing.T) { model := &config.AgentModelConfig{ - Primary: "gpt-4", - + Primary: "gpt-4", Fallbacks: []string{}, // explicitly empty = disable - } - cfg := testCfg([]config.AgentConfig{ {ID: "no-fallback", Default: true, Model: model}, }) - cfg.Agents.Defaults.ModelFallbacks = []string{"should-not-inherit"} - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) agent, _ := registry.GetAgent("no-fallback") - if len(agent.Fallbacks) != 0 { t.Errorf("expected 0 fallbacks (explicit empty), got %d: %v", len(agent.Fallbacks), agent.Fallbacks) } diff --git a/pkg/channels/manager_ext_test.go b/pkg/channels/manager_ext_test.go new file mode 100644 index 000000000..cec14c047 --- /dev/null +++ b/pkg/channels/manager_ext_test.go @@ -0,0 +1,881 @@ +package channels + +import ( + "context" + "fmt" + "github.com/sipeed/picoclaw/pkg/bus" + "golang.org/x/time/rate" + "sync/atomic" + "testing" + "time" +) + +// mockEditorWithSendID implements MessageEditor and MessageSenderWithID. +type mockEditorWithSendID struct { + mockChannel + editFn func(ctx context.Context, chatID, messageID, content string) error + sendWithID func(ctx context.Context, chatID, content string) (string, error) +} + +func (m *mockEditorWithSendID) SendWithID(ctx context.Context, chatID, content string) (string, error) { + return m.sendWithID(ctx, chatID, content) +} + +func TestHandleStatusSend_EditsPlaceholder(t *testing.T) { + m := newTestManager() + var editCalled bool + var editedContent string + + ch := &mockEditorWithSendID{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + editFn: func(_ context.Context, _, messageID, content string) error { + editCalled = true + editedContent = content + if messageID != "ph-42" { + t.Fatalf("expected messageID ph-42, got %s", messageID) + } + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + t.Fatal("SendWithID should not be called when placeholder exists") + return "", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + m.RecordPlaceholder("test", "123", "ph-42") + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "status update 1", IsStatus: true} + m.handleStatusSend(context.Background(), "test", w, msg) + + if !editCalled { + t.Fatal("expected EditMessage to be called on placeholder") + } + if editedContent != "status update 1" { + t.Fatalf("expected content 'status update 1', got %s", editedContent) + } +} + +func TestHandleStatusSend_EditsTrackedStatus(t *testing.T) { + m := newTestManager() + var editCalled bool + + ch := &mockEditorWithSendID{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + editFn: func(_ context.Context, _, messageID, _ string) error { + editCalled = true + if messageID != "status-99" { + t.Fatalf("expected messageID status-99, got %s", messageID) + } + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + t.Fatal("SendWithID should not be called when statusMsgID exists") + return "", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + m.statusMsgIDs.Store("test:123", statusMsgEntry{messageID: "status-99", createdAt: time.Now()}) + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "update 2", IsStatus: true} + m.handleStatusSend(context.Background(), "test", w, msg) + + if !editCalled { + t.Fatal("expected EditMessage to be called on tracked status message") + } +} + +func TestHandleStatusSend_SendsNewAndTracks(t *testing.T) { + m := newTestManager() + var sendWithIDCalled bool + + ch := &mockEditorWithSendID{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + editFn: func(_ context.Context, _, _, _ string) error { + return nil + }, + sendWithID: func(_ context.Context, chatID, content string) (string, error) { + sendWithIDCalled = true + if chatID != "123" { + t.Fatalf("expected chatID 123, got %s", chatID) + } + return "new-msg-1", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "first status", IsStatus: true} + m.handleStatusSend(context.Background(), "test", w, msg) + + if !sendWithIDCalled { + t.Fatal("expected SendWithID to be called") + } + + v, ok := m.statusMsgIDs.Load("test:123") + if !ok { + t.Fatal("expected statusMsgIDs to contain tracked entry") + } + entry := v.(statusMsgEntry) + if entry.messageID != "new-msg-1" { + t.Fatalf("expected messageID new-msg-1, got %s", entry.messageID) + } +} + +func TestHandleTaskStatusSend_EditsExisting(t *testing.T) { + m := newTestManager() + var editCalled bool + + ch := &mockEditorWithSendID{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + editFn: func(_ context.Context, _, messageID, content string) error { + editCalled = true + if messageID != "task-msg-1" { + t.Fatalf("expected messageID task-msg-1, got %s", messageID) + } + if content != "task progress 50%" { + t.Fatalf("expected content 'task progress 50%%', got %s", content) + } + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + t.Fatal("SendWithID should not be called when task message exists") + return "", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + m.taskMsgIDs.Store( + taskStatusKey("test", "123", "task-abc"), + statusMsgEntry{messageID: "task-msg-1", createdAt: time.Now()}, + ) + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "task progress 50%", + IsTaskStatus: true, + TaskID: "task-abc", + } + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + if !editCalled { + t.Fatal("expected EditMessage to be called") + } +} + +func TestHandleTaskStatusSend_SendsNewAndTracks(t *testing.T) { + m := newTestManager() + var sendWithIDCalled bool + + ch := &mockEditorWithSendID{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + editFn: func(_ context.Context, _, _, _ string) error { return nil }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + sendWithIDCalled = true + return "new-task-msg", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "task started", + IsTaskStatus: true, + TaskID: "task-xyz", + } + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + if !sendWithIDCalled { + t.Fatal("expected SendWithID to be called") + } + + v, ok := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-xyz")) + if !ok { + t.Fatal("expected taskMsgIDs to contain tracked entry") + } + entry := v.(statusMsgEntry) + if entry.messageID != "new-task-msg" { + t.Fatalf("expected messageID new-task-msg, got %s", entry.messageID) + } +} + +func TestHandleTaskStatusSend_FallbackToSend(t *testing.T) { + m := newTestManager() + var sendCalled bool + + ch := &mockChannel{ + sendFn: func(_ context.Context, msg bus.OutboundMessage) error { + sendCalled = true + if msg.Content != "task status" { + t.Fatalf("expected content 'task status', got %s", msg.Content) + } + return nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "task status", + IsTaskStatus: true, + TaskID: "task-fallback", + } + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + if !sendCalled { + t.Fatal("expected fallback Send to be called") + } +} + +func TestPreSend_EditsStatusMessage(t *testing.T) { + m := newTestManager() + var editCalled bool + + ch := &mockMessageEditor{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + editFn: func(_ context.Context, _, messageID, _ string) error { + editCalled = true + if messageID != "status-msg-77" { + t.Fatalf("expected messageID status-msg-77, got %s", messageID) + } + return nil + }, + } + + m.statusMsgIDs.Store("test:123", statusMsgEntry{messageID: "status-msg-77", createdAt: time.Now()}) + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final response"} + edited := m.preSend(context.Background(), "test", msg, ch) + + if !edited { + t.Fatal("expected preSend to return true (status message edited)") + } + if !editCalled { + t.Fatal("expected EditMessage to be called") + } + + if _, loaded := m.statusMsgIDs.Load("test:123"); loaded { + t.Fatal("expected statusMsgIDs entry to be deleted after preSend") + } +} + +func TestRunWorker_RoutesStatusMessages(t *testing.T) { + m := newTestManager() + + var regularSendCount atomic.Int32 + var sendWithIDCount atomic.Int32 + + ch := &mockEditorWithSendID{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { + regularSendCount.Add(1) + return nil + }, + }, + editFn: func(_ context.Context, _, _, _ string) error { + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + sendWithIDCount.Add(1) + return "tracked-1", nil + }, + } + + w := &channelWorker{ + ch: ch, + queue: make(chan bus.OutboundMessage, 10), + done: make(chan struct{}), + limiter: rate.NewLimiter(rate.Inf, 1), + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + go m.runWorker(ctx, "test", w) + + w.queue <- bus.OutboundMessage{Channel: "test", ChatID: "1", Content: "status", IsStatus: true} + + w.queue <- bus.OutboundMessage{Channel: "test", ChatID: "2", Content: "task", IsTaskStatus: true, TaskID: "t1"} + + w.queue <- bus.OutboundMessage{Channel: "test", ChatID: "3", Content: "hello"} + + time.Sleep(200 * time.Millisecond) + + if regularSendCount.Load() != 1 { + t.Fatalf("expected 1 regular Send call, got %d", regularSendCount.Load()) + } + if sendWithIDCount.Load() != 2 { + t.Fatalf("expected 2 SendWithID calls (status + task), got %d", sendWithIDCount.Load()) + } +} + +func TestStatusMsgTTLJanitor(t *testing.T) { + m := newTestManager() + + m.statusMsgIDs.Store("test:old", statusMsgEntry{ + messageID: "old-status", + createdAt: time.Now().Add(-10 * time.Minute), + }) + m.taskMsgIDs.Store("task-old", statusMsgEntry{ + messageID: "old-task", + createdAt: time.Now().Add(-60 * time.Minute), + }) + + m.statusMsgIDs.Store("test:fresh", statusMsgEntry{ + messageID: "fresh-status", + createdAt: time.Now(), + }) + + now := time.Now() + m.statusMsgIDs.Range(func(key, value any) bool { + if entry, ok := value.(statusMsgEntry); ok { + if now.Sub(entry.createdAt) > statusMsgTTL { + m.statusMsgIDs.Delete(key) + } + } + return true + }) + m.taskMsgIDs.Range(func(key, value any) bool { + if entry, ok := value.(statusMsgEntry); ok { + if now.Sub(entry.createdAt) > taskMsgTTL { + m.taskMsgIDs.Delete(key) + } + } + return true + }) + + if _, loaded := m.statusMsgIDs.Load("test:old"); loaded { + t.Fatal("expected old status entry to be evicted") + } + if _, loaded := m.taskMsgIDs.Load("task-old"); loaded { + t.Fatal("expected old task entry to be evicted") + } + if _, loaded := m.statusMsgIDs.Load("test:fresh"); !loaded { + t.Fatal("expected fresh status entry to survive") + } +} + +// mockDraftSender implements DraftSender + MessageSenderWithID + MessageEditor. +type mockDraftSender struct { + mockChannel + draftFn func(ctx context.Context, chatID string, draftID int, content string) error + editFn func(ctx context.Context, chatID, messageID, content string) error + sendWithID func(ctx context.Context, chatID, content string) (string, error) +} + +func (m *mockDraftSender) SendDraft(ctx context.Context, chatID string, draftID int, content string) error { + return m.draftFn(ctx, chatID, draftID, content) +} + +func (m *mockDraftSender) SendWithID(ctx context.Context, chatID, content string) (string, error) { + return m.sendWithID(ctx, chatID, content) +} + +func TestHandleStatusSend_UsesDraftSender(t *testing.T) { + m := newTestManager() + var draftCalled bool + var draftContent string + var draftDID int + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, chatID string, draftID int, content string) error { + draftCalled = true + draftContent = content + draftDID = draftID + return nil + }, + editFn: func(_ context.Context, _, _, _ string) error { + t.Fatal("EditMessage should not be called when draft succeeds") + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + t.Fatal("SendWithID should not be called when draft succeeds") + return "", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "streaming preview", IsStatus: true} + m.handleStatusSend(context.Background(), "test", w, msg) + + if !draftCalled { + t.Fatal("expected SendDraft to be called") + } + if draftContent != "streaming preview" { + t.Fatalf("expected draft content 'streaming preview', got %s", draftContent) + } + if draftDID == 0 { + t.Fatal("expected non-zero draftID") + } + + draftCalled = false + var secondDID int + ch.draftFn = func(_ context.Context, _ string, draftID int, _ string) error { + draftCalled = true + secondDID = draftID + return nil + } + msg.Content = "streaming preview updated" + m.handleStatusSend(context.Background(), "test", w, msg) + + if !draftCalled { + t.Fatal("expected SendDraft to be called again") + } + if secondDID != draftDID { + t.Fatalf("expected same draftID %d, got %d", draftDID, secondDID) + } +} + +func TestHandleStatusSend_DraftFails_FallsToEdit(t *testing.T) { + m := newTestManager() + var editCalled bool + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, _ string, _ int, _ string) error { + return fmt.Errorf("draft not supported in group") + }, + editFn: func(_ context.Context, _, _, _ string) error { + editCalled = true + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + return "msg-1", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "preview", IsStatus: true} + m.handleStatusSend(context.Background(), "test", w, msg) + + if editCalled { + t.Fatal("expected EditMessage NOT to be called (no placeholder)") + } +} + +func TestHandleStatusSend_DraftFailure_DoesNotClobberTrackedMessageID(t *testing.T) { + m := newTestManager() + var sendWithIDCount int + var editCount int + var editedMessageID string + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, _ string, _ int, _ string) error { + return fmt.Errorf("draft unsupported") + }, + editFn: func(_ context.Context, _, messageID, _ string) error { + editCount++ + editedMessageID = messageID + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + sendWithIDCount++ + return "msg-1", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{Channel: "test", ChatID: "group-main", Content: "preview-1", IsStatus: true} + m.handleStatusSend(context.Background(), "test", w, msg) + + msg.Content = "preview-2" + m.handleStatusSend(context.Background(), "test", w, msg) + + if sendWithIDCount != 1 { + t.Fatalf("expected SendWithID to be called once, got %d", sendWithIDCount) + } + if editCount != 1 { + t.Fatalf("expected EditMessage to be called once, got %d", editCount) + } + if editedMessageID != "msg-1" { + t.Fatalf("expected EditMessage target msg-1, got %s", editedMessageID) + } +} + +func TestHandleTaskStatusSend_UsesDraftSender(t *testing.T) { + m := newTestManager() + var draftCalled bool + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, _ string, _ int, _ string) error { + draftCalled = true + return nil + }, + editFn: func(_ context.Context, _, _, _ string) error { + t.Fatal("EditMessage should not be called when draft succeeds") + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + t.Fatal("SendWithID should not be called when draft succeeds") + return "", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "task progress 50%", + IsTaskStatus: true, + TaskID: "task-draft", + } + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + if !draftCalled { + t.Fatal("expected SendDraft to be called for task status") + } +} + +func TestHandleTaskStatusSend_Final_UpdatesDraftInPlace(t *testing.T) { + m := newTestManager() + var draftUpdateCalled bool + var draftUpdateDraftID int + var draftUpdateContent string + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { + t.Fatal("Send should not be called when draft update succeeds") + return nil + }, + }, + draftFn: func(_ context.Context, chatID string, draftID int, content string) error { + draftUpdateCalled = true + draftUpdateDraftID = draftID + draftUpdateContent = content + if chatID != "123" { + t.Fatalf("expected chatID 123, got %s", chatID) + } + return nil + }, + editFn: func(_ context.Context, _, _, _ string) error { return nil }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + t.Fatal("SendWithID should not be called when draft update succeeds") + return "", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + m.taskMsgIDs.Store(taskStatusKey("test", "123", "task-final"), statusMsgEntry{draftID: 42, createdAt: time.Now()}) + m.statusEditTimes.Store(taskStatusKey("test", "123", "task-final"), time.Now()) + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "task completed", + IsTaskStatus: true, + TaskID: "task-final", + Final: true, + } + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + if !draftUpdateCalled { + t.Fatal("expected SendDraft to update draft with final content") + } + if draftUpdateDraftID != 42 { + t.Fatalf("expected draftID 42, got %d", draftUpdateDraftID) + } + if draftUpdateContent != "task completed" { + t.Fatalf("expected draft content 'task completed', got %q", draftUpdateContent) + } + if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-final")); loaded { + t.Fatal("expected taskMsgIDs entry to be deleted for final task status") + } + if _, loaded := m.statusEditTimes.Load(taskStatusKey("test", "123", "task-final")); loaded { + t.Fatal("expected statusEditTimes entry to be deleted for final task status") + } +} + +func TestHandleTaskStatusSend_DraftStreaming_IsolatedByChatThread(t *testing.T) { + m := newTestManager() + + type draftCall struct { + chatID string + draftID int + content string + } + calls := make([]draftCall, 0, 2) + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, chatID string, draftID int, content string) error { + calls = append(calls, draftCall{chatID: chatID, draftID: draftID, content: content}) + return nil + }, + editFn: func(_ context.Context, _, _, _ string) error { return nil }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msgA := bus.OutboundMessage{ + Channel: "test", + ChatID: "-100/10", + Content: "A:10%", + IsTaskStatus: true, + TaskID: "shared-task", + } + msgB := bus.OutboundMessage{ + Channel: "test", + ChatID: "-100/20", + Content: "B:10%", + IsTaskStatus: true, + TaskID: "shared-task", + } + + m.handleTaskStatusSend(context.Background(), "test", w, msgA) + m.handleTaskStatusSend(context.Background(), "test", w, msgB) + + if len(calls) != 2 { + t.Fatalf("expected 2 SendDraft calls, got %d", len(calls)) + } + if calls[0].chatID == calls[1].chatID { + t.Fatalf("expected different chat threads, got %q and %q", calls[0].chatID, calls[1].chatID) + } + if calls[0].draftID == calls[1].draftID { + t.Fatalf("expected distinct draft IDs per thread key, both got %d", calls[0].draftID) + } + + if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/10", "shared-task")); !loaded { + t.Fatal("expected taskMsgIDs entry for thread A") + } + if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/20", "shared-task")); !loaded { + t.Fatal("expected taskMsgIDs entry for thread B") + } +} + +func TestHandleTaskStatusSend_DraftFailure_DoesNotClobberTrackedMessageID(t *testing.T) { + m := newTestManager() + var sendWithIDCount int + var editCount int + var editedMessageID string + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, _ string, _ int, _ string) error { + return fmt.Errorf("draft unsupported") + }, + editFn: func(_ context.Context, _, messageID, _ string) error { + editCount++ + editedMessageID = messageID + return nil + }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { + sendWithIDCount++ + return "task-msg-1", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "group-main", + Content: "task-10%", + IsTaskStatus: true, + TaskID: "task-1", + } + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + msg.Content = "task-20%" + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + if sendWithIDCount != 1 { + t.Fatalf("expected SendWithID to be called once, got %d", sendWithIDCount) + } + if editCount != 1 { + t.Fatalf("expected EditMessage to be called once, got %d", editCount) + } + if editedMessageID != "task-msg-1" { + t.Fatalf("expected EditMessage target task-msg-1, got %s", editedMessageID) + } +} + +func TestPreSend_ClearsDraftState(t *testing.T) { + m := newTestManager() + + ch := &mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + } + + m.statusMsgIDs.Store("test:123", statusMsgEntry{draftID: 42, createdAt: time.Now()}) + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final response"} + edited := m.preSend(context.Background(), "test", msg, ch) + + if edited { + t.Fatal("expected preSend to return false for draft-based status (sendMessage replaces draft)") + } + + if _, loaded := m.statusMsgIDs.Load("test:123"); loaded { + t.Fatal("expected draft status entry to be deleted after preSend") + } +} + +func TestGenerateDraftID_Stable(t *testing.T) { + id1 := generateDraftID("telegram:123") + id2 := generateDraftID("telegram:123") + if id1 != id2 { + t.Fatalf("expected stable draft ID, got %d vs %d", id1, id2) + } + if id1 == 0 { + t.Fatal("expected non-zero draft ID") + } + + id3 := generateDraftID("telegram:456") + if id1 == id3 { + t.Fatalf("expected different draft IDs for different keys, both got %d", id1) + } +} + +// TestPreSend_DismissesDraftBeforeSend verifies that preSend explicitly +// dismisses a draft-based status bubble (via SendDraft with empty text) +// before proceeding to send the permanent message. This prevents ghost +// draft bubbles when a user message arrives between the last draft update +// and the final sendMessage. +// TestPreSend_DismissesDraftBeforeSend verifies that preSend explicitly +// dismisses a draft-based status bubble (via SendDraft with empty text) +// before proceeding to send the permanent message. This prevents ghost +// draft bubbles when a user message arrives between the last draft update +// and the final sendMessage. +func TestPreSend_DismissesDraftBeforeSend(t *testing.T) { + m := newTestManager() + + var dismissCalled bool + var dismissContent string + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, _ string, _ int, content string) error { + dismissCalled = true + dismissContent = content + return nil + }, + editFn: func(_ context.Context, _, _, _ string) error { return nil }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil }, + } + + m.statusMsgIDs.Store("test:123", statusMsgEntry{draftID: 42, createdAt: time.Now()}) + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final response"} + edited := m.preSend(context.Background(), "test", msg, ch) + + if edited { + t.Fatal("expected preSend to return false for draft-based status") + } + if !dismissCalled { + t.Fatal("expected preSend to call SendDraft to dismiss the draft") + } + if dismissContent != "" { + t.Fatalf("expected empty dismiss content, got %q", dismissContent) + } +} + +// TestRecordTypingStop_CleansUpOldEntry verifies that recording a new +// typing stop function calls the previous stop first. +// TestRecordTypingStop_CleansUpOldEntry verifies that recording a new +// typing stop function calls the previous stop first. +func TestRecordTypingStop_CleansUpOldEntry(t *testing.T) { + m := newTestManager() + + var oldStopped atomic.Bool + + m.RecordTypingStop("tg", "42", func() { oldStopped.Store(true) }) + + m.RecordTypingStop("tg", "42", func() {}) + + if !oldStopped.Load() { + t.Fatal("expected old typing stop to be called when new entry is recorded") + } +} + +// TestRecordReactionUndo_CleansUpOldEntry verifies that recording a new +// reaction undo function calls the previous undo first. +// TestRecordReactionUndo_CleansUpOldEntry verifies that recording a new +// reaction undo function calls the previous undo first. +func TestRecordReactionUndo_CleansUpOldEntry(t *testing.T) { + m := newTestManager() + + var oldUndone atomic.Bool + + m.RecordReactionUndo("tg", "42", func() { oldUndone.Store(true) }) + + m.RecordReactionUndo("tg", "42", func() {}) + + if !oldUndone.Load() { + t.Fatal("expected old reaction undo to be called when new entry is recorded") + } +} + +// TestPreSend_DraftDismiss_ClearsEditTimes verifies that dismissing a draft +// in preSend also clears the statusEditTimes entry for that key, preventing +// stale throttle state from affecting the next processing cycle. +// TestPreSend_DraftDismiss_ClearsEditTimes verifies that dismissing a draft +// in preSend also clears the statusEditTimes entry for that key, preventing +// stale throttle state from affecting the next processing cycle. +func TestPreSend_DraftDismiss_ClearsEditTimes(t *testing.T) { + m := newTestManager() + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, _ string, _ int, _ string) error { return nil }, + editFn: func(_ context.Context, _, _, _ string) error { return nil }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil }, + } + + key := "test:123" + m.statusMsgIDs.Store(key, statusMsgEntry{draftID: 42, createdAt: time.Now()}) + m.statusEditTimes.Store(key, time.Now()) + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final"} + m.preSend(context.Background(), "test", msg, ch) + + if _, loaded := m.statusEditTimes.Load(key); loaded { + t.Fatal("expected statusEditTimes to be cleared after draft dismiss") + } +} diff --git a/pkg/channels/manager_test.go b/pkg/channels/manager_test.go index c7efaaea2..e0f55288a 100644 --- a/pkg/channels/manager_test.go +++ b/pkg/channels/manager_test.go @@ -17,16 +17,32 @@ import ( // mockChannel is a test double that delegates Send to a configurable function. type mockChannel struct { BaseChannel - sendFn func(ctx context.Context, msg bus.OutboundMessage) error + sendFn func(ctx context.Context, msg bus.OutboundMessage) error + sentMessages []bus.OutboundMessage + placeholdersSent int + editedMessages int + lastPlaceholderID string } func (m *mockChannel) Send(ctx context.Context, msg bus.OutboundMessage) error { + m.sentMessages = append(m.sentMessages, msg) return m.sendFn(ctx, msg) } func (m *mockChannel) Start(ctx context.Context) error { return nil } func (m *mockChannel) Stop(ctx context.Context) error { return nil } +func (m *mockChannel) SendPlaceholder(ctx context.Context, chatID string) (string, error) { + m.placeholdersSent++ + m.lastPlaceholderID = "mock-ph-123" + return m.lastPlaceholderID, nil +} + +func (m *mockChannel) EditMessage(ctx context.Context, chatID, messageID, content string) error { + m.editedMessages++ + return nil +} + // newTestManager creates a minimal Manager suitable for unit tests. func newTestManager() *Manager { return &Manager{ @@ -600,6 +616,37 @@ func TestRecordTypingStop_ConcurrentSafe(t *testing.T) { wg.Wait() } +func TestRecordTypingStop_ReplacesExistingStop(t *testing.T) { + m := newTestManager() + var oldStopCalls int + var newStopCalls int + + m.RecordTypingStop("test", "123", func() { + oldStopCalls++ + }) + + m.RecordTypingStop("test", "123", func() { + newStopCalls++ + }) + + if oldStopCalls != 1 { + t.Fatalf("expected previous typing stop to be called once when replaced, got %d", oldStopCalls) + } + if newStopCalls != 0 { + t.Fatalf("expected replacement typing stop to stay active until preSend, got %d calls", newStopCalls) + } + + msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "hello"} + m.preSend(context.Background(), "test", msg, &mockChannel{}) + + if newStopCalls != 1 { + t.Fatalf("expected replacement typing stop to be called by preSend, got %d", newStopCalls) + } + if oldStopCalls != 1 { + t.Fatalf("expected previous typing stop to not be called again, got %d", oldStopCalls) + } +} + func TestSendWithRetry_PreSendEditsPlaceholder(t *testing.T) { m := newTestManager() var sendCalled bool @@ -861,893 +908,285 @@ func TestBuildMediaScope_WithMessageID(t *testing.T) { } } -// --- Status / TaskStatus message handling tests --- +func TestManager_PlaceholderConsumedByResponse(t *testing.T) { + mgr := &Manager{ + channels: make(map[string]Channel), + workers: make(map[string]*channelWorker), + placeholders: sync.Map{}, + } -// mockEditorWithSendID implements MessageEditor and MessageSenderWithID. -type mockEditorWithSendID struct { - mockChannel - editFn func(ctx context.Context, chatID, messageID, content string) error - sendWithID func(ctx context.Context, chatID, content string) (string, error) -} - -func (m *mockEditorWithSendID) EditMessage(ctx context.Context, chatID, messageID, content string) error { - return m.editFn(ctx, chatID, messageID, content) -} - -func (m *mockEditorWithSendID) SendWithID(ctx context.Context, chatID, content string) (string, error) { - return m.sendWithID(ctx, chatID, content) -} - -func TestHandleStatusSend_EditsPlaceholder(t *testing.T) { - m := newTestManager() - var editCalled bool - var editedContent string - - ch := &mockEditorWithSendID{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - editFn: func(_ context.Context, _, messageID, content string) error { - editCalled = true - editedContent = content - if messageID != "ph-42" { - t.Fatalf("expected messageID ph-42, got %s", messageID) - } + mockCh := &mockChannel{ + sendFn: func(ctx context.Context, msg bus.OutboundMessage) error { return nil }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - t.Fatal("SendWithID should not be called when placeholder exists") - return "", nil - }, + } + worker := newChannelWorker("mock", mockCh) + mgr.channels["mock"] = mockCh + mgr.workers["mock"] = worker + + ctx := context.Background() + key := "mock:chat-1" + + // Simulate a placeholder recorded by base.go HandleMessage + mgr.RecordPlaceholder("mock", "chat-1", "ph-123") + + if _, ok := mgr.placeholders.Load(key); !ok { + t.Fatal("expected placeholder to be recorded") } - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - // Register a placeholder - m.RecordPlaceholder("test", "123", "ph-42") - - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "status update 1", IsStatus: true} - m.handleStatusSend(context.Background(), "test", w, msg) - - if !editCalled { - t.Fatal("expected EditMessage to be called on placeholder") + // Transcription feedback arrives first — it should consume the placeholder + // and be delivered via EditMessage, not Send. + msgTranscript := bus.OutboundMessage{ + Channel: "mock", + ChatID: "chat-1", + Content: "Transcript: hello", } - if editedContent != "status update 1" { - t.Fatalf("expected content 'status update 1', got %s", editedContent) + mgr.sendWithRetry(ctx, "mock", worker, msgTranscript) + + if mockCh.editedMessages != 1 { + t.Errorf("expected 1 edited message (placeholder consumed by transcript), got %d", mockCh.editedMessages) + } + if len(mockCh.sentMessages) != 0 { + t.Errorf("expected 0 normal messages (transcript used edit), got %d", len(mockCh.sentMessages)) + } + + // Placeholder should be gone now + if _, ok := mgr.placeholders.Load(key); ok { + t.Error("expected placeholder to be removed after being consumed") + } + + // Final LLM response arrives — no placeholder left, so it goes through Send + msgFinal := bus.OutboundMessage{ + Channel: "mock", + ChatID: "chat-1", + Content: "Final Answer", + } + mgr.sendWithRetry(ctx, "mock", worker, msgFinal) + + if len(mockCh.sentMessages) != 1 { + t.Errorf("expected 1 normal message sent, got %d", len(mockCh.sentMessages)) } } -func TestHandleStatusSend_EditsTrackedStatus(t *testing.T) { +func TestSendMessage_Synchronous(t *testing.T) { m := newTestManager() - var editCalled bool - ch := &mockEditorWithSendID{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - editFn: func(_ context.Context, _, messageID, _ string) error { - editCalled = true - if messageID != "status-99" { - t.Fatalf("expected messageID status-99, got %s", messageID) - } - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - t.Fatal("SendWithID should not be called when statusMsgID exists") - return "", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - // Pre-store a tracked status message - m.statusMsgIDs.Store("test:123", statusMsgEntry{messageID: "status-99", createdAt: time.Now()}) - - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "update 2", IsStatus: true} - m.handleStatusSend(context.Background(), "test", w, msg) - - if !editCalled { - t.Fatal("expected EditMessage to be called on tracked status message") - } -} - -func TestHandleStatusSend_SendsNewAndTracks(t *testing.T) { - m := newTestManager() - var sendWithIDCalled bool - - ch := &mockEditorWithSendID{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - editFn: func(_ context.Context, _, _, _ string) error { - return nil - }, - sendWithID: func(_ context.Context, chatID, content string) (string, error) { - sendWithIDCalled = true - if chatID != "123" { - t.Fatalf("expected chatID 123, got %s", chatID) - } - return "new-msg-1", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - // No placeholder, no tracked status -> should use SendWithID - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "first status", IsStatus: true} - m.handleStatusSend(context.Background(), "test", w, msg) - - if !sendWithIDCalled { - t.Fatal("expected SendWithID to be called") - } - - // Verify tracked - v, ok := m.statusMsgIDs.Load("test:123") - if !ok { - t.Fatal("expected statusMsgIDs to contain tracked entry") - } - entry := v.(statusMsgEntry) - if entry.messageID != "new-msg-1" { - t.Fatalf("expected messageID new-msg-1, got %s", entry.messageID) - } -} - -func TestHandleTaskStatusSend_EditsExisting(t *testing.T) { - m := newTestManager() - var editCalled bool - - ch := &mockEditorWithSendID{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - editFn: func(_ context.Context, _, messageID, content string) error { - editCalled = true - if messageID != "task-msg-1" { - t.Fatalf("expected messageID task-msg-1, got %s", messageID) - } - if content != "task progress 50%" { - t.Fatalf("expected content 'task progress 50%%', got %s", content) - } - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - t.Fatal("SendWithID should not be called when task message exists") - return "", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - // Pre-store task message - m.taskMsgIDs.Store( - taskStatusKey("test", "123", "task-abc"), - statusMsgEntry{messageID: "task-msg-1", createdAt: time.Now()}, - ) - - msg := bus.OutboundMessage{ - Channel: "test", - ChatID: "123", - Content: "task progress 50%", - IsTaskStatus: true, - TaskID: "task-abc", - } - m.handleTaskStatusSend(context.Background(), "test", w, msg) - - if !editCalled { - t.Fatal("expected EditMessage to be called") - } -} - -func TestHandleTaskStatusSend_SendsNewAndTracks(t *testing.T) { - m := newTestManager() - var sendWithIDCalled bool - - ch := &mockEditorWithSendID{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - editFn: func(_ context.Context, _, _, _ string) error { return nil }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - sendWithIDCalled = true - return "new-task-msg", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - msg := bus.OutboundMessage{ - Channel: "test", - ChatID: "123", - Content: "task started", - IsTaskStatus: true, - TaskID: "task-xyz", - } - m.handleTaskStatusSend(context.Background(), "test", w, msg) - - if !sendWithIDCalled { - t.Fatal("expected SendWithID to be called") - } - - v, ok := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-xyz")) - if !ok { - t.Fatal("expected taskMsgIDs to contain tracked entry") - } - entry := v.(statusMsgEntry) - if entry.messageID != "new-task-msg" { - t.Fatalf("expected messageID new-task-msg, got %s", entry.messageID) - } -} - -func TestHandleTaskStatusSend_FallbackToSend(t *testing.T) { - m := newTestManager() - var sendCalled bool - - // Channel without SendWithID -- only has Send + var received []bus.OutboundMessage ch := &mockChannel{ sendFn: func(_ context.Context, msg bus.OutboundMessage) error { - sendCalled = true - if msg.Content != "task status" { - t.Fatalf("expected content 'task status', got %s", msg.Content) - } + received = append(received, msg) return nil }, } - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - msg := bus.OutboundMessage{ - Channel: "test", - ChatID: "123", - Content: "task status", - IsTaskStatus: true, - TaskID: "task-fallback", - } - m.handleTaskStatusSend(context.Background(), "test", w, msg) - - if !sendCalled { - t.Fatal("expected fallback Send to be called") - } -} - -func TestPreSend_EditsStatusMessage(t *testing.T) { - m := newTestManager() - var editCalled bool - - ch := &mockMessageEditor{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - editFn: func(_ context.Context, _, messageID, _ string) error { - editCalled = true - if messageID != "status-msg-77" { - t.Fatalf("expected messageID status-msg-77, got %s", messageID) - } - return nil - }, - } - - // Store a tracked status message - m.statusMsgIDs.Store("test:123", statusMsgEntry{messageID: "status-msg-77", createdAt: time.Now()}) - - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final response"} - edited := m.preSend(context.Background(), "test", msg, ch) - - if !edited { - t.Fatal("expected preSend to return true (status message edited)") - } - if !editCalled { - t.Fatal("expected EditMessage to be called") - } - - // Verify status message was consumed (LoadAndDelete) - if _, loaded := m.statusMsgIDs.Load("test:123"); loaded { - t.Fatal("expected statusMsgIDs entry to be deleted after preSend") - } -} - -func TestRunWorker_RoutesStatusMessages(t *testing.T) { - m := newTestManager() - - var regularSendCount atomic.Int32 - var sendWithIDCount atomic.Int32 - - ch := &mockEditorWithSendID{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { - regularSendCount.Add(1) - return nil - }, - }, - editFn: func(_ context.Context, _, _, _ string) error { - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - sendWithIDCount.Add(1) - return "tracked-1", nil - }, - } - w := &channelWorker{ ch: ch, - queue: make(chan bus.OutboundMessage, 10), - done: make(chan struct{}), limiter: rate.NewLimiter(rate.Inf, 1), } - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - go m.runWorker(ctx, "test", w) - - // Send a status message for chatID "1" (routed to handleStatusSend -> SendWithID) - w.queue <- bus.OutboundMessage{Channel: "test", ChatID: "1", Content: "status", IsStatus: true} - // Send a task status message for chatID "2" (routed to handleTaskStatusSend -> SendWithID) - w.queue <- bus.OutboundMessage{Channel: "test", ChatID: "2", Content: "task", IsTaskStatus: true, TaskID: "t1"} - // Send a regular message for chatID "3" (no tracked status -> regular Send) - w.queue <- bus.OutboundMessage{Channel: "test", ChatID: "3", Content: "hello"} - - time.Sleep(200 * time.Millisecond) - - if regularSendCount.Load() != 1 { - t.Fatalf("expected 1 regular Send call, got %d", regularSendCount.Load()) - } - if sendWithIDCount.Load() != 2 { - t.Fatalf("expected 2 SendWithID calls (status + task), got %d", sendWithIDCount.Load()) - } -} - -func TestStatusMsgTTLJanitor(t *testing.T) { - m := newTestManager() - - // Store entries with timestamps in the past - m.statusMsgIDs.Store("test:old", statusMsgEntry{ - messageID: "old-status", - createdAt: time.Now().Add(-10 * time.Minute), - }) - m.taskMsgIDs.Store("task-old", statusMsgEntry{ - messageID: "old-task", - createdAt: time.Now().Add(-60 * time.Minute), - }) - // Store a fresh entry that should survive - m.statusMsgIDs.Store("test:fresh", statusMsgEntry{ - messageID: "fresh-status", - createdAt: time.Now(), - }) - - // Simulate janitor logic - now := time.Now() - m.statusMsgIDs.Range(func(key, value any) bool { - if entry, ok := value.(statusMsgEntry); ok { - if now.Sub(entry.createdAt) > statusMsgTTL { - m.statusMsgIDs.Delete(key) - } - } - return true - }) - m.taskMsgIDs.Range(func(key, value any) bool { - if entry, ok := value.(statusMsgEntry); ok { - if now.Sub(entry.createdAt) > taskMsgTTL { - m.taskMsgIDs.Delete(key) - } - } - return true - }) - - if _, loaded := m.statusMsgIDs.Load("test:old"); loaded { - t.Fatal("expected old status entry to be evicted") - } - if _, loaded := m.taskMsgIDs.Load("task-old"); loaded { - t.Fatal("expected old task entry to be evicted") - } - if _, loaded := m.statusMsgIDs.Load("test:fresh"); !loaded { - t.Fatal("expected fresh status entry to survive") - } -} - -// --- DraftSender tests --- - -// mockDraftSender implements DraftSender + MessageSenderWithID + MessageEditor. -type mockDraftSender struct { - mockChannel - draftFn func(ctx context.Context, chatID string, draftID int, content string) error - editFn func(ctx context.Context, chatID, messageID, content string) error - sendWithID func(ctx context.Context, chatID, content string) (string, error) -} - -func (m *mockDraftSender) SendDraft(ctx context.Context, chatID string, draftID int, content string) error { - return m.draftFn(ctx, chatID, draftID, content) -} - -func (m *mockDraftSender) EditMessage(ctx context.Context, chatID, messageID, content string) error { - return m.editFn(ctx, chatID, messageID, content) -} - -func (m *mockDraftSender) SendWithID(ctx context.Context, chatID, content string) (string, error) { - return m.sendWithID(ctx, chatID, content) -} - -func TestHandleStatusSend_UsesDraftSender(t *testing.T) { - m := newTestManager() - var draftCalled bool - var draftContent string - var draftDID int - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - draftFn: func(_ context.Context, chatID string, draftID int, content string) error { - draftCalled = true - draftContent = content - draftDID = draftID - return nil - }, - editFn: func(_ context.Context, _, _, _ string) error { - t.Fatal("EditMessage should not be called when draft succeeds") - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - t.Fatal("SendWithID should not be called when draft succeeds") - return "", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "streaming preview", IsStatus: true} - m.handleStatusSend(context.Background(), "test", w, msg) - - if !draftCalled { - t.Fatal("expected SendDraft to be called") - } - if draftContent != "streaming preview" { - t.Fatalf("expected draft content 'streaming preview', got %s", draftContent) - } - if draftDID == 0 { - t.Fatal("expected non-zero draftID") - } - - // Second call should reuse the same draftID - draftCalled = false - var secondDID int - ch.draftFn = func(_ context.Context, _ string, draftID int, _ string) error { - draftCalled = true - secondDID = draftID - return nil - } - msg.Content = "streaming preview updated" - m.handleStatusSend(context.Background(), "test", w, msg) - - if !draftCalled { - t.Fatal("expected SendDraft to be called again") - } - if secondDID != draftDID { - t.Fatalf("expected same draftID %d, got %d", draftDID, secondDID) - } -} - -func TestHandleStatusSend_DraftFails_FallsToEdit(t *testing.T) { - m := newTestManager() - var editCalled bool - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - draftFn: func(_ context.Context, _ string, _ int, _ string) error { - return fmt.Errorf("draft not supported in group") - }, - editFn: func(_ context.Context, _, _, _ string) error { - editCalled = true - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - return "msg-1", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - // No existing placeholder/status — draft fails, then SendWithID - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "preview", IsStatus: true} - m.handleStatusSend(context.Background(), "test", w, msg) - - // Draft failed, so it should fall through; no placeholder → no edit → SendWithID - if editCalled { - t.Fatal("expected EditMessage NOT to be called (no placeholder)") - } -} - -func TestHandleStatusSend_DraftFailure_DoesNotClobberTrackedMessageID(t *testing.T) { - m := newTestManager() - var sendWithIDCount int - var editCount int - var editedMessageID string - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - draftFn: func(_ context.Context, _ string, _ int, _ string) error { - return fmt.Errorf("draft unsupported") - }, - editFn: func(_ context.Context, _, messageID, _ string) error { - editCount++ - editedMessageID = messageID - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - sendWithIDCount++ - return "msg-1", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - msg := bus.OutboundMessage{Channel: "test", ChatID: "group-main", Content: "preview-1", IsStatus: true} - m.handleStatusSend(context.Background(), "test", w, msg) - - msg.Content = "preview-2" - m.handleStatusSend(context.Background(), "test", w, msg) - - if sendWithIDCount != 1 { - t.Fatalf("expected SendWithID to be called once, got %d", sendWithIDCount) - } - if editCount != 1 { - t.Fatalf("expected EditMessage to be called once, got %d", editCount) - } - if editedMessageID != "msg-1" { - t.Fatalf("expected EditMessage target msg-1, got %s", editedMessageID) - } -} - -func TestHandleTaskStatusSend_UsesDraftSender(t *testing.T) { - m := newTestManager() - var draftCalled bool - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - draftFn: func(_ context.Context, _ string, _ int, _ string) error { - draftCalled = true - return nil - }, - editFn: func(_ context.Context, _, _, _ string) error { - t.Fatal("EditMessage should not be called when draft succeeds") - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - t.Fatal("SendWithID should not be called when draft succeeds") - return "", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + m.channels["test"] = ch + m.workers["test"] = w msg := bus.OutboundMessage{ - Channel: "test", - ChatID: "123", - Content: "task progress 50%", - IsTaskStatus: true, - TaskID: "task-draft", + Channel: "test", + ChatID: "123", + Content: "hello world", + ReplyToMessageID: "msg-456", } - m.handleTaskStatusSend(context.Background(), "test", w, msg) - if !draftCalled { - t.Fatal("expected SendDraft to be called for task status") + err := m.SendMessage(context.Background(), msg) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + // SendMessage is synchronous — message should already be delivered + if len(received) != 1 { + t.Fatalf("expected 1 message sent, got %d", len(received)) + } + if received[0].ReplyToMessageID != "msg-456" { + t.Fatalf("expected ReplyToMessageID msg-456, got %s", received[0].ReplyToMessageID) + } + if received[0].Content != "hello world" { + t.Fatalf("expected content 'hello world', got %s", received[0].Content) } } -func TestHandleTaskStatusSend_Final_UpdatesDraftInPlace(t *testing.T) { +func TestSendMessage_UnknownChannel(t *testing.T) { m := newTestManager() - var draftUpdateCalled bool - var draftUpdateDraftID int - var draftUpdateContent string - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { - t.Fatal("Send should not be called when draft update succeeds") - return nil - }, - }, - draftFn: func(_ context.Context, chatID string, draftID int, content string) error { - draftUpdateCalled = true - draftUpdateDraftID = draftID - draftUpdateContent = content - if chatID != "123" { - t.Fatalf("expected chatID 123, got %s", chatID) - } - return nil - }, - editFn: func(_ context.Context, _, _, _ string) error { return nil }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - t.Fatal("SendWithID should not be called when draft update succeeds") - return "", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - m.taskMsgIDs.Store(taskStatusKey("test", "123", "task-final"), statusMsgEntry{draftID: 42, createdAt: time.Now()}) - m.statusEditTimes.Store(taskStatusKey("test", "123", "task-final"), time.Now()) msg := bus.OutboundMessage{ - Channel: "test", - ChatID: "123", - Content: "task completed", - IsTaskStatus: true, - TaskID: "task-final", - Final: true, + Channel: "nonexistent", + ChatID: "123", + Content: "hello", } - m.handleTaskStatusSend(context.Background(), "test", w, msg) - if !draftUpdateCalled { - t.Fatal("expected SendDraft to update draft with final content") - } - if draftUpdateDraftID != 42 { - t.Fatalf("expected draftID 42, got %d", draftUpdateDraftID) - } - if draftUpdateContent != "task completed" { - t.Fatalf("expected draft content 'task completed', got %q", draftUpdateContent) - } - if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-final")); loaded { - t.Fatal("expected taskMsgIDs entry to be deleted for final task status") - } - if _, loaded := m.statusEditTimes.Load(taskStatusKey("test", "123", "task-final")); loaded { - t.Fatal("expected statusEditTimes entry to be deleted for final task status") + err := m.SendMessage(context.Background(), msg) + if err == nil { + t.Fatal("expected error for unknown channel") } } -func TestHandleTaskStatusSend_DraftStreaming_IsolatedByChatThread(t *testing.T) { - m := newTestManager() - - type draftCall struct { - chatID string - draftID int - content string - } - calls := make([]draftCall, 0, 2) - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - draftFn: func(_ context.Context, chatID string, draftID int, content string) error { - calls = append(calls, draftCall{chatID: chatID, draftID: draftID, content: content}) - return nil - }, - editFn: func(_ context.Context, _, _, _ string) error { return nil }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - msgA := bus.OutboundMessage{ - Channel: "test", - ChatID: "-100/10", - Content: "A:10%", - IsTaskStatus: true, - TaskID: "shared-task", - } - msgB := bus.OutboundMessage{ - Channel: "test", - ChatID: "-100/20", - Content: "B:10%", - IsTaskStatus: true, - TaskID: "shared-task", - } - - m.handleTaskStatusSend(context.Background(), "test", w, msgA) - m.handleTaskStatusSend(context.Background(), "test", w, msgB) - - if len(calls) != 2 { - t.Fatalf("expected 2 SendDraft calls, got %d", len(calls)) - } - if calls[0].chatID == calls[1].chatID { - t.Fatalf("expected different chat threads, got %q and %q", calls[0].chatID, calls[1].chatID) - } - if calls[0].draftID == calls[1].draftID { - t.Fatalf("expected distinct draft IDs per thread key, both got %d", calls[0].draftID) - } - - if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/10", "shared-task")); !loaded { - t.Fatal("expected taskMsgIDs entry for thread A") - } - if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/20", "shared-task")); !loaded { - t.Fatal("expected taskMsgIDs entry for thread B") - } -} - -func TestHandleTaskStatusSend_DraftFailure_DoesNotClobberTrackedMessageID(t *testing.T) { - m := newTestManager() - var sendWithIDCount int - var editCount int - var editedMessageID string - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - draftFn: func(_ context.Context, _ string, _ int, _ string) error { - return fmt.Errorf("draft unsupported") - }, - editFn: func(_ context.Context, _, messageID, _ string) error { - editCount++ - editedMessageID = messageID - return nil - }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { - sendWithIDCount++ - return "task-msg-1", nil - }, - } - - w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} - - msg := bus.OutboundMessage{ - Channel: "test", - ChatID: "group-main", - Content: "task-10%", - IsTaskStatus: true, - TaskID: "task-1", - } - m.handleTaskStatusSend(context.Background(), "test", w, msg) - - msg.Content = "task-20%" - m.handleTaskStatusSend(context.Background(), "test", w, msg) - - if sendWithIDCount != 1 { - t.Fatalf("expected SendWithID to be called once, got %d", sendWithIDCount) - } - if editCount != 1 { - t.Fatalf("expected EditMessage to be called once, got %d", editCount) - } - if editedMessageID != "task-msg-1" { - t.Fatalf("expected EditMessage target task-msg-1, got %s", editedMessageID) - } -} - -func TestPreSend_ClearsDraftState(t *testing.T) { +func TestSendMessage_NoWorker(t *testing.T) { m := newTestManager() ch := &mockChannel{ sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, } + m.channels["test"] = ch + // No worker registered - // Store a draft-based status entry (draftID != 0, messageID empty) - m.statusMsgIDs.Store("test:123", statusMsgEntry{draftID: 42, createdAt: time.Now()}) - - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final response"} - edited := m.preSend(context.Background(), "test", msg, ch) - - // Draft-based entries don't trigger edit; the final sendMessage replaces the draft - if edited { - t.Fatal("expected preSend to return false for draft-based status (sendMessage replaces draft)") + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "hello", } - // Verify draft state was consumed - if _, loaded := m.statusMsgIDs.Load("test:123"); loaded { - t.Fatal("expected draft status entry to be deleted after preSend") + err := m.SendMessage(context.Background(), msg) + if err == nil { + t.Fatal("expected error when no worker exists") } } -func TestGenerateDraftID_Stable(t *testing.T) { - id1 := generateDraftID("telegram:123") - id2 := generateDraftID("telegram:123") - if id1 != id2 { - t.Fatalf("expected stable draft ID, got %d vs %d", id1, id2) - } - if id1 == 0 { - t.Fatal("expected non-zero draft ID") - } - - // Different key should produce different ID - id3 := generateDraftID("telegram:456") - if id1 == id3 { - t.Fatalf("expected different draft IDs for different keys, both got %d", id1) - } -} - -// TestPreSend_DismissesDraftBeforeSend verifies that preSend explicitly -// dismisses a draft-based status bubble (via SendDraft with empty text) -// before proceeding to send the permanent message. This prevents ghost -// draft bubbles when a user message arrives between the last draft update -// and the final sendMessage. -func TestPreSend_DismissesDraftBeforeSend(t *testing.T) { +func TestSendMessage_WithRetry(t *testing.T) { m := newTestManager() - var dismissCalled bool - var dismissContent string - - ch := &mockDraftSender{ - mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, - }, - draftFn: func(_ context.Context, _ string, _ int, content string) error { - dismissCalled = true - dismissContent = content + var callCount int + ch := &mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { + callCount++ + if callCount == 1 { + return fmt.Errorf("transient: %w", ErrTemporary) + } return nil }, - editFn: func(_ context.Context, _, _, _ string) error { return nil }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil }, } - // Store a draft-based status entry (simulates active streaming) - m.statusMsgIDs.Store("test:123", statusMsgEntry{draftID: 42, createdAt: time.Now()}) - - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final response"} - edited := m.preSend(context.Background(), "test", msg, ch) - - if edited { - t.Fatal("expected preSend to return false for draft-based status") + w := &channelWorker{ + ch: ch, + limiter: rate.NewLimiter(rate.Inf, 1), } - if !dismissCalled { - t.Fatal("expected preSend to call SendDraft to dismiss the draft") + m.channels["test"] = ch + m.workers["test"] = w + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "retry me", } - if dismissContent != "" { - t.Fatalf("expected empty dismiss content, got %q", dismissContent) + + err := m.SendMessage(context.Background(), msg) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if callCount != 2 { + t.Fatalf("expected 2 Send calls (1 failure + 1 success), got %d", callCount) } } -// TestRecordTypingStop_CleansUpOldEntry verifies that recording a new -// typing stop function calls the previous stop first. -func TestRecordTypingStop_CleansUpOldEntry(t *testing.T) { +func TestSendMessage_WithSplitting(t *testing.T) { m := newTestManager() - var oldStopped atomic.Bool - - m.RecordTypingStop("tg", "42", func() { oldStopped.Store(true) }) - - // Record a new one — old stop should fire - m.RecordTypingStop("tg", "42", func() {}) - - if !oldStopped.Load() { - t.Fatal("expected old typing stop to be called when new entry is recorded") - } -} - -// TestRecordReactionUndo_CleansUpOldEntry verifies that recording a new -// reaction undo function calls the previous undo first. -func TestRecordReactionUndo_CleansUpOldEntry(t *testing.T) { - m := newTestManager() - - var oldUndone atomic.Bool - - m.RecordReactionUndo("tg", "42", func() { oldUndone.Store(true) }) - - // Record a new one — old undo should fire - m.RecordReactionUndo("tg", "42", func() {}) - - if !oldUndone.Load() { - t.Fatal("expected old reaction undo to be called when new entry is recorded") - } -} - -// TestPreSend_DraftDismiss_ClearsEditTimes verifies that dismissing a draft -// in preSend also clears the statusEditTimes entry for that key, preventing -// stale throttle state from affecting the next processing cycle. -func TestPreSend_DraftDismiss_ClearsEditTimes(t *testing.T) { - m := newTestManager() - - ch := &mockDraftSender{ + var received []string + ch := &mockChannelWithLength{ mockChannel: mockChannel{ - sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + sendFn: func(_ context.Context, msg bus.OutboundMessage) error { + received = append(received, msg.Content) + return nil + }, }, - draftFn: func(_ context.Context, _ string, _ int, _ string) error { return nil }, - editFn: func(_ context.Context, _, _, _ string) error { return nil }, - sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil }, + maxLen: 5, } - key := "test:123" - m.statusMsgIDs.Store(key, statusMsgEntry{draftID: 42, createdAt: time.Now()}) - m.statusEditTimes.Store(key, time.Now()) + w := &channelWorker{ + ch: ch, + limiter: rate.NewLimiter(rate.Inf, 1), + } + m.channels["test"] = ch + m.workers["test"] = w - msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final"} - m.preSend(context.Background(), "test", msg, ch) + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "hello world", + } - if _, loaded := m.statusEditTimes.Load(key); loaded { - t.Fatal("expected statusEditTimes to be cleared after draft dismiss") + err := m.SendMessage(context.Background(), msg) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if len(received) < 2 { + t.Fatalf("expected message to be split into at least 2 chunks, got %d", len(received)) + } +} + +func TestSendMessage_PreservesOrdering(t *testing.T) { + m := newTestManager() + + var order []string + ch := &mockChannel{ + sendFn: func(_ context.Context, msg bus.OutboundMessage) error { + order = append(order, msg.Content) + return nil + }, + } + + w := &channelWorker{ + ch: ch, + limiter: rate.NewLimiter(rate.Inf, 1), + } + m.channels["test"] = ch + m.workers["test"] = w + + // Send two messages sequentially — they must arrive in order + _ = m.SendMessage(context.Background(), bus.OutboundMessage{ + Channel: "test", ChatID: "1", Content: "first", + }) + _ = m.SendMessage(context.Background(), bus.OutboundMessage{ + Channel: "test", ChatID: "1", Content: "second", + }) + + if len(order) != 2 { + t.Fatalf("expected 2 messages, got %d", len(order)) + } + if order[0] != "first" || order[1] != "second" { + t.Fatalf("expected [first, second], got %v", order) + } +} + +func TestManager_SendPlaceholder(t *testing.T) { + mgr := &Manager{ + channels: make(map[string]Channel), + workers: make(map[string]*channelWorker), + placeholders: sync.Map{}, + } + + mockCh := &mockChannel{ + sendFn: func(ctx context.Context, msg bus.OutboundMessage) error { + return nil + }, + } + mgr.channels["mock"] = mockCh + + ctx := context.Background() + + // SendPlaceholder should send a placeholder and record it + ok := mgr.SendPlaceholder(ctx, "mock", "chat-1") + if !ok { + t.Fatal("expected SendPlaceholder to succeed") + } + if mockCh.placeholdersSent != 1 { + t.Errorf("expected 1 placeholder sent, got %d", mockCh.placeholdersSent) + } + + key := "mock:chat-1" + if _, loaded := mgr.placeholders.Load(key); !loaded { + t.Error("expected placeholder to be recorded in manager") + } + + // SendPlaceholder on unknown channel should return false + ok = mgr.SendPlaceholder(ctx, "unknown", "chat-1") + if ok { + t.Error("expected SendPlaceholder to fail for unknown channel") } } diff --git a/pkg/channels/matrix/matrix_test.go b/pkg/channels/matrix/matrix_test.go index e76db0d3e..07a35c021 100644 --- a/pkg/channels/matrix/matrix_test.go +++ b/pkg/channels/matrix/matrix_test.go @@ -2,14 +2,19 @@ package matrix import ( "context" + "net/http" + "net/http/httptest" "os" "path/filepath" + "strings" "testing" "time" "maunium.net/go/mautrix" "maunium.net/go/mautrix/event" "maunium.net/go/mautrix/id" + + "github.com/sipeed/picoclaw/pkg/config" ) func TestMatrixLocalpartMentionRegexp(t *testing.T) { @@ -194,6 +199,50 @@ func TestMatrixMediaExt(t *testing.T) { } } +func TestDownloadMedia_WritesResponseToTempFile(t *testing.T) { + const wantBody = "matrix-media-payload" + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/_matrix/client/v1/media/download/matrix.test/abc123") { + t.Fatalf("unexpected download path: %s", r.URL.Path) + } + w.Header().Set("Content-Type", "image/png") + _, _ = w.Write([]byte(wantBody)) + })) + defer server.Close() + + client, err := mautrix.NewClient(server.URL, id.UserID("@picoclaw:matrix.test"), "") + if err != nil { + t.Fatalf("NewClient: %v", err) + } + + ch := &MatrixChannel{client: client} + msg := &event.MessageEventContent{ + MsgType: event.MsgImage, + Body: "image.png", + URL: id.ContentURIString("mxc://matrix.test/abc123"), + Info: &event.FileInfo{MimeType: "image/png"}, + } + + path, err := ch.downloadMedia(context.Background(), msg, "image") + if err != nil { + t.Fatalf("downloadMedia: %v", err) + } + defer os.Remove(path) + + if ext := filepath.Ext(path); ext != ".png" { + t.Fatalf("temp file extension=%q want=.png", ext) + } + + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(got) != wantBody { + t.Fatalf("file contents=%q want=%q", string(got), wantBody) + } +} + func TestExtractInboundContent_ImageNoURLFallback(t *testing.T) { ch := &MatrixChannel{} msg := &event.MessageEventContent{ @@ -289,3 +338,50 @@ func TestMatrixOutboundContent(t *testing.T) { t.Fatalf("unexpected fallback body: %q", noCaption.Body) } } + +func TestMarkdownToHTML(t *testing.T) { + tests := []struct { + name string + input string + contains string + }{ + {"bold", "**hello**", "hello"}, + {"italic", "_world_", "world"}, + {"header", "### Title", ""}, + {"inline code", "`x`", "x"}, + {"plain text", "just text", "just text"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := markdownToHTML(tt.input) + if !strings.Contains(got, tt.contains) { + t.Fatalf("markdownToHTML(%q) = %q, want it to contain %q", tt.input, got, tt.contains) + } + }) + } +} + +func TestMessageContent(t *testing.T) { + richtext := &MatrixChannel{config: config.MatrixConfig{MessageFormat: "richtext"}} + plain := &MatrixChannel{config: config.MatrixConfig{MessageFormat: "plain"}} + defaultt := &MatrixChannel{config: config.MatrixConfig{}} + + for _, c := range []*MatrixChannel{richtext, defaultt} { + mc := c.messageContent("**hi**") + if mc.Format != event.FormatHTML { + t.Errorf("format %q: expected FormatHTML, got %q", c.config.MessageFormat, mc.Format) + } + if !strings.Contains(mc.FormattedBody, "hi") { + t.Errorf("format %q: FormattedBody %q missing ", c.config.MessageFormat, mc.FormattedBody) + } + if mc.Body != "**hi**" { + t.Errorf("format %q: Body should remain plain, got %q", c.config.MessageFormat, mc.Body) + } + } + + mc := plain.messageContent("**hi**") + if mc.Format != "" || mc.FormattedBody != "" { + t.Errorf("plain: expected no formatting, got format=%q formattedBody=%q", mc.Format, mc.FormattedBody) + } +} diff --git a/pkg/channels/telegram/telegram_ext_test.go b/pkg/channels/telegram/telegram_ext_test.go new file mode 100644 index 000000000..c8819a7fa --- /dev/null +++ b/pkg/channels/telegram/telegram_ext_test.go @@ -0,0 +1,52 @@ +package telegram + +import ( + "testing" +) + +func TestParseChatID(t *testing.T) { + tests := []struct { + name string + input string + wantCID int64 + wantTID int + wantErr bool + }{ + {name: "plain private", input: "12345", wantCID: 12345, wantTID: 0}, + {name: "group topic", input: "-100123/45", wantCID: -100123, wantTID: 45}, + {name: "trim spaces", input: " -100200/7 ", wantCID: -100200, wantTID: 7}, + {name: "topic zero", input: "-100/0", wantCID: -100, wantTID: 0}, + {name: "empty", input: "", wantErr: true}, + {name: "bad chat", input: "abc/def", wantErr: true}, + {name: "missing topic", input: "-100/", wantErr: true}, + {name: "too many parts", input: "-100/1/2", wantErr: true}, + {name: "negative topic", input: "-100/-1", wantErr: true}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + gotCID, gotTID, err := parseChatID(tc.input) + if tc.wantErr { + if err == nil { + t.Fatalf("parseChatID(%q) expected error, got nil", tc.input) + } + return + } + if err != nil { + t.Fatalf("parseChatID(%q) unexpected error: %v", tc.input, err) + } + if gotCID != tc.wantCID || gotTID != tc.wantTID { + t.Fatalf("parseChatID(%q) = (%d, %d), want (%d, %d)", tc.input, gotCID, gotTID, tc.wantCID, tc.wantTID) + } + }) + } +} + +func TestFormatChatID(t *testing.T) { + if got := formatChatID(-100, 42); got != "-100/42" { + t.Fatalf("formatChatID(-100, 42) = %q, want %q", got, "-100/42") + } + if got := formatChatID(12345, 0); got != "12345" { + t.Fatalf("formatChatID(12345, 0) = %q, want %q", got, "12345") + } +} diff --git a/pkg/channels/telegram/telegram_test.go b/pkg/channels/telegram/telegram_test.go index 83588f926..c2186d0a3 100644 --- a/pkg/channels/telegram/telegram_test.go +++ b/pkg/channels/telegram/telegram_test.go @@ -1,50 +1,462 @@ package telegram -import "testing" +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" -func TestParseChatID(t *testing.T) { - tests := []struct { - name string - input string - wantCID int64 - wantTID int - wantErr bool - }{ - {name: "plain private", input: "12345", wantCID: 12345, wantTID: 0}, - {name: "group topic", input: "-100123/45", wantCID: -100123, wantTID: 45}, - {name: "trim spaces", input: " -100200/7 ", wantCID: -100200, wantTID: 7}, - {name: "topic zero", input: "-100/0", wantCID: -100, wantTID: 0}, - {name: "empty", input: "", wantErr: true}, - {name: "bad chat", input: "abc/def", wantErr: true}, - {name: "missing topic", input: "-100/", wantErr: true}, - {name: "too many parts", input: "-100/1/2", wantErr: true}, - {name: "negative topic", input: "-100/-1", wantErr: true}, - } + "github.com/mymmrac/telego" + ta "github.com/mymmrac/telego/telegoapi" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - gotCID, gotTID, err := parseChatID(tc.input) - if tc.wantErr { - if err == nil { - t.Fatalf("parseChatID(%q) expected error, got nil", tc.input) - } - return - } - if err != nil { - t.Fatalf("parseChatID(%q) unexpected error: %v", tc.input, err) - } - if gotCID != tc.wantCID || gotTID != tc.wantTID { - t.Fatalf("parseChatID(%q) = (%d, %d), want (%d, %d)", tc.input, gotCID, gotTID, tc.wantCID, tc.wantTID) - } - }) + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/channels" +) + +const testToken = "1234567890:aaaabbbbaaaabbbbaaaabbbbaaaabbbbccc" + +// stubCaller implements ta.Caller for testing. +type stubCaller struct { + calls []stubCall + callFn func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) +} + +type stubCall struct { + URL string + Data *ta.RequestData +} + +func (s *stubCaller) Call(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + s.calls = append(s.calls, stubCall{URL: url, Data: data}) + return s.callFn(ctx, url, data) +} + +// stubConstructor implements ta.RequestConstructor for testing. +type stubConstructor struct{} + +func (s *stubConstructor) JSONRequest(parameters any) (*ta.RequestData, error) { + return &ta.RequestData{}, nil +} + +func (s *stubConstructor) MultipartRequest( + parameters map[string]string, + files map[string]ta.NamedReader, +) (*ta.RequestData, error) { + return &ta.RequestData{}, nil +} + +// successResponse returns a ta.Response that telego will treat as a successful SendMessage. +func successResponse(t *testing.T) *ta.Response { + t.Helper() + msg := &telego.Message{MessageID: 1} + b, err := json.Marshal(msg) + require.NoError(t, err) + return &ta.Response{Ok: true, Result: b} +} + +// newTestChannel creates a TelegramChannel with a mocked bot for unit testing. +func newTestChannel(t *testing.T, caller *stubCaller) *TelegramChannel { + t.Helper() + + bot, err := telego.NewBot(testToken, + telego.WithAPICaller(caller), + telego.WithRequestConstructor(&stubConstructor{}), + telego.WithDiscardLogger(), + ) + require.NoError(t, err) + + base := channels.NewBaseChannel("telegram", nil, nil, nil, + channels.WithMaxMessageLength(4000), + ) + base.SetRunning(true) + + return &TelegramChannel{ + BaseChannel: base, + bot: bot, + chatIDs: make(map[string]int64), } } -func TestFormatChatID(t *testing.T) { - if got := formatChatID(-100, 42); got != "-100/42" { - t.Fatalf("formatChatID(-100, 42) = %q, want %q", got, "-100/42") - } - if got := formatChatID(12345, 0); got != "12345" { - t.Fatalf("formatChatID(12345, 0) = %q, want %q", got, "12345") +func TestSend_EmptyContent(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + t.Fatal("SendMessage should not be called for empty content") + return nil, nil + }, } + ch := newTestChannel(t, caller) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: "", + }) + + assert.NoError(t, err) + assert.Empty(t, caller.calls, "no API calls should be made for empty content") +} + +func TestSend_ShortMessage_SingleCall(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return successResponse(t), nil + }, + } + ch := newTestChannel(t, caller) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: "Hello, world!", + }) + + assert.NoError(t, err) + assert.Len(t, caller.calls, 1, "short message should result in exactly one SendMessage call") +} + +func TestSend_LongMessage_SingleCall(t *testing.T) { + // With WithMaxMessageLength(4000), the Manager pre-splits messages before + // they reach Send(). A message at exactly 4000 chars should go through + // as a single SendMessage call (no re-split needed since HTML expansion + // won't exceed 4096 for plain text). + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return successResponse(t), nil + }, + } + ch := newTestChannel(t, caller) + + longContent := strings.Repeat("a", 4000) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: longContent, + }) + + assert.NoError(t, err) + assert.Len(t, caller.calls, 1, "pre-split message within limit should result in one SendMessage call") +} + +func TestSend_HTMLFallback_PerChunk(t *testing.T) { + callCount := 0 + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + callCount++ + // Fail on odd calls (HTML attempt), succeed on even calls (plain text fallback) + if callCount%2 == 1 { + return nil, errors.New("Bad Request: can't parse entities") + } + return successResponse(t), nil + }, + } + ch := newTestChannel(t, caller) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: "Hello **world**", + }) + + assert.NoError(t, err) + // One short message → 1 HTML attempt (fail) + 1 plain text fallback (success) = 2 calls + assert.Equal(t, 2, len(caller.calls), "should have HTML attempt + plain text fallback") +} + +func TestSend_HTMLFallback_BothFail(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return nil, errors.New("send failed") + }, + } + ch := newTestChannel(t, caller) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: "Hello", + }) + + assert.Error(t, err) + assert.True(t, errors.Is(err, channels.ErrTemporary), "error should wrap ErrTemporary") + assert.Equal(t, 2, len(caller.calls), "should have HTML attempt + plain text attempt") +} + +func TestSend_LongMessage_HTMLFallback_StopsOnError(t *testing.T) { + // With a long message that gets split into 2 chunks, if both HTML and + // plain text fail on the first chunk, Send should return early. + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return nil, errors.New("send failed") + }, + } + ch := newTestChannel(t, caller) + + longContent := strings.Repeat("x", 4001) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: longContent, + }) + + assert.Error(t, err) + // Should fail on the first chunk (2 calls: HTML + fallback), never reaching the second chunk. + assert.Equal(t, 2, len(caller.calls), "should stop after first chunk fails both HTML and plain text") +} + +func TestSend_MarkdownShortButHTMLLong_MultipleCalls(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return successResponse(t), nil + }, + } + ch := newTestChannel(t, caller) + + // Create markdown whose length is <= 4000 but whose HTML expansion is much longer. + // "**a** " (6 chars) becomes "a " (9 chars) in HTML, so repeating it many times + // yields HTML that exceeds Telegram's limit while markdown stays within it. + markdownContent := strings.Repeat("**a** ", 600) // 3600 chars markdown, HTML ~5400+ chars + assert.LessOrEqual(t, len([]rune(markdownContent)), 4000, "markdown content must not exceed chunk size") + + htmlExpanded := markdownToTelegramHTML(markdownContent) + assert.Greater( + t, len([]rune(htmlExpanded)), 4096, + "HTML expansion must exceed Telegram limit for this test to be meaningful", + ) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: markdownContent, + }) + + assert.NoError(t, err) + assert.Greater( + t, len(caller.calls), 1, + "markdown-short but HTML-long message should be split into multiple SendMessage calls", + ) +} + +func TestSend_NotRunning(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + t.Fatal("should not be called") + return nil, nil + }, + } + ch := newTestChannel(t, caller) + ch.SetRunning(false) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "12345", + Content: "Hello", + }) + + assert.ErrorIs(t, err, channels.ErrNotRunning) + assert.Empty(t, caller.calls) +} + +func TestSend_InvalidChatID(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + t.Fatal("should not be called") + return nil, nil + }, + } + ch := newTestChannel(t, caller) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "not-a-number", + Content: "Hello", + }) + + assert.Error(t, err) + assert.True(t, errors.Is(err, channels.ErrSendFailed), "error should wrap ErrSendFailed") + assert.Empty(t, caller.calls) +} + +func TestParseTelegramChatID_Plain(t *testing.T) { + cid, tid, err := parseTelegramChatID("12345") + assert.NoError(t, err) + assert.Equal(t, int64(12345), cid) + assert.Equal(t, 0, tid) +} + +func TestParseTelegramChatID_NegativeGroup(t *testing.T) { + cid, tid, err := parseTelegramChatID("-1001234567890") + assert.NoError(t, err) + assert.Equal(t, int64(-1001234567890), cid) + assert.Equal(t, 0, tid) +} + +func TestParseTelegramChatID_WithThreadID(t *testing.T) { + cid, tid, err := parseTelegramChatID("-1001234567890/42") + assert.NoError(t, err) + assert.Equal(t, int64(-1001234567890), cid) + assert.Equal(t, 42, tid) +} + +func TestParseTelegramChatID_GeneralTopic(t *testing.T) { + cid, tid, err := parseTelegramChatID("-100123/1") + assert.NoError(t, err) + assert.Equal(t, int64(-100123), cid) + assert.Equal(t, 1, tid) +} + +func TestParseTelegramChatID_Invalid(t *testing.T) { + _, _, err := parseTelegramChatID("not-a-number") + assert.Error(t, err) +} + +func TestParseTelegramChatID_InvalidThreadID(t *testing.T) { + _, _, err := parseTelegramChatID("-100123/not-a-thread") + assert.Error(t, err) + assert.Contains(t, err.Error(), "invalid thread ID") +} + +func TestSend_WithForumThreadID(t *testing.T) { + caller := &stubCaller{ + callFn: func(ctx context.Context, url string, data *ta.RequestData) (*ta.Response, error) { + return successResponse(t), nil + }, + } + ch := newTestChannel(t, caller) + + err := ch.Send(context.Background(), bus.OutboundMessage{ + ChatID: "-1001234567890/42", + Content: "Hello from topic", + }) + + assert.NoError(t, err) + assert.Len(t, caller.calls, 1) +} + +func TestHandleMessage_ForumTopic_SetsMetadata(t *testing.T) { + messageBus := bus.NewMessageBus() + ch := &TelegramChannel{ + BaseChannel: channels.NewBaseChannel("telegram", nil, messageBus, nil), + chatIDs: make(map[string]int64), + ctx: context.Background(), + } + + msg := &telego.Message{ + Text: "hello from topic", + MessageID: 10, + MessageThreadID: 42, + Chat: telego.Chat{ + ID: -1001234567890, + Type: "supergroup", + IsForum: true, + }, + From: &telego.User{ + ID: 7, + FirstName: "Alice", + }, + } + + err := ch.handleMessage(context.Background(), msg) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + inbound, ok := messageBus.ConsumeInbound(ctx) + require.True(t, ok, "expected inbound message") + + // Composite chatID should include thread ID + assert.Equal(t, "-1001234567890/42", inbound.ChatID) + + // Peer ID should include thread ID for session key isolation + assert.Equal(t, "group", inbound.Peer.Kind) + assert.Equal(t, "-1001234567890/42", inbound.Peer.ID) + + // Parent peer metadata should be set for agent binding + assert.Equal(t, "topic", inbound.Metadata["parent_peer_kind"]) + assert.Equal(t, "42", inbound.Metadata["parent_peer_id"]) +} + +func TestHandleMessage_NoForum_NoThreadMetadata(t *testing.T) { + messageBus := bus.NewMessageBus() + ch := &TelegramChannel{ + BaseChannel: channels.NewBaseChannel("telegram", nil, messageBus, nil), + chatIDs: make(map[string]int64), + ctx: context.Background(), + } + + msg := &telego.Message{ + Text: "regular group message", + MessageID: 11, + Chat: telego.Chat{ + ID: -100999, + Type: "group", + }, + From: &telego.User{ + ID: 8, + FirstName: "Bob", + }, + } + + err := ch.handleMessage(context.Background(), msg) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + inbound, ok := messageBus.ConsumeInbound(ctx) + require.True(t, ok) + + // Plain chatID without thread suffix + assert.Equal(t, "-100999", inbound.ChatID) + + // Peer ID should be raw chat ID (no thread suffix) + assert.Equal(t, "group", inbound.Peer.Kind) + assert.Equal(t, "-100999", inbound.Peer.ID) + + // No parent peer metadata + assert.Empty(t, inbound.Metadata["parent_peer_kind"]) + assert.Empty(t, inbound.Metadata["parent_peer_id"]) +} + +func TestHandleMessage_ReplyThread_NonForum_NoIsolation(t *testing.T) { + messageBus := bus.NewMessageBus() + ch := &TelegramChannel{ + BaseChannel: channels.NewBaseChannel("telegram", nil, messageBus, nil), + chatIDs: make(map[string]int64), + ctx: context.Background(), + } + + // In regular groups, reply threads set MessageThreadID to the original + // message ID. This should NOT trigger per-thread session isolation. + msg := &telego.Message{ + Text: "reply in thread", + MessageID: 20, + MessageThreadID: 15, + Chat: telego.Chat{ + ID: -100999, + Type: "supergroup", + IsForum: false, + }, + From: &telego.User{ + ID: 9, + FirstName: "Carol", + }, + } + + err := ch.handleMessage(context.Background(), msg) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + + inbound, ok := messageBus.ConsumeInbound(ctx) + require.True(t, ok) + + // chatID should NOT include thread suffix for non-forum groups + assert.Equal(t, "-100999", inbound.ChatID) + + // Peer ID should be raw chat ID (shared session for whole group) + assert.Equal(t, "group", inbound.Peer.Kind) + assert.Equal(t, "-100999", inbound.Peer.ID) + + // No parent peer metadata + assert.Empty(t, inbound.Metadata["parent_peer_kind"]) + assert.Empty(t, inbound.Metadata["parent_peer_id"]) } diff --git a/pkg/channels/wecom/app_test.go b/pkg/channels/wecom/app_test.go index 7f230494f..7d07041ad 100644 --- a/pkg/channels/wecom/app_test.go +++ b/pkg/channels/wecom/app_test.go @@ -209,7 +209,7 @@ func TestWeComAppVerifySignature(t *testing.T) { } }) - t.Run("empty token skips verification", func(t *testing.T) { + t.Run("empty token rejects verification (fail-closed)", func(t *testing.T) { cfgEmpty := config.WeComAppConfig{ CorpID: "test_corp_id", CorpSecret: "test_secret", @@ -218,8 +218,8 @@ func TestWeComAppVerifySignature(t *testing.T) { } chEmpty, _ := NewWeComAppChannel(cfgEmpty, msgBus) - if !verifySignature(chEmpty.config.Token, "any_sig", "any_ts", "any_nonce", "any_msg") { - t.Error("empty token should skip verification and return true") + if verifySignature(chEmpty.config.Token, "any_sig", "any_ts", "any_nonce", "any_msg") { + t.Error("empty token should reject verification (fail-closed)") } }) } diff --git a/pkg/channels/wecom/bot_test.go b/pkg/channels/wecom/bot_test.go index c053578b1..d223bb6b6 100644 --- a/pkg/channels/wecom/bot_test.go +++ b/pkg/channels/wecom/bot_test.go @@ -189,8 +189,7 @@ func TestWeComBotVerifySignature(t *testing.T) { } }) - t.Run("empty token skips verification", func(t *testing.T) { - // Create a channel manually with empty token to test the behavior + t.Run("empty token rejects verification (fail-closed)", func(t *testing.T) { cfgEmpty := config.WeComConfig{ Token: "", WebhookURL: "https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=test", @@ -199,8 +198,8 @@ func TestWeComBotVerifySignature(t *testing.T) { config: cfgEmpty, } - if !verifySignature(chEmpty.config.Token, "any_sig", "any_ts", "any_nonce", "any_msg") { - t.Error("empty token should skip verification and return true") + if verifySignature(chEmpty.config.Token, "any_sig", "any_ts", "any_nonce", "any_msg") { + t.Error("empty token should reject verification (fail-closed)") } }) } diff --git a/pkg/channels/wecom/common.go b/pkg/channels/wecom/common.go index 6510e6f81..9a622a2fc 100644 --- a/pkg/channels/wecom/common.go +++ b/pkg/channels/wecom/common.go @@ -31,7 +31,7 @@ func computeSignature(token, timestamp, nonce, encrypt string) string { // This is a common function used by both WeCom Bot and WeCom App func verifySignature(token, msgSignature, timestamp, nonce, msgEncrypt string) bool { if token == "" { - return true // Skip verification if token is not set + return false } return computeSignature(token, timestamp, nonce, msgEncrypt) == msgSignature } diff --git a/pkg/config/config_ext_test.go b/pkg/config/config_ext_test.go new file mode 100644 index 000000000..351ce3826 --- /dev/null +++ b/pkg/config/config_ext_test.go @@ -0,0 +1,123 @@ +package config + +import ( + "encoding/json" + "testing" +) + +func TestAgentDefaults_PlanModel_StringParse(t *testing.T) { + jsonData := `{ + "agents": { + "defaults": { + "workspace": "~/.picoclaw/workspace", + "model": "glm-4.7", + "plan_model": "anthropic/claude-sonnet-4-6", + "plan_model_fallbacks": ["openai/gpt-4o"], + "max_tokens": 8192, + "max_tool_iterations": 20 + } + } + }` + + cfg := DefaultConfig() + if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { + t.Fatalf("unmarshal: %v", err) + } + + if cfg.Agents.Defaults.PlanModel != "anthropic/claude-sonnet-4-6" { + t.Errorf("PlanModel = %q, want 'anthropic/claude-sonnet-4-6'", cfg.Agents.Defaults.PlanModel) + } + if len(cfg.Agents.Defaults.PlanModelFallbacks) != 1 || + cfg.Agents.Defaults.PlanModelFallbacks[0] != "openai/gpt-4o" { + t.Errorf("PlanModelFallbacks = %v, want [openai/gpt-4o]", cfg.Agents.Defaults.PlanModelFallbacks) + } +} + +func TestAgentConfig_PlanModel_ObjectParse(t *testing.T) { + jsonData := `{ + "agents": { + "defaults": { + "workspace": "~/.picoclaw/workspace", + "model": "glm-4.7", + "max_tokens": 8192, + "max_tool_iterations": 20 + }, + "list": [ + { + "id": "main", + "plan_model": "anthropic/claude-sonnet-4-6" + }, + { + "id": "advanced", + "plan_model": { + "primary": "anthropic/claude-opus-4", + "fallbacks": ["anthropic/claude-sonnet-4-6"] + } + } + ] + } + }` + + cfg := DefaultConfig() + if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { + t.Fatalf("unmarshal: %v", err) + } + + if len(cfg.Agents.List) != 2 { + t.Fatalf("agents.list len = %d, want 2", len(cfg.Agents.List)) + } + + main := cfg.Agents.List[0] + if main.PlanModel == nil || main.PlanModel.Primary != "anthropic/claude-sonnet-4-6" { + t.Errorf("main.PlanModel = %+v, want primary 'anthropic/claude-sonnet-4-6'", main.PlanModel) + } + + adv := cfg.Agents.List[1] + if adv.PlanModel == nil || adv.PlanModel.Primary != "anthropic/claude-opus-4" { + t.Errorf("advanced.PlanModel = %+v, want primary 'anthropic/claude-opus-4'", adv.PlanModel) + } + if len(adv.PlanModel.Fallbacks) != 1 || adv.PlanModel.Fallbacks[0] != "anthropic/claude-sonnet-4-6" { + t.Errorf("advanced.PlanModel.Fallbacks = %v", adv.PlanModel.Fallbacks) + } +} + +func TestAgentConfig_PlanModel_OverridesDefaults(t *testing.T) { + jsonData := `{ + "agents": { + "defaults": { + "workspace": "~/.picoclaw/workspace", + "model": "glm-4.7", + "plan_model": "default-plan-model", + "plan_model_fallbacks": ["default-fallback"], + "max_tokens": 8192, + "max_tool_iterations": 20 + }, + "list": [ + { + "id": "custom", + "plan_model": { + "primary": "custom-plan-model", + "fallbacks": ["custom-fallback"] + } + } + ] + } + }` + + cfg := DefaultConfig() + if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { + t.Fatalf("unmarshal: %v", err) + } + + custom := cfg.Agents.List[0] + if custom.PlanModel == nil || custom.PlanModel.Primary != "custom-plan-model" { + t.Errorf("custom.PlanModel.Primary = %v, want 'custom-plan-model'", custom.PlanModel) + } + if len(custom.PlanModel.Fallbacks) != 1 || custom.PlanModel.Fallbacks[0] != "custom-fallback" { + t.Errorf("custom.PlanModel.Fallbacks = %v, want [custom-fallback]", custom.PlanModel.Fallbacks) + } + + if cfg.Agents.Defaults.PlanModel != "default-plan-model" { + t.Errorf("defaults.PlanModel = %q, want 'default-plan-model'", cfg.Agents.Defaults.PlanModel) + } +} diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 86155530c..1c93028c7 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -296,7 +296,7 @@ func TestDefaultConfig_WebTools(t *testing.T) { if cfg.Tools.Web.Brave.MaxResults != 5 { t.Error("Expected Brave MaxResults 5, got ", cfg.Tools.Web.Brave.MaxResults) } - if cfg.Tools.Web.Brave.APIKey != "" { + if len(cfg.Tools.Web.Brave.APIKeys) != 0 { t.Error("Brave API key should be empty by default") } if cfg.Tools.Web.DuckDuckGo.MaxResults != 5 { @@ -384,6 +384,13 @@ func TestDefaultConfig_OpenAIWebSearchEnabled(t *testing.T) { } } +func TestDefaultConfig_ExecAllowRemoteEnabled(t *testing.T) { + cfg := DefaultConfig() + if !cfg.Tools.Exec.AllowRemote { + t.Fatal("DefaultConfig().Tools.Exec.AllowRemote should be true") + } +} + func TestLoadConfig_OpenAIWebSearchDefaultsTrueWhenUnset(t *testing.T) { dir := t.TempDir() configPath := filepath.Join(dir, "config.json") @@ -400,6 +407,22 @@ func TestLoadConfig_OpenAIWebSearchDefaultsTrueWhenUnset(t *testing.T) { } } +func TestLoadConfig_ExecAllowRemoteDefaultsTrueWhenUnset(t *testing.T) { + dir := t.TempDir() + configPath := filepath.Join(dir, "config.json") + if err := os.WriteFile(configPath, []byte(`{"tools":{"exec":{"enable_deny_patterns":true}}}`), 0o600); err != nil { + t.Fatalf("WriteFile() error: %v", err) + } + + cfg, err := LoadConfig(configPath) + if err != nil { + t.Fatalf("LoadConfig() error: %v", err) + } + if !cfg.Tools.Exec.AllowRemote { + t.Fatal("tools.exec.allow_remote should remain true when unset in config file") + } +} + func TestLoadConfig_OpenAIWebSearchCanBeDisabled(t *testing.T) { dir := t.TempDir() configPath := filepath.Join(dir, "config.json") @@ -416,133 +439,12 @@ func TestLoadConfig_OpenAIWebSearchCanBeDisabled(t *testing.T) { } } -func TestAgentDefaults_PlanModel_StringParse(t *testing.T) { - jsonData := `{ - "agents": { - "defaults": { - "workspace": "~/.picoclaw/workspace", - "model": "glm-4.7", - "plan_model": "anthropic/claude-sonnet-4-6", - "plan_model_fallbacks": ["openai/gpt-4o"], - "max_tokens": 8192, - "max_tool_iterations": 20 - } - } - }` - - cfg := DefaultConfig() - if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { - t.Fatalf("unmarshal: %v", err) - } - - if cfg.Agents.Defaults.PlanModel != "anthropic/claude-sonnet-4-6" { - t.Errorf("PlanModel = %q, want 'anthropic/claude-sonnet-4-6'", cfg.Agents.Defaults.PlanModel) - } - if len(cfg.Agents.Defaults.PlanModelFallbacks) != 1 || - cfg.Agents.Defaults.PlanModelFallbacks[0] != "openai/gpt-4o" { - t.Errorf("PlanModelFallbacks = %v, want [openai/gpt-4o]", cfg.Agents.Defaults.PlanModelFallbacks) - } -} - -func TestAgentConfig_PlanModel_ObjectParse(t *testing.T) { - jsonData := `{ - "agents": { - "defaults": { - "workspace": "~/.picoclaw/workspace", - "model": "glm-4.7", - "max_tokens": 8192, - "max_tool_iterations": 20 - }, - "list": [ - { - "id": "main", - "plan_model": "anthropic/claude-sonnet-4-6" - }, - { - "id": "advanced", - "plan_model": { - "primary": "anthropic/claude-opus-4", - "fallbacks": ["anthropic/claude-sonnet-4-6"] - } - } - ] - } - }` - - cfg := DefaultConfig() - if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { - t.Fatalf("unmarshal: %v", err) - } - - if len(cfg.Agents.List) != 2 { - t.Fatalf("agents.list len = %d, want 2", len(cfg.Agents.List)) - } - - // String form - main := cfg.Agents.List[0] - if main.PlanModel == nil || main.PlanModel.Primary != "anthropic/claude-sonnet-4-6" { - t.Errorf("main.PlanModel = %+v, want primary 'anthropic/claude-sonnet-4-6'", main.PlanModel) - } - - // Object form with fallbacks - adv := cfg.Agents.List[1] - if adv.PlanModel == nil || adv.PlanModel.Primary != "anthropic/claude-opus-4" { - t.Errorf("advanced.PlanModel = %+v, want primary 'anthropic/claude-opus-4'", adv.PlanModel) - } - if len(adv.PlanModel.Fallbacks) != 1 || adv.PlanModel.Fallbacks[0] != "anthropic/claude-sonnet-4-6" { - t.Errorf("advanced.PlanModel.Fallbacks = %v", adv.PlanModel.Fallbacks) - } -} - -func TestAgentConfig_PlanModel_OverridesDefaults(t *testing.T) { - jsonData := `{ - "agents": { - "defaults": { - "workspace": "~/.picoclaw/workspace", - "model": "glm-4.7", - "plan_model": "default-plan-model", - "plan_model_fallbacks": ["default-fallback"], - "max_tokens": 8192, - "max_tool_iterations": 20 - }, - "list": [ - { - "id": "custom", - "plan_model": { - "primary": "custom-plan-model", - "fallbacks": ["custom-fallback"] - } - } - ] - } - }` - - cfg := DefaultConfig() - if err := json.Unmarshal([]byte(jsonData), cfg); err != nil { - t.Fatalf("unmarshal: %v", err) - } - - // Agent-level plan_model should override defaults - custom := cfg.Agents.List[0] - if custom.PlanModel == nil || custom.PlanModel.Primary != "custom-plan-model" { - t.Errorf("custom.PlanModel.Primary = %v, want 'custom-plan-model'", custom.PlanModel) - } - if len(custom.PlanModel.Fallbacks) != 1 || custom.PlanModel.Fallbacks[0] != "custom-fallback" { - t.Errorf("custom.PlanModel.Fallbacks = %v, want [custom-fallback]", custom.PlanModel.Fallbacks) - } - - // Defaults should still be intact - if cfg.Agents.Defaults.PlanModel != "default-plan-model" { - t.Errorf("defaults.PlanModel = %q, want 'default-plan-model'", cfg.Agents.Defaults.PlanModel) - } -} - func TestLoadConfig_WebToolsProxy(t *testing.T) { tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, "config.json") configJSON := `{ "agents": {"defaults":{"workspace":"./workspace","model":"gpt4","max_tokens":8192,"max_tool_iterations":20}}, - "model_list": [{"model_name":"gpt4","model":"openai/gpt-5.2","api_key":"x"}], + "model_list": [{"model_name":"gpt4","model":"openai/gpt-5.4","api_key":"x"}], "tools": {"web":{"proxy":"http://127.0.0.1:7890"}} }` if err := os.WriteFile(configPath, []byte(configJSON), 0o600); err != nil { @@ -603,3 +505,119 @@ func TestDefaultConfig_WorkspacePath_WithPicoclawHome(t *testing.T) { t.Errorf("Workspace path with PICOCLAW_HOME = %q, want %q", cfg.Agents.Defaults.Workspace, want) } } + +// TestFlexibleStringSlice_UnmarshalText tests UnmarshalText with various comma separators +func TestFlexibleStringSlice_UnmarshalText(t *testing.T) { + tests := []struct { + name string + input string + expected []string + }{ + { + name: "English commas only", + input: "123,456,789", + expected: []string{"123", "456", "789"}, + }, + { + name: "Chinese commas only", + input: "123，456，789", + expected: []string{"123", "456", "789"}, + }, + { + name: "Mixed English and Chinese commas", + input: "123,456，789", + expected: []string{"123", "456", "789"}, + }, + { + name: "Single value", + input: "123", + expected: []string{"123"}, + }, + { + name: "Values with whitespace", + input: " 123 , 456 , 789 ", + expected: []string{"123", "456", "789"}, + }, + { + name: "Empty string", + input: "", + expected: nil, + }, + { + name: "Only commas - English", + input: ",,", + expected: []string{}, + }, + { + name: "Only commas - Chinese", + input: "，，", + expected: []string{}, + }, + { + name: "Mixed commas with empty parts", + input: "123,,456，，789", + expected: []string{"123", "456", "789"}, + }, + { + name: "Complex mixed values", + input: "user1@example.com，user2@test.com, admin@domain.org", + expected: []string{"user1@example.com", "user2@test.com", "admin@domain.org"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var f FlexibleStringSlice + err := f.UnmarshalText([]byte(tt.input)) + if err != nil { + t.Fatalf("UnmarshalText(%q) error = %v", tt.input, err) + } + + if tt.expected == nil { + if f != nil { + t.Errorf("UnmarshalText(%q) = %v, want nil", tt.input, f) + } + return + } + + if len(f) != len(tt.expected) { + t.Errorf("UnmarshalText(%q) length = %d, want %d", tt.input, len(f), len(tt.expected)) + return + } + + for i, v := range tt.expected { + if f[i] != v { + t.Errorf("UnmarshalText(%q)[%d] = %q, want %q", tt.input, i, f[i], v) + } + } + }) + } +} + +// TestFlexibleStringSlice_UnmarshalText_EmptySliceConsistency tests nil vs empty slice behavior +func TestFlexibleStringSlice_UnmarshalText_EmptySliceConsistency(t *testing.T) { + t.Run("Empty string returns nil", func(t *testing.T) { + var f FlexibleStringSlice + err := f.UnmarshalText([]byte("")) + if err != nil { + t.Fatalf("UnmarshalText error = %v", err) + } + if f != nil { + t.Errorf("Empty string should return nil, got %v", f) + } + }) + + t.Run("Commas only returns empty slice", func(t *testing.T) { + var f FlexibleStringSlice + err := f.UnmarshalText([]byte(",,,")) + if err != nil { + t.Fatalf("UnmarshalText error = %v", err) + } + if f == nil { + t.Error("Commas only should return empty slice, not nil") + } + if len(f) != 0 { + t.Errorf("Expected empty slice, got %v", f) + } + }) +} diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index cbdebd43c..09eecbbdf 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -413,6 +413,7 @@ func DefaultConfig() *Config { Enabled: true, }, EnableDenyPatterns: true, + AllowRemote: true, TimeoutSeconds: 60, }, Skills: SkillsToolsConfig{ diff --git a/pkg/config/migration.go b/pkg/config/migration.go index ade9bf677..bb04fd2a8 100644 --- a/pkg/config/migration.go +++ b/pkg/config/migration.go @@ -45,7 +45,7 @@ func ConvertProvidersToModelList(cfg *Config) []ModelConfig { p := cfg.Providers - result := make([]ModelConfig, 0, 20) + var result []ModelConfig // Track if we've applied the legacy model name fix (only for first provider) legacyModelNameApplied := false @@ -61,7 +61,7 @@ func ConvertProvidersToModelList(cfg *Config) []ModelConfig { } return ModelConfig{ ModelName: "openai", - Model: "openai/gpt-5.2", + Model: "openai/gpt-5.4", APIKey: p.OpenAI.APIKey, APIBase: p.OpenAI.APIBase, Proxy: p.OpenAI.Proxy, @@ -335,7 +335,7 @@ func ConvertProvidersToModelList(cfg *Config) []ModelConfig { } return ModelConfig{ ModelName: "github-copilot", - Model: "github-copilot/gpt-5.2", + Model: "github-copilot/gpt-5.4", APIBase: p.GitHubCopilot.APIBase, ConnectMode: p.GitHubCopilot.ConnectMode, }, true diff --git a/pkg/config/migration_test.go b/pkg/config/migration_test.go index d3019aab0..0665ededa 100644 --- a/pkg/config/migration_test.go +++ b/pkg/config/migration_test.go @@ -31,8 +31,8 @@ func TestConvertProvidersToModelList_OpenAI(t *testing.T) { if result[0].ModelName != "openai" { t.Errorf("ModelName = %q, want %q", result[0].ModelName, "openai") } - if result[0].Model != "openai/gpt-5.2" { - t.Errorf("Model = %q, want %q", result[0].Model, "openai/gpt-5.2") + if result[0].Model != "openai/gpt-5.4" { + t.Errorf("Model = %q, want %q", result[0].Model, "openai/gpt-5.4") } if result[0].APIKey != "sk-test-key" { t.Errorf("APIKey = %q, want %q", result[0].APIKey, "sk-test-key") @@ -162,14 +162,15 @@ func TestConvertProvidersToModelList_AllProviders(t *testing.T) { Qwen: ProviderConfig{APIKey: "key17"}, Mistral: ProviderConfig{APIKey: "key18"}, Avian: ProviderConfig{APIKey: "key19"}, + LongCat: ProviderConfig{APIKey: "key-longcat"}, }, } result := ConvertProvidersToModelList(cfg) - // All 21 providers should be converted - if len(result) != 21 { - t.Errorf("len(result) = %d, want 21", len(result)) + // All 22 providers should be converted + if len(result) != 22 { + t.Errorf("len(result) = %d, want 22", len(result)) } } @@ -383,8 +384,8 @@ func TestConvertProvidersToModelList_MultipleProviders_PreservesUserModel(t *tes for _, mc := range result { switch mc.ModelName { case "openai": - if mc.Model != "openai/gpt-5.2" { - t.Errorf("OpenAI Model = %q, want %q (default)", mc.Model, "openai/gpt-5.2") + if mc.Model != "openai/gpt-5.4" { + t.Errorf("OpenAI Model = %q, want %q (default)", mc.Model, "openai/gpt-5.4") } case "deepseek": if mc.Model != "deepseek/deepseek-reasoner" { @@ -557,9 +558,9 @@ func TestConvertProvidersToModelList_NoProviderField_NoModel(t *testing.T) { // Tests for buildModelWithProtocol helper function func TestBuildModelWithProtocol_NoPrefix(t *testing.T) { - result := buildModelWithProtocol("openai", "gpt-5.2") - if result != "openai/gpt-5.2" { - t.Errorf("buildModelWithProtocol(openai, gpt-5.2) = %q, want %q", result, "openai/gpt-5.2") + result := buildModelWithProtocol("openai", "gpt-5.4") + if result != "openai/gpt-5.4" { + t.Errorf("buildModelWithProtocol(openai, gpt-5.4) = %q, want %q", result, "openai/gpt-5.4") } } diff --git a/pkg/heartbeat/service_ext_test.go b/pkg/heartbeat/service_ext_test.go new file mode 100644 index 000000000..8c2338073 --- /dev/null +++ b/pkg/heartbeat/service_ext_test.go @@ -0,0 +1,103 @@ +package heartbeat + +import ( + "github.com/sipeed/picoclaw/pkg/tools" + "os" + "path/filepath" + "testing" +) + +// TestExecuteHeartbeat_NoSendResponse verifies that heartbeat results +// do not trigger sendResponse (dedup: response is included in task status instead). +// TestExecuteHeartbeat_NoSendResponse verifies that heartbeat results +// do not trigger sendResponse (dedup: response is included in task status instead). +func TestExecuteHeartbeat_NoSendResponse(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "heartbeat-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + hs := NewHeartbeatService(tmpDir, 30, true) + hs.stopChan = make(chan struct{}) + + hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult { + return &tools.ToolResult{ + ForUser: "Task result for user", + ForLLM: "Task result for LLM", + Silent: false, + IsError: false, + Async: false, + } + }) + + os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644) + + hs.executeHeartbeat() + + hs.mu.RLock() + notified := !hs.lastNotifiedAt.IsZero() + hs.mu.RUnlock() + if !notified { + t.Error("Expected lastNotifiedAt to be set after heartbeat completion") + } +} + +func TestExecuteHeartbeat_TargetPriority_ExplicitTarget(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "heartbeat-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + hs := NewHeartbeatService(tmpDir, 30, true) + hs.stopChan = make(chan struct{}) + hs.SetHeartbeatThreadID(77) + if err := hs.state.SetHeartbeatTarget("slack:C12345/999"); err != nil { + t.Fatalf("SetHeartbeatTarget failed: %v", err) + } + if err := hs.state.SetLastHeartbeatTarget("telegram:-100500"); err != nil { + t.Fatalf("SetLastHeartbeatTarget failed: %v", err) + } + + var gotChannel, gotChatID string + hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult { + gotChannel, gotChatID = channel, chatID + return tools.SilentResult("ok") + }) + os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644) + + hs.executeHeartbeat() + + if gotChannel != "slack" || gotChatID != "C12345/999" { + t.Fatalf("handler target = %s:%s, want slack:C12345/999", gotChannel, gotChatID) + } +} + +func TestExecuteHeartbeat_TargetPriority_TelegramThread(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "heartbeat-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + hs := NewHeartbeatService(tmpDir, 30, true) + hs.stopChan = make(chan struct{}) + hs.SetHeartbeatThreadID(77) + if err := hs.state.SetLastHeartbeatTarget("telegram:-100500"); err != nil { + t.Fatalf("SetLastHeartbeatTarget failed: %v", err) + } + + var gotChannel, gotChatID string + hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult { + gotChannel, gotChatID = channel, chatID + return tools.SilentResult("ok") + }) + os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644) + + hs.executeHeartbeat() + + if gotChannel != "telegram" || gotChatID != "-100500/77" { + t.Fatalf("handler target = %s:%s, want telegram:-100500/77", gotChannel, gotChatID) + } +} diff --git a/pkg/heartbeat/service_test.go b/pkg/heartbeat/service_test.go index 8b34ebf6c..3b7eeeefb 100644 --- a/pkg/heartbeat/service_test.go +++ b/pkg/heartbeat/service_test.go @@ -184,42 +184,6 @@ func TestLogPath(t *testing.T) { } } -// TestExecuteHeartbeat_NoSendResponse verifies that heartbeat results -// do not trigger sendResponse (dedup: response is included in task status instead). -func TestExecuteHeartbeat_NoSendResponse(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "heartbeat-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - - hs := NewHeartbeatService(tmpDir, 30, true) - hs.stopChan = make(chan struct{}) - - hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult { - return &tools.ToolResult{ - ForUser: "Task result for user", - ForLLM: "Task result for LLM", - Silent: false, - IsError: false, - Async: false, - } - }) - - os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644) - - // Execute heartbeat — since bus is nil, sendResponse would log but not crash. - // The key assertion is that lastNotifiedAt is still updated (flow reaches end). - hs.executeHeartbeat() - - hs.mu.RLock() - notified := !hs.lastNotifiedAt.IsZero() - hs.mu.RUnlock() - if !notified { - t.Error("Expected lastNotifiedAt to be set after heartbeat completion") - } -} - // TestHeartbeatFilePath verifies HEARTBEAT.md is at workspace root func TestHeartbeatFilePath(t *testing.T) { tmpDir, err := os.MkdirTemp("", "heartbeat-test-*") @@ -239,62 +203,3 @@ func TestHeartbeatFilePath(t *testing.T) { t.Errorf("Expected HEARTBEAT.md at %s, but it doesn't exist", expectedPath) } } - -func TestExecuteHeartbeat_TargetPriority_ExplicitTarget(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "heartbeat-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - - hs := NewHeartbeatService(tmpDir, 30, true) - hs.stopChan = make(chan struct{}) - hs.SetHeartbeatThreadID(77) - if err := hs.state.SetHeartbeatTarget("slack:C12345/999"); err != nil { - t.Fatalf("SetHeartbeatTarget failed: %v", err) - } - if err := hs.state.SetLastHeartbeatTarget("telegram:-100500"); err != nil { - t.Fatalf("SetLastHeartbeatTarget failed: %v", err) - } - - var gotChannel, gotChatID string - hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult { - gotChannel, gotChatID = channel, chatID - return tools.SilentResult("ok") - }) - os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644) - - hs.executeHeartbeat() - - if gotChannel != "slack" || gotChatID != "C12345/999" { - t.Fatalf("handler target = %s:%s, want slack:C12345/999", gotChannel, gotChatID) - } -} - -func TestExecuteHeartbeat_TargetPriority_TelegramThread(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "heartbeat-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - - hs := NewHeartbeatService(tmpDir, 30, true) - hs.stopChan = make(chan struct{}) - hs.SetHeartbeatThreadID(77) - if err := hs.state.SetLastHeartbeatTarget("telegram:-100500"); err != nil { - t.Fatalf("SetLastHeartbeatTarget failed: %v", err) - } - - var gotChannel, gotChatID string - hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult { - gotChannel, gotChatID = channel, chatID - return tools.SilentResult("ok") - }) - os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644) - - hs.executeHeartbeat() - - if gotChannel != "telegram" || gotChatID != "-100500/77" { - t.Fatalf("handler target = %s:%s, want telegram:-100500/77", gotChannel, gotChatID) - } -} diff --git a/pkg/logger/logger_ext_test.go b/pkg/logger/logger_ext_test.go new file mode 100644 index 000000000..665dcf2f5 --- /dev/null +++ b/pkg/logger/logger_ext_test.go @@ -0,0 +1,289 @@ +package logger + +import ( + "testing" + "time" +) + +func TestRingBuffer_PushAndRecent(t *testing.T) { + rb := newLogRingBuffer(5) + + for i := 0; i < 3; i++ { + rb.push(LogEntry{Message: "msg" + string(rune('A'+i))}) + } + + got := rb.recent(0) + if len(got) != 3 { + t.Fatalf("expected 3, got %d", len(got)) + } + if got[0].Message != "msgA" || got[2].Message != "msgC" { + t.Errorf("unexpected order: %v", got) + } +} + +func TestRingBuffer_Wrap(t *testing.T) { + rb := newLogRingBuffer(3) + for i := 0; i < 5; i++ { + rb.push(LogEntry{Message: string(rune('A' + i))}) + } + + got := rb.recent(0) + if len(got) != 3 { + t.Fatalf("expected 3, got %d", len(got)) + } + + if got[0].Message != "C" || got[1].Message != "D" || got[2].Message != "E" { + t.Errorf("expected [C,D,E], got [%s,%s,%s]", got[0].Message, got[1].Message, got[2].Message) + } +} + +func TestRingBuffer_RecentLimit(t *testing.T) { + rb := newLogRingBuffer(10) + for i := 0; i < 8; i++ { + rb.push(LogEntry{Message: string(rune('A' + i))}) + } + + got := rb.recent(3) + if len(got) != 3 { + t.Fatalf("expected 3, got %d", len(got)) + } + if got[0].Message != "F" || got[2].Message != "H" { + t.Errorf("expected last 3 entries, got %v", got) + } +} + +func TestRecentLogs_FilterByLevel(t *testing.T) { + initialLevel := GetLevel() + defer SetLevel(initialLevel) + SetLevel(DEBUG) + + DebugC("test", "debug msg") + InfoC("test", "info msg") + WarnC("test", "warn msg") + ErrorC("test", "error msg") + + got := RecentLogs(WARN, "", 100) + for _, e := range got { + if e.Level == "DEBUG" || e.Level == "INFO" { + t.Errorf("unexpected level %s in result with minLevel=WARN", e.Level) + } + } +} + +func TestRecentLogs_FilterByComponent(t *testing.T) { + initialLevel := GetLevel() + defer SetLevel(initialLevel) + SetLevel(DEBUG) + + InfoC("alpha", "from alpha") + InfoC("beta", "from beta") + InfoC("alpha", "another from alpha") + + got := RecentLogs(DEBUG, "alpha", 100) + for _, e := range got { + if e.Component != "alpha" { + t.Errorf("unexpected component %s in result with component=alpha", e.Component) + } + } +} + +func TestRecentLogs_CallerStripped(t *testing.T) { + initialLevel := GetLevel() + defer SetLevel(initialLevel) + SetLevel(DEBUG) + + InfoC("test", "caller test") + + got := RecentLogs(DEBUG, "", 100) + for _, e := range got { + if e.Caller != "" { + t.Errorf("Caller should be stripped, got %q", e.Caller) + } + } +} + +func TestSubscribe_ReceivesEntries(t *testing.T) { + initialLevel := GetLevel() + defer SetLevel(initialLevel) + SetLevel(DEBUG) + + sub := Subscribe(nil) + defer Unsubscribe(sub) + + InfoC("sub-test", "hello subscriber") + + select { + case entry := <-sub.Ch: + if entry.Message != "hello subscriber" { + t.Errorf("expected 'hello subscriber', got %q", entry.Message) + } + case <-time.After(time.Second): + t.Error("timed out waiting for log entry") + } +} + +func TestSubscribe_FilterApplied(t *testing.T) { + initialLevel := GetLevel() + defer SetLevel(initialLevel) + SetLevel(DEBUG) + + sub := Subscribe(func(e LogEntry) bool { + return e.Component == "target" + }) + defer Unsubscribe(sub) + + InfoC("other", "should be filtered out") + InfoC("target", "should arrive") + + select { + case entry := <-sub.Ch: + if entry.Component != "target" { + t.Errorf("expected component=target, got %q", entry.Component) + } + case <-time.After(time.Second): + t.Error("timed out waiting for filtered entry") + } +} + +func TestUnsubscribe_ClosesChannel(t *testing.T) { + sub := Subscribe(nil) + Unsubscribe(sub) + + _, ok := <-sub.Ch + if ok { + t.Error("expected channel to be closed after Unsubscribe") + } +} + +func TestSanitizeFields(t *testing.T) { + tests := []struct { + name string + input map[string]any + maskedK []string // keys that should be "***" + safeK []string // keys that should keep original value + }{ + { + name: "nil fields", + input: nil, + maskedK: nil, + }, + { + name: "empty fields", + input: map[string]any{}, + maskedK: nil, + }, + { + name: "sensitive keys masked", + input: map[string]any{ + "token": "abc123", + "api_key": "sk-xxx", + "secret": "s3cr3t", + "password": "pass", + "authorization": "Bearer tok", + }, + maskedK: []string{"token", "api_key", "secret", "password", "authorization"}, + }, + { + name: "case insensitive", + input: map[string]any{ + "Token": "abc", + "API_KEY": "xyz", + "Secret": "s", + "PASSWORD": "p", + "Authorization": "a", + "Credential": "c", + }, + maskedK: []string{"Token", "API_KEY", "Secret", "PASSWORD", "Authorization", "Credential"}, + }, + { + name: "safe keys preserved", + input: map[string]any{"error": "something failed", "count": 42, "user_id": "12345", "component": "test"}, + safeK: []string{"error", "count", "user_id", "component"}, + }, + { + name: "mixed keys", + input: map[string]any{ + "token": "sensitive", + "msg_signature": "safe", + "corp_secret": "sensitive2", + "nonce": "safe2", + }, + maskedK: []string{"token", "corp_secret"}, + safeK: []string{"msg_signature", "nonce"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := SanitizeFields(tt.input) + for _, k := range tt.maskedK { + if v, ok := result[k]; !ok || v != "***" { + t.Errorf("expected key %q to be masked, got %v", k, v) + } + } + for _, k := range tt.safeK { + if result[k] != tt.input[k] { + t.Errorf("expected key %q to be preserved as %v, got %v", k, tt.input[k], result[k]) + } + } + }) + } +} + +func TestSanitizeFieldsDoesNotMutateOriginal(t *testing.T) { + original := map[string]any{"token": "secret_value", "name": "test"} + _ = SanitizeFields(original) + if original["token"] != "secret_value" { + t.Error("SanitizeFields should not mutate the original map") + } +} + +func TestRecentLogsSanitizesFields(t *testing.T) { + initialLevel := GetLevel() + defer SetLevel(initialLevel) + SetLevel(DEBUG) + + InfoCF("sanitize-test", "log with sensitive fields", map[string]any{ + "token": "my-secret-token", + "api_key": "sk-12345", + "user_id": "safe-value", + }) + + got := RecentLogs(DEBUG, "sanitize-test", 100) + if len(got) == 0 { + t.Fatal("expected at least one log entry") + } + + last := got[len(got)-1] + if last.Fields["token"] != "***" { + t.Errorf("expected token to be masked, got %v", last.Fields["token"]) + } + if last.Fields["api_key"] != "***" { + t.Errorf("expected api_key to be masked, got %v", last.Fields["api_key"]) + } + if last.Fields["user_id"] != "safe-value" { + t.Errorf("expected user_id to be preserved, got %v", last.Fields["user_id"]) + } +} + +func TestParseLevel(t *testing.T) { + tests := []struct { + input string + want LogLevel + }{ + {"debug", DEBUG}, + {"DEBUG", DEBUG}, + {"info", INFO}, + {"WARN", WARN}, + {"error", ERROR}, + {"fatal", FATAL}, + {"unknown", INFO}, + {"", INFO}, + } + for _, tt := range tests { + got := ParseLevel(tt.input) + if got != tt.want { + t.Errorf("ParseLevel(%q) = %d, want %d", tt.input, got, tt.want) + } + } +} diff --git a/pkg/logger/logger_test.go b/pkg/logger/logger_test.go index 741b61209..6e6f8dfa8 100644 --- a/pkg/logger/logger_test.go +++ b/pkg/logger/logger_test.go @@ -2,7 +2,6 @@ package logger import ( "testing" - "time" ) func TestLogLevelFiltering(t *testing.T) { @@ -138,289 +137,3 @@ func TestLoggerHelperFunctions(t *testing.T) { DebugC("test", "Debug with component") WarnF("Warning with fields", map[string]any{"key": "value"}) } - -// ── Ring buffer tests ── - -func TestRingBuffer_PushAndRecent(t *testing.T) { - rb := newLogRingBuffer(5) - - for i := 0; i < 3; i++ { - rb.push(LogEntry{Message: "msg" + string(rune('A'+i))}) - } - - got := rb.recent(0) - if len(got) != 3 { - t.Fatalf("expected 3, got %d", len(got)) - } - if got[0].Message != "msgA" || got[2].Message != "msgC" { - t.Errorf("unexpected order: %v", got) - } -} - -func TestRingBuffer_Wrap(t *testing.T) { - rb := newLogRingBuffer(3) - for i := 0; i < 5; i++ { - rb.push(LogEntry{Message: string(rune('A' + i))}) - } - - got := rb.recent(0) - if len(got) != 3 { - t.Fatalf("expected 3, got %d", len(got)) - } - // Should have C, D, E (oldest two dropped) - if got[0].Message != "C" || got[1].Message != "D" || got[2].Message != "E" { - t.Errorf("expected [C,D,E], got [%s,%s,%s]", got[0].Message, got[1].Message, got[2].Message) - } -} - -func TestRingBuffer_RecentLimit(t *testing.T) { - rb := newLogRingBuffer(10) - for i := 0; i < 8; i++ { - rb.push(LogEntry{Message: string(rune('A' + i))}) - } - - got := rb.recent(3) - if len(got) != 3 { - t.Fatalf("expected 3, got %d", len(got)) - } - if got[0].Message != "F" || got[2].Message != "H" { - t.Errorf("expected last 3 entries, got %v", got) - } -} - -func TestRecentLogs_FilterByLevel(t *testing.T) { - initialLevel := GetLevel() - defer SetLevel(initialLevel) - SetLevel(DEBUG) - - // Log messages at different levels - DebugC("test", "debug msg") - InfoC("test", "info msg") - WarnC("test", "warn msg") - ErrorC("test", "error msg") - - got := RecentLogs(WARN, "", 100) - for _, e := range got { - if e.Level == "DEBUG" || e.Level == "INFO" { - t.Errorf("unexpected level %s in result with minLevel=WARN", e.Level) - } - } -} - -func TestRecentLogs_FilterByComponent(t *testing.T) { - initialLevel := GetLevel() - defer SetLevel(initialLevel) - SetLevel(DEBUG) - - InfoC("alpha", "from alpha") - InfoC("beta", "from beta") - InfoC("alpha", "another from alpha") - - got := RecentLogs(DEBUG, "alpha", 100) - for _, e := range got { - if e.Component != "alpha" { - t.Errorf("unexpected component %s in result with component=alpha", e.Component) - } - } -} - -func TestRecentLogs_CallerStripped(t *testing.T) { - initialLevel := GetLevel() - defer SetLevel(initialLevel) - SetLevel(DEBUG) - - InfoC("test", "caller test") - - got := RecentLogs(DEBUG, "", 100) - for _, e := range got { - if e.Caller != "" { - t.Errorf("Caller should be stripped, got %q", e.Caller) - } - } -} - -func TestSubscribe_ReceivesEntries(t *testing.T) { - initialLevel := GetLevel() - defer SetLevel(initialLevel) - SetLevel(DEBUG) - - sub := Subscribe(nil) - defer Unsubscribe(sub) - - InfoC("sub-test", "hello subscriber") - - select { - case entry := <-sub.Ch: - if entry.Message != "hello subscriber" { - t.Errorf("expected 'hello subscriber', got %q", entry.Message) - } - case <-time.After(time.Second): - t.Error("timed out waiting for log entry") - } -} - -func TestSubscribe_FilterApplied(t *testing.T) { - initialLevel := GetLevel() - defer SetLevel(initialLevel) - SetLevel(DEBUG) - - sub := Subscribe(func(e LogEntry) bool { - return e.Component == "target" - }) - defer Unsubscribe(sub) - - InfoC("other", "should be filtered out") - InfoC("target", "should arrive") - - select { - case entry := <-sub.Ch: - if entry.Component != "target" { - t.Errorf("expected component=target, got %q", entry.Component) - } - case <-time.After(time.Second): - t.Error("timed out waiting for filtered entry") - } -} - -func TestUnsubscribe_ClosesChannel(t *testing.T) { - sub := Subscribe(nil) - Unsubscribe(sub) - - _, ok := <-sub.Ch - if ok { - t.Error("expected channel to be closed after Unsubscribe") - } -} - -func TestSanitizeFields(t *testing.T) { - tests := []struct { - name string - input map[string]any - maskedK []string // keys that should be "***" - safeK []string // keys that should keep original value - }{ - { - name: "nil fields", - input: nil, - maskedK: nil, - }, - { - name: "empty fields", - input: map[string]any{}, - maskedK: nil, - }, - { - name: "sensitive keys masked", - input: map[string]any{ - "token": "abc123", - "api_key": "sk-xxx", - "secret": "s3cr3t", - "password": "pass", - "authorization": "Bearer tok", - }, - maskedK: []string{"token", "api_key", "secret", "password", "authorization"}, - }, - { - name: "case insensitive", - input: map[string]any{ - "Token": "abc", - "API_KEY": "xyz", - "Secret": "s", - "PASSWORD": "p", - "Authorization": "a", - "Credential": "c", - }, - maskedK: []string{"Token", "API_KEY", "Secret", "PASSWORD", "Authorization", "Credential"}, - }, - { - name: "safe keys preserved", - input: map[string]any{"error": "something failed", "count": 42, "user_id": "12345", "component": "test"}, - safeK: []string{"error", "count", "user_id", "component"}, - }, - { - name: "mixed keys", - input: map[string]any{ - "token": "sensitive", - "msg_signature": "safe", - "corp_secret": "sensitive2", - "nonce": "safe2", - }, - maskedK: []string{"token", "corp_secret"}, - safeK: []string{"msg_signature", "nonce"}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - result := SanitizeFields(tt.input) - for _, k := range tt.maskedK { - if v, ok := result[k]; !ok || v != "***" { - t.Errorf("expected key %q to be masked, got %v", k, v) - } - } - for _, k := range tt.safeK { - if result[k] != tt.input[k] { - t.Errorf("expected key %q to be preserved as %v, got %v", k, tt.input[k], result[k]) - } - } - }) - } -} - -func TestSanitizeFieldsDoesNotMutateOriginal(t *testing.T) { - original := map[string]any{"token": "secret_value", "name": "test"} - _ = SanitizeFields(original) - if original["token"] != "secret_value" { - t.Error("SanitizeFields should not mutate the original map") - } -} - -func TestRecentLogsSanitizesFields(t *testing.T) { - initialLevel := GetLevel() - defer SetLevel(initialLevel) - SetLevel(DEBUG) - - InfoCF("sanitize-test", "log with sensitive fields", map[string]any{ - "token": "my-secret-token", - "api_key": "sk-12345", - "user_id": "safe-value", - }) - - got := RecentLogs(DEBUG, "sanitize-test", 100) - if len(got) == 0 { - t.Fatal("expected at least one log entry") - } - - last := got[len(got)-1] - if last.Fields["token"] != "***" { - t.Errorf("expected token to be masked, got %v", last.Fields["token"]) - } - if last.Fields["api_key"] != "***" { - t.Errorf("expected api_key to be masked, got %v", last.Fields["api_key"]) - } - if last.Fields["user_id"] != "safe-value" { - t.Errorf("expected user_id to be preserved, got %v", last.Fields["user_id"]) - } -} - -func TestParseLevel(t *testing.T) { - tests := []struct { - input string - want LogLevel - }{ - {"debug", DEBUG}, - {"DEBUG", DEBUG}, - {"info", INFO}, - {"WARN", WARN}, - {"error", ERROR}, - {"fatal", FATAL}, - {"unknown", INFO}, - {"", INFO}, - } - for _, tt := range tests { - got := ParseLevel(tt.input) - if got != tt.want { - t.Errorf("ParseLevel(%q) = %d, want %d", tt.input, got, tt.want) - } - } -} diff --git a/pkg/memory/migration_test.go b/pkg/memory/migration_test.go index 3170758b7..4466c96f9 100644 --- a/pkg/memory/migration_test.go +++ b/pkg/memory/migration_test.go @@ -382,3 +382,55 @@ func TestMigrateFromJSON_NonexistentDir(t *testing.T) { t.Errorf("expected 0, got %d", count) } } + +func TestMigrateFromJSON_SkipsMetaJSONFiles(t *testing.T) { + sessionsDir := t.TempDir() + store, err := NewJSONLStore(sessionsDir) + if err != nil { + t.Fatalf("NewJSONLStore: %v", err) + } + ctx := context.Background() + + if addErr := store.AddMessage(ctx, "agent:main:pico:direct:pico:test", "user", "keep me"); addErr != nil { + t.Fatalf("AddMessage: %v", addErr) + } + if summaryErr := store.SetSummary(ctx, "agent:main:pico:direct:pico:test", "keep summary"); summaryErr != nil { + t.Fatalf("SetSummary: %v", summaryErr) + } + + metaPath := filepath.Join(sessionsDir, "agent_main_pico_direct_pico_test.meta.json") + if _, statErr := os.Stat(metaPath); statErr != nil { + t.Fatalf("meta file missing before migration: %v", statErr) + } + + count, err := MigrateFromJSON(ctx, sessionsDir, store) + if err != nil { + t.Fatalf("MigrateFromJSON: %v", err) + } + if count != 0 { + t.Fatalf("expected 0 migrated, got %d", count) + } + + history, err := store.GetHistory(ctx, "agent:main:pico:direct:pico:test") + if err != nil { + t.Fatalf("GetHistory: %v", err) + } + if len(history) != 1 || history[0].Content != "keep me" { + t.Fatalf("history = %+v, want preserved single message", history) + } + + summary, err := store.GetSummary(ctx, "agent:main:pico:direct:pico:test") + if err != nil { + t.Fatalf("GetSummary: %v", err) + } + if summary != "keep summary" { + t.Fatalf("summary = %q, want %q", summary, "keep summary") + } + + if _, statErr := os.Stat(metaPath); statErr != nil { + t.Fatalf("meta file should remain in place: %v", statErr) + } + if _, statErr := os.Stat(metaPath + ".migrated"); !os.IsNotExist(statErr) { + t.Fatalf("meta file should not be renamed, stat err = %v", statErr) + } +} diff --git a/pkg/migrate/sources/openclaw/openclaw_config.go b/pkg/migrate/sources/openclaw/openclaw_config.go index 19d63bb77..92fcf0f62 100644 --- a/pkg/migrate/sources/openclaw/openclaw_config.go +++ b/pkg/migrate/sources/openclaw/openclaw_config.go @@ -1107,6 +1107,7 @@ func (c ToolsConfig) ToStandardTools() config.ToolsConfig { Exec: config.ExecConfig{ EnableDenyPatterns: c.Exec.EnableDenyPatterns, CustomDenyPatterns: c.Exec.CustomDenyPatterns, + AllowRemote: config.DefaultConfig().Tools.Exec.AllowRemote, }, } } diff --git a/pkg/migrate/sources/openclaw/openclaw_config_test.go b/pkg/migrate/sources/openclaw/openclaw_config_test.go index 3a7d0c686..802693825 100644 --- a/pkg/migrate/sources/openclaw/openclaw_config_test.go +++ b/pkg/migrate/sources/openclaw/openclaw_config_test.go @@ -290,6 +290,20 @@ func TestConvertToPicoClaw(t *testing.T) { } } +func TestToStandardConfig_ExecAllowRemoteDefaultsTrue(t *testing.T) { + cfg := (&PicoClawConfig{ + Tools: ToolsConfig{ + Exec: ExecConfig{ + EnableDenyPatterns: true, + }, + }, + }).ToStandardConfig() + + if !cfg.Tools.Exec.AllowRemote { + t.Fatal("ToStandardConfig() should preserve the default tools.exec.allow_remote=true") + } +} + func TestConvertToPicoClawWithQQAndDingTalk(t *testing.T) { tmpDir := t.TempDir() configPath := filepath.Join(tmpDir, "openclaw.json") diff --git a/pkg/providers/anthropic/provider_test.go b/pkg/providers/anthropic/provider_test.go index 2c8a8e6bf..b1aed17b5 100644 --- a/pkg/providers/anthropic/provider_test.go +++ b/pkg/providers/anthropic/provider_test.go @@ -9,8 +9,6 @@ import ( "github.com/anthropics/anthropic-sdk-go" anthropicoption "github.com/anthropics/anthropic-sdk-go/option" - - "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" ) func TestBuildParams_BasicMessage(t *testing.T) { @@ -86,13 +84,13 @@ func TestBuildParams_WithTools(t *testing.T) { Function: ToolFunctionDefinition{ Name: "get_weather", Description: "Get weather for a city", - Parameters: protocoltypes.MustMarshalParameters(map[string]any{ + Parameters: map[string]any{ "type": "object", "properties": map[string]any{ "city": map[string]any{"type": "string"}, }, "required": []any{"city"}, - }), + }, }, }, } diff --git a/pkg/providers/antigravity_provider_test.go b/pkg/providers/antigravity_provider_test.go index a7eebbedc..238765321 100644 --- a/pkg/providers/antigravity_provider_test.go +++ b/pkg/providers/antigravity_provider_test.go @@ -12,7 +12,7 @@ func TestBuildRequestUsesFunctionFieldsWhenToolCallNameMissing(t *testing.T) { ID: "call_read_file_123", Function: &FunctionCall{ Name: "read_file", - Arguments: map[string]any{"path": "README.md"}, + Arguments: `{"path":"README.md"}`, }, }}, }, diff --git a/pkg/providers/claude_cli_provider_ext_test.go b/pkg/providers/claude_cli_provider_ext_test.go new file mode 100644 index 000000000..8c0164c75 --- /dev/null +++ b/pkg/providers/claude_cli_provider_ext_test.go @@ -0,0 +1,283 @@ +package providers + +import ( + "strings" + "testing" +) + +func TestExtractXMLToolCalls_Single(t *testing.T) { + text := ` + +echo hello + +` + + calls := extractXMLToolCalls(text) + if len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(calls)) + } + if calls[0].Name != "exec" { + t.Errorf("Name = %q, want %q", calls[0].Name, "exec") + } + if calls[0].Arguments["command"] != "echo hello" { + t.Errorf("Arguments[command] = %v, want %q", calls[0].Arguments["command"], "echo hello") + } + if calls[0].Function == nil || calls[0].Function.Name != "exec" { + t.Errorf("Function.Name should be exec") + } +} + +func TestExtractXMLToolCalls_Multiple(t *testing.T) { + text := ` + +golang testing + + +go test ./... +30 + +` + + calls := extractXMLToolCalls(text) + if len(calls) != 2 { + t.Fatalf("expected 2 tool calls, got %d", len(calls)) + } + if calls[0].Name != "web_search" { + t.Errorf("[0].Name = %q, want %q", calls[0].Name, "web_search") + } + if calls[1].Name != "exec" { + t.Errorf("[1].Name = %q, want %q", calls[1].Name, "exec") + } + if calls[1].Arguments["timeout"] != "30" { + t.Errorf("[1].Arguments[timeout] = %v, want %q", calls[1].Arguments["timeout"], "30") + } +} + +func TestExtractXMLToolCalls_NoXML(t *testing.T) { + calls := extractXMLToolCalls("just regular text") + if len(calls) != 0 { + t.Errorf("expected 0 tool calls, got %d", len(calls)) + } +} + +func TestStripXMLToolCalls(t *testing.T) { + text := `Let me run that. + + +echo hello + + +Done.` + + got := stripXMLToolCalls(text) + if strings.Contains(got, "toolcall") { + t.Errorf("should remove XML block, got %q", got) + } + if !strings.Contains(got, "Let me run that.") { + t.Errorf("should keep text before, got %q", got) + } + if !strings.Contains(got, "Done.") { + t.Errorf("should keep text after, got %q", got) + } +} + +func TestExtractXMLToolCalls_MismatchedCloseTag(t *testing.T) { + + text := ` + +/home/user/project/pyproject.toml + +` + + calls := extractXMLToolCalls(text) + if len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(calls)) + } + if calls[0].Name != "readfile" { + t.Errorf("Name = %q, want %q", calls[0].Name, "readfile") + } + if calls[0].Arguments["path"] != "/home/user/project/pyproject.toml" { + t.Errorf("Arguments[path] = %v, want pyproject.toml path", calls[0].Arguments["path"]) + } +} + +func TestStripXMLToolCalls_MismatchedCloseTag(t *testing.T) { + text := `今テスト走らせるね。` + + ` + + +cd /home/user && pytest + +` + + got := stripXMLToolCalls(text) + if strings.Contains(got, "toolcall") || strings.Contains(got, "tool_call") { + t.Errorf("should remove XML block, got %q", got) + } + if !strings.Contains(got, "今テスト走らせるね。") { + t.Errorf("should keep text before, got %q", got) + } +} + +func TestExtractXMLToolCalls_UnderscoreOpenTag(t *testing.T) { + + text := ` + +ls -la + +` + + calls := extractXMLToolCalls(text) + if len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(calls)) + } + if calls[0].Name != "exec" { + t.Errorf("Name = %q, want %q", calls[0].Name, "exec") + } + if calls[0].Arguments["command"] != "ls -la" { + t.Errorf("Arguments[command] = %v, want %q", calls[0].Arguments["command"], "ls -la") + } +} + +func TestExtractXMLToolCalls_HyphenTag(t *testing.T) { + + text := ` + +/etc/hosts + +` + + calls := extractXMLToolCalls(text) + if len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(calls)) + } + if calls[0].Name != "read_file" { + t.Errorf("Name = %q, want %q", calls[0].Name, "read_file") + } +} + +func TestStripXMLToolCalls_UnderscoreOpenTag(t *testing.T) { + text := `Here is the result. + + +ls + + +Finished.` + + got := stripXMLToolCalls(text) + if strings.Contains(got, "tool_call") || strings.Contains(got, "toolcall") { + t.Errorf("should remove XML block, got %q", got) + } + if !strings.Contains(got, "Here is the result.") { + t.Errorf("should keep text before, got %q", got) + } + if !strings.Contains(got, "Finished.") { + t.Errorf("should keep text after, got %q", got) + } +} + +func TestExtractXMLToolCalls_OrphanedClosingTag(t *testing.T) { + + text := "了解！確認するね。\n[TOOLCALL]\n\n/home/user/workspace\n\n" + + calls := extractXMLToolCalls(text) + if len(calls) != 1 { + t.Fatalf("expected 1 tool call, got %d", len(calls)) + } + if calls[0].Name != "listdir" { + t.Errorf("Name = %q, want %q", calls[0].Name, "listdir") + } + if calls[0].Arguments["path"] != "/home/user/workspace" { + t.Errorf("Arguments[path] = %v, want /home/user/workspace", calls[0].Arguments["path"]) + } +} + +func TestStripXMLToolCalls_OrphanedClosingTag(t *testing.T) { + text := "了解！確認するね。\n[TOOLCALL]\n\n/home/user\n\n" + got := stripXMLToolCalls(text) + if strings.Contains(got, "invoke") || strings.Contains(got, "TOOLCALL") || strings.Contains(got, "minimax") { + t.Errorf("should remove orphaned closing tag block, got %q", got) + } + if !strings.Contains(got, "了解") { + t.Errorf("should keep user-facing text, got %q", got) + } +} + +func TestStripXMLToolCalls_NoXML(t *testing.T) { + text := "Just regular text." + got := stripXMLToolCalls(text) + if got != text { + t.Errorf("stripXMLToolCalls() = %q, want %q", got, text) + } +} + +func TestNormalizeAlpha(t *testing.T) { + tests := []struct { + input, want string + }{ + {"toolcall", "toolcall"}, + {"tool_call", "toolcall"}, + {"Tool-Call", "toolcall"}, + {"ReadFile", "readfile"}, + {"read_file", "readfile"}, + {"EXEC", "exec"}, + {"web123search", "websearch"}, + {"", ""}, + } + for _, tt := range tests { + got := normalizeAlpha(tt.input) + if got != tt.want { + t.Errorf("normalizeAlpha(%q) = %q, want %q", tt.input, got, tt.want) + } + } +} + +func TestLevenshtein(t *testing.T) { + tests := []struct { + a, b string + want int + }{ + {"", "", 0}, + {"abc", "", 3}, + {"", "abc", 3}, + {"toolcall", "toolcall", 0}, + {"toolcall", "tool_call", 1}, + {"toolcall", "tool-call", 1}, + {"toolcall", "ToolCall", 2}, + {"kitten", "sitting", 3}, + } + for _, tt := range tests { + got := levenshtein(tt.a, tt.b) + if got != tt.want { + t.Errorf("levenshtein(%q, %q) = %d, want %d", tt.a, tt.b, got, tt.want) + } + } +} + +func TestIsToolCallTag(t *testing.T) { + + for _, name := range []string{"toolcall", "tool_call", "tool-call", "ToolCall", "Toolcall", "toolCall", "TOOLCALL"} { + if !isToolCallTag(name) { + t.Errorf("isToolCallTag(%q) = false, want true", name) + } + } + + for _, name := range []string{"function_call", "FunctionCall", "functioncall", "FUNCTION_CALL"} { + if !isToolCallTag(name) { + t.Errorf("isToolCallTag(%q) = false, want true", name) + } + } + + for _, name := range []string{"tool_use", "ToolUse", "tooluse", "TOOL_USE"} { + if !isToolCallTag(name) { + t.Errorf("isToolCallTag(%q) = false, want true", name) + } + } + + for _, name := range []string{"invoke", "parameter", "function", "result", "hello", "content"} { + if isToolCallTag(name) { + t.Errorf("isToolCallTag(%q) = true, want false", name) + } + } +} diff --git a/pkg/providers/claude_cli_provider_test.go b/pkg/providers/claude_cli_provider_test.go index 11c33ab2c..d4d648f5a 100644 --- a/pkg/providers/claude_cli_provider_test.go +++ b/pkg/providers/claude_cli_provider_test.go @@ -619,12 +619,12 @@ func TestBuildSystemPrompt_WithTools(t *testing.T) { Function: ToolFunctionDefinition{ Name: "get_weather", Description: "Get weather for a location", - Parameters: MustMarshalParameters(map[string]any{ + Parameters: map[string]any{ "type": "object", "properties": map[string]any{ "location": map[string]any{"type": "string"}, }, - }), + }, }, }, } @@ -917,9 +917,9 @@ func TestExtractToolCalls_ToolCallArgumentsParsing(t *testing.T) { if got[0].Arguments["name"] != "test" { t.Errorf("Arguments[name] = %v, want test", got[0].Arguments["name"]) } - // Verify parsed arguments are also set on FunctionCall - if len(got[0].Function.Arguments) == 0 { - t.Error("Function.Arguments should contain parsed JSON arguments") + // Verify raw arguments string is preserved in FunctionCall + if got[0].Function.Arguments == "" { + t.Error("Function.Arguments should contain raw JSON string") } } @@ -984,282 +984,3 @@ func TestFindMatchingBrace(t *testing.T) { } } } - -// --- XML tool call extract/strip tests --- - -func TestExtractXMLToolCalls_Single(t *testing.T) { - text := ` - -echo hello - -` - - calls := extractXMLToolCalls(text) - if len(calls) != 1 { - t.Fatalf("expected 1 tool call, got %d", len(calls)) - } - if calls[0].Name != "exec" { - t.Errorf("Name = %q, want %q", calls[0].Name, "exec") - } - if calls[0].Arguments["command"] != "echo hello" { - t.Errorf("Arguments[command] = %v, want %q", calls[0].Arguments["command"], "echo hello") - } - if calls[0].Function == nil || calls[0].Function.Name != "exec" { - t.Errorf("Function.Name should be exec") - } -} - -func TestExtractXMLToolCalls_Multiple(t *testing.T) { - text := ` - -golang testing - - -go test ./... -30 - -` - - calls := extractXMLToolCalls(text) - if len(calls) != 2 { - t.Fatalf("expected 2 tool calls, got %d", len(calls)) - } - if calls[0].Name != "web_search" { - t.Errorf("[0].Name = %q, want %q", calls[0].Name, "web_search") - } - if calls[1].Name != "exec" { - t.Errorf("[1].Name = %q, want %q", calls[1].Name, "exec") - } - if calls[1].Arguments["timeout"] != "30" { - t.Errorf("[1].Arguments[timeout] = %v, want %q", calls[1].Arguments["timeout"], "30") - } -} - -func TestExtractXMLToolCalls_NoXML(t *testing.T) { - calls := extractXMLToolCalls("just regular text") - if len(calls) != 0 { - t.Errorf("expected 0 tool calls, got %d", len(calls)) - } -} - -func TestStripXMLToolCalls(t *testing.T) { - text := `Let me run that. - - -echo hello - - -Done.` - - got := stripXMLToolCalls(text) - if strings.Contains(got, "toolcall") { - t.Errorf("should remove XML block, got %q", got) - } - if !strings.Contains(got, "Let me run that.") { - t.Errorf("should keep text before, got %q", got) - } - if !strings.Contains(got, "Done.") { - t.Errorf("should keep text after, got %q", got) - } -} - -func TestExtractXMLToolCalls_MismatchedCloseTag(t *testing.T) { - // MiniMax uses but closes with (underscore) - text := ` - -/home/user/project/pyproject.toml - -` - - calls := extractXMLToolCalls(text) - if len(calls) != 1 { - t.Fatalf("expected 1 tool call, got %d", len(calls)) - } - if calls[0].Name != "readfile" { - t.Errorf("Name = %q, want %q", calls[0].Name, "readfile") - } - if calls[0].Arguments["path"] != "/home/user/project/pyproject.toml" { - t.Errorf("Arguments[path] = %v, want pyproject.toml path", calls[0].Arguments["path"]) - } -} - -func TestStripXMLToolCalls_MismatchedCloseTag(t *testing.T) { - text := `今テスト走らせるね。` + //nolint:gosmopolitan // CJK test data - ` - - -cd /home/user && pytest - -` - - got := stripXMLToolCalls(text) - if strings.Contains(got, "toolcall") || strings.Contains(got, "tool_call") { - t.Errorf("should remove XML block, got %q", got) - } - if !strings.Contains(got, "今テスト走らせるね。") { //nolint:gosmopolitan // CJK test data - t.Errorf("should keep text before, got %q", got) - } -} - -func TestExtractXMLToolCalls_UnderscoreOpenTag(t *testing.T) { - // Opening tag also uses underscore: - text := ` - -ls -la - -` - - calls := extractXMLToolCalls(text) - if len(calls) != 1 { - t.Fatalf("expected 1 tool call, got %d", len(calls)) - } - if calls[0].Name != "exec" { - t.Errorf("Name = %q, want %q", calls[0].Name, "exec") - } - if calls[0].Arguments["command"] != "ls -la" { - t.Errorf("Arguments[command] = %v, want %q", calls[0].Arguments["command"], "ls -la") - } -} - -func TestExtractXMLToolCalls_HyphenTag(t *testing.T) { - // Hypothetical: - text := ` - -/etc/hosts - -` - - calls := extractXMLToolCalls(text) - if len(calls) != 1 { - t.Fatalf("expected 1 tool call, got %d", len(calls)) - } - if calls[0].Name != "read_file" { - t.Errorf("Name = %q, want %q", calls[0].Name, "read_file") - } -} - -func TestStripXMLToolCalls_UnderscoreOpenTag(t *testing.T) { - text := `Here is the result. - - -ls - - -Finished.` - - got := stripXMLToolCalls(text) - if strings.Contains(got, "tool_call") || strings.Contains(got, "toolcall") { - t.Errorf("should remove XML block, got %q", got) - } - if !strings.Contains(got, "Here is the result.") { - t.Errorf("should keep text before, got %q", got) - } - if !strings.Contains(got, "Finished.") { - t.Errorf("should keep text after, got %q", got) - } -} - -func TestExtractXMLToolCalls_OrphanedClosingTag(t *testing.T) { - // LLM emits [TOOLCALL] marker + with orphaned closing tag (no opening tag) - text := "了解！確認するね。\n[TOOLCALL]\n\n/home/user/workspace\n\n" //nolint:gosmopolitan // CJK test data - - calls := extractXMLToolCalls(text) - if len(calls) != 1 { - t.Fatalf("expected 1 tool call, got %d", len(calls)) - } - if calls[0].Name != "listdir" { - t.Errorf("Name = %q, want %q", calls[0].Name, "listdir") - } - if calls[0].Arguments["path"] != "/home/user/workspace" { - t.Errorf("Arguments[path] = %v, want /home/user/workspace", calls[0].Arguments["path"]) - } -} - -func TestStripXMLToolCalls_OrphanedClosingTag(t *testing.T) { - text := "了解！確認するね。\n[TOOLCALL]\n\n/home/user\n\n" //nolint:gosmopolitan // CJK test data - got := stripXMLToolCalls(text) - if strings.Contains(got, "invoke") || strings.Contains(got, "TOOLCALL") || strings.Contains(got, "minimax") { - t.Errorf("should remove orphaned closing tag block, got %q", got) - } - if !strings.Contains(got, "了解") { //nolint:gosmopolitan // CJK test data - t.Errorf("should keep user-facing text, got %q", got) - } -} - -func TestStripXMLToolCalls_NoXML(t *testing.T) { - text := "Just regular text." - got := stripXMLToolCalls(text) - if got != text { - t.Errorf("stripXMLToolCalls() = %q, want %q", got, text) - } -} - -func TestNormalizeAlpha(t *testing.T) { - tests := []struct { - input, want string - }{ - {"toolcall", "toolcall"}, - {"tool_call", "toolcall"}, - {"Tool-Call", "toolcall"}, - {"ReadFile", "readfile"}, - {"read_file", "readfile"}, - {"EXEC", "exec"}, - {"web123search", "websearch"}, - {"", ""}, - } - for _, tt := range tests { - got := normalizeAlpha(tt.input) - if got != tt.want { - t.Errorf("normalizeAlpha(%q) = %q, want %q", tt.input, got, tt.want) - } - } -} - -func TestLevenshtein(t *testing.T) { - tests := []struct { - a, b string - want int - }{ - {"", "", 0}, - {"abc", "", 3}, - {"", "abc", 3}, - {"toolcall", "toolcall", 0}, - {"toolcall", "tool_call", 1}, - {"toolcall", "tool-call", 1}, - {"toolcall", "ToolCall", 2}, // T and C - {"kitten", "sitting", 3}, - } - for _, tt := range tests { - got := levenshtein(tt.a, tt.b) - if got != tt.want { - t.Errorf("levenshtein(%q, %q) = %d, want %d", tt.a, tt.b, got, tt.want) - } - } -} - -func TestIsToolCallTag(t *testing.T) { - // Should match — toolcall variants - for _, name := range []string{"toolcall", "tool_call", "tool-call", "ToolCall", "Toolcall", "toolCall", "TOOLCALL"} { - if !isToolCallTag(name) { - t.Errorf("isToolCallTag(%q) = false, want true", name) - } - } - // Should match — function_call variants - for _, name := range []string{"function_call", "FunctionCall", "functioncall", "FUNCTION_CALL"} { - if !isToolCallTag(name) { - t.Errorf("isToolCallTag(%q) = false, want true", name) - } - } - // Should match — tool_use variants - for _, name := range []string{"tool_use", "ToolUse", "tooluse", "TOOL_USE"} { - if !isToolCallTag(name) { - t.Errorf("isToolCallTag(%q) = false, want true", name) - } - } - // Should NOT match - for _, name := range []string{"invoke", "parameter", "function", "result", "hello", "content"} { - if isToolCallTag(name) { - t.Errorf("isToolCallTag(%q) = true, want false", name) - } - } -} diff --git a/pkg/providers/codex_cli_provider_test.go b/pkg/providers/codex_cli_provider_test.go index e537d7e31..0f66e25f4 100644 --- a/pkg/providers/codex_cli_provider_test.go +++ b/pkg/providers/codex_cli_provider_test.go @@ -76,8 +76,8 @@ func TestParseJSONLEvents_ToolCallExtraction(t *testing.T) { if resp.ToolCalls[0].ID != "call_1" { t.Errorf("ToolCalls[0].ID = %q, want %q", resp.ToolCalls[0].ID, "call_1") } - if resp.ToolCalls[0].Function.Arguments["path"] != "/tmp/test.txt" { - t.Errorf("ToolCalls[0].Function.Arguments[path] = %v", resp.ToolCalls[0].Function.Arguments["path"]) + if resp.ToolCalls[0].Function.Arguments != `{"path":"/tmp/test.txt"}` { + t.Errorf("ToolCalls[0].Function.Arguments = %q", resp.ToolCalls[0].Function.Arguments) } // Content should have the tool call JSON stripped if strings.Contains(resp.Content, "tool_calls") { @@ -292,12 +292,12 @@ func TestBuildPrompt_WithTools(t *testing.T) { Function: ToolFunctionDefinition{ Name: "get_weather", Description: "Get current weather", - Parameters: MustMarshalParameters(map[string]any{ + Parameters: map[string]any{ "type": "object", "properties": map[string]any{ "city": map[string]any{"type": "string"}, }, - }), + }, }, }, } @@ -490,7 +490,7 @@ echo '{"type":"turn.completed"}'` } messages := []Message{{Role: "user", Content: "test"}} - _, err := p.Chat(context.Background(), messages, nil, "gpt-5.2-codex", nil) + _, err := p.Chat(context.Background(), messages, nil, "gpt-5.3-codex", nil) if err != nil { t.Fatalf("Chat() error: %v", err) } @@ -502,7 +502,7 @@ echo '{"type":"turn.completed"}'` } args := string(argsData) - if !strings.Contains(args, "-m gpt-5.2-codex") { + if !strings.Contains(args, "-m gpt-5.3-codex") { t.Errorf("args should contain model flag, got: %s", args) } if !strings.Contains(args, "-C /tmp/test-workspace") { diff --git a/pkg/providers/codex_provider_test.go b/pkg/providers/codex_provider_test.go index c0db1c381..dd5ad2637 100644 --- a/pkg/providers/codex_provider_test.go +++ b/pkg/providers/codex_provider_test.go @@ -79,7 +79,7 @@ func TestBuildCodexParams_ToolCallFunctionFallback(t *testing.T) { Type: "function", Function: &FunctionCall{ Name: "read_file", - Arguments: map[string]any{"path": "README.md"}, + Arguments: `{"path":"README.md"}`, }, }, }, @@ -114,12 +114,12 @@ func TestBuildCodexParams_WithTools(t *testing.T) { Function: ToolFunctionDefinition{ Name: "get_weather", Description: "Get weather", - Parameters: MustMarshalParameters(map[string]any{ + Parameters: map[string]any{ "type": "object", "properties": map[string]any{ "city": map[string]any{"type": "string"}, }, - }), + }, }, }, } @@ -166,9 +166,9 @@ func TestBuildCodexParams_WebSearchFunctionReplacedWithBuiltin(t *testing.T) { Function: ToolFunctionDefinition{ Name: "web_search", Description: "local web search", - Parameters: MustMarshalParameters(map[string]any{ + Parameters: map[string]any{ "type": "object", - }), + }, }, }, { @@ -176,9 +176,9 @@ func TestBuildCodexParams_WebSearchFunctionReplacedWithBuiltin(t *testing.T) { Function: ToolFunctionDefinition{ Name: "read_file", Description: "read file", - Parameters: MustMarshalParameters(map[string]any{ + Parameters: map[string]any{ "type": "object", - }), + }, }, }, } @@ -568,7 +568,7 @@ func TestCodexProvider_ChatRoundTrip_ModelFallbackFromUnsupported(t *testing.T) provider.client = createOpenAITestClient(server.URL, "test-token", "acc-123") messages := []Message{{Role: "user", Content: "Hello"}} - resp, err := provider.Chat(t.Context(), messages, nil, "gpt-5.2", nil) + resp, err := provider.Chat(t.Context(), messages, nil, "gpt-5.3-codex", nil) if err != nil { t.Fatalf("Chat() error: %v", err) } @@ -599,7 +599,7 @@ func TestResolveCodexModel(t *testing.T) { wantFallback: true, }, {name: "non-openai prefixed", input: "glm-4.7", wantModel: codexDefaultModel, wantFallback: true}, - {name: "openai prefix", input: "openai/gpt-5.2", wantModel: "gpt-5.2", wantFallback: false}, + {name: "openai prefix", input: "openai/gpt-5.3-codex", wantModel: "gpt-5.3-codex", wantFallback: false}, {name: "direct gpt", input: "gpt-4o", wantModel: "gpt-4o", wantFallback: false}, } diff --git a/pkg/providers/factory_ext_test.go b/pkg/providers/factory_ext_test.go new file mode 100644 index 000000000..465968dc3 --- /dev/null +++ b/pkg/providers/factory_ext_test.go @@ -0,0 +1,73 @@ +package providers + +import ( + "github.com/sipeed/picoclaw/pkg/auth" + "github.com/sipeed/picoclaw/pkg/config" + "testing" +) + +func TestCreateProviderByName_OpenAI_OAuth(t *testing.T) { + originalGetCredential := getCredential + t.Cleanup(func() { getCredential = originalGetCredential }) + + getCredential = func(provider string) (*auth.AuthCredential, error) { + if provider != "openai" { + t.Fatalf("provider = %q, want openai", provider) + } + return &auth.AuthCredential{ + AccessToken: "openai-token", + AccountID: "acct_test", + }, nil + } + + cfg := config.DefaultConfig() + cfg.Providers.OpenAI.AuthMethod = "oauth" + + provider, err := CreateProviderByName(cfg, "openai") + if err != nil { + t.Fatalf("CreateProviderByName() error = %v", err) + } + + if _, ok := provider.(*CodexProvider); !ok { + t.Fatalf("provider type = %T, want *CodexProvider", provider) + } +} + +func TestCreateProviderByName_VLLM(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Providers.VLLM.APIKey = "test-vllm-key" + cfg.Providers.VLLM.APIBase = "https://api.example.com/v1" + + provider, err := CreateProviderByName(cfg, "vllm") + if err != nil { + t.Fatalf("CreateProviderByName() error = %v", err) + } + + if _, ok := provider.(*HTTPProvider); !ok { + t.Fatalf("provider type = %T, want *HTTPProvider", provider) + } +} + +func TestCreateProviderByName_Unknown(t *testing.T) { + cfg := config.DefaultConfig() + + _, err := CreateProviderByName(cfg, "nonexistent-provider") + if err == nil { + t.Fatal("expected error for unknown provider, got nil") + } +} + +func TestCreateProviderByName_CaseInsensitive(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Providers.VLLM.APIKey = "test-key" + cfg.Providers.VLLM.APIBase = "https://example.com/v1" + + provider, err := CreateProviderByName(cfg, "VLLM") + if err != nil { + t.Fatalf("CreateProviderByName() error = %v", err) + } + + if _, ok := provider.(*HTTPProvider); !ok { + t.Fatalf("provider type = %T, want *HTTPProvider", provider) + } +} diff --git a/pkg/providers/factory_provider_test.go b/pkg/providers/factory_provider_test.go index 17bc55d25..6c7bb4795 100644 --- a/pkg/providers/factory_provider_test.go +++ b/pkg/providers/factory_provider_test.go @@ -113,6 +113,7 @@ func TestCreateProviderFromConfig_DefaultAPIBase(t *testing.T) { {"vllm", "vllm"}, {"deepseek", "deepseek"}, {"ollama", "ollama"}, + {"longcat", "longcat"}, } for _, tt := range tests { @@ -162,6 +163,29 @@ func TestCreateProviderFromConfig_LiteLLM(t *testing.T) { } } +func TestCreateProviderFromConfig_LongCat(t *testing.T) { + cfg := &config.ModelConfig{ + ModelName: "test-longcat", + Model: "longcat/LongCat-Flash-Thinking", + APIKey: "test-key", + APIBase: "https://api.longcat.chat/openai", + } + + provider, modelID, err := CreateProviderFromConfig(cfg) + if err != nil { + t.Fatalf("CreateProviderFromConfig() error = %v", err) + } + if provider == nil { + t.Fatal("CreateProviderFromConfig() returned nil provider") + } + if modelID != "LongCat-Flash-Thinking" { + t.Errorf("modelID = %q, want %q", modelID, "LongCat-Flash-Thinking") + } + if _, ok := provider.(*HTTPProvider); !ok { + t.Fatalf("expected *HTTPProvider, got %T", provider) + } +} + func TestCreateProviderFromConfig_Anthropic(t *testing.T) { cfg := &config.ModelConfig{ ModelName: "test-anthropic", diff --git a/pkg/providers/factory_test.go b/pkg/providers/factory_test.go index 64cbc211b..91469f25b 100644 --- a/pkg/providers/factory_test.go +++ b/pkg/providers/factory_test.go @@ -178,6 +178,26 @@ func TestResolveProviderSelection(t *testing.T) { wantAPIBase: "https://api.moonshot.cn/v1", wantProxy: "http://127.0.0.1:7890", }, + { + name: "explicit longcat provider uses defaults", + setup: func(cfg *config.Config) { + cfg.Agents.Defaults.Provider = "longcat" + cfg.Providers.LongCat.APIKey = "longcat-key" + cfg.Providers.LongCat.Proxy = "http://127.0.0.1:7890" + }, + wantType: providerTypeHTTPCompat, + wantAPIBase: "https://api.longcat.chat/openai", + wantProxy: "http://127.0.0.1:7890", + }, + { + name: "longcat model fallback uses longcat base default", + setup: func(cfg *config.Config) { + cfg.Agents.Defaults.Model = "longcat/LongCat-Flash-Thinking" + cfg.Providers.LongCat.APIKey = "longcat-key" + }, + wantType: providerTypeHTTPCompat, + wantAPIBase: "https://api.longcat.chat/openai", + }, { name: "missing keys returns model config error", setup: func(cfg *config.Config) { @@ -329,69 +349,3 @@ func TestCreateProviderReturnsCodexProviderForOpenAIOAuth(t *testing.T) { // which is not yet implemented in the new factory_provider.go t.Skip("OpenAI OAuth via model_list not yet implemented") } - -func TestCreateProviderByName_OpenAI_OAuth(t *testing.T) { - originalGetCredential := getCredential - t.Cleanup(func() { getCredential = originalGetCredential }) - - getCredential = func(provider string) (*auth.AuthCredential, error) { - if provider != "openai" { - t.Fatalf("provider = %q, want openai", provider) - } - return &auth.AuthCredential{ - AccessToken: "openai-token", - AccountID: "acct_test", - }, nil - } - - cfg := config.DefaultConfig() - cfg.Providers.OpenAI.AuthMethod = "oauth" - - provider, err := CreateProviderByName(cfg, "openai") - if err != nil { - t.Fatalf("CreateProviderByName() error = %v", err) - } - - if _, ok := provider.(*CodexProvider); !ok { - t.Fatalf("provider type = %T, want *CodexProvider", provider) - } -} - -func TestCreateProviderByName_VLLM(t *testing.T) { - cfg := config.DefaultConfig() - cfg.Providers.VLLM.APIKey = "test-vllm-key" - cfg.Providers.VLLM.APIBase = "https://api.example.com/v1" - - provider, err := CreateProviderByName(cfg, "vllm") - if err != nil { - t.Fatalf("CreateProviderByName() error = %v", err) - } - - if _, ok := provider.(*HTTPProvider); !ok { - t.Fatalf("provider type = %T, want *HTTPProvider", provider) - } -} - -func TestCreateProviderByName_Unknown(t *testing.T) { - cfg := config.DefaultConfig() - - _, err := CreateProviderByName(cfg, "nonexistent-provider") - if err == nil { - t.Fatal("expected error for unknown provider, got nil") - } -} - -func TestCreateProviderByName_CaseInsensitive(t *testing.T) { - cfg := config.DefaultConfig() - cfg.Providers.VLLM.APIKey = "test-key" - cfg.Providers.VLLM.APIBase = "https://example.com/v1" - - provider, err := CreateProviderByName(cfg, "VLLM") - if err != nil { - t.Fatalf("CreateProviderByName() error = %v", err) - } - - if _, ok := provider.(*HTTPProvider); !ok { - t.Fatalf("provider type = %T, want *HTTPProvider", provider) - } -} diff --git a/pkg/providers/openai_compat/provider_ext_test.go b/pkg/providers/openai_compat/provider_ext_test.go new file mode 100644 index 000000000..2c25422db --- /dev/null +++ b/pkg/providers/openai_compat/provider_ext_test.go @@ -0,0 +1,456 @@ +package openai_compat + +import ( + "context" + "encoding/json" + "fmt" + "github.com/sipeed/picoclaw/pkg/providers/protocoltypes" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestProviderChat_StripsGroqAndOllamaPrefixes(t *testing.T) { + tests := []struct { + name string + input string + wantModel string + }{ + { + name: "strips groq prefix and keeps nested model", + input: "groq/openai/gpt-oss-120b", + wantModel: "openai/gpt-oss-120b", + }, + { + name: "strips ollama prefix", + input: "ollama/qwen2.5:14b", + wantModel: "qwen2.5:14b", + }, + { + name: "strips deepseek prefix", + input: "deepseek/deepseek-chat", + wantModel: "deepseek-chat", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var requestBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "ok"}, + "finish_reason": "stop", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, tt.input, nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + if requestBody["model"] != tt.wantModel { + t.Fatalf("model = %v, want %s", requestBody["model"], tt.wantModel) + } + }) + } +} + +func TestNormalizeModel_OpenAIPrefix(t *testing.T) { + if got := normalizeModel("openai/gpt-5.2", "https://api.openai.com/v1"); got != "gpt-5.2" { + t.Fatalf("normalizeModel(openai/gpt-5.2) = %q, want %q", got, "gpt-5.2") + } +} + +func TestProviderChat_StreamingTextResponse(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/text/chatcompletion_v2" { + http.Error(w, "not found", http.StatusNotFound) + return + } + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + if body["stream"] != true { + t.Error("expected stream=true in request body") + } + + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + chunks := []string{ + `data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"content":" world"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`, + `data: [DONE]`, + } + for _, c := range chunks { + fmt.Fprintln(w, c) + fmt.Fprintln(w) + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProvider("key", server.URL, "", + WithEndpointPath("/text/chatcompletion_v2"), + WithStream(true), + ) + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "MiniMax-M1", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if out.Content != "Hello world" { + t.Fatalf("Content = %q, want %q", out.Content, "Hello world") + } + if out.FinishReason != "stop" { + t.Fatalf("FinishReason = %q, want %q", out.FinishReason, "stop") + } + if out.Usage == nil || out.Usage.TotalTokens != 7 { + t.Fatalf("Usage.TotalTokens = %v, want 7", out.Usage) + } +} + +func TestProviderChat_StreamingToolCalls(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + chunks := []string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":"}}]},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"SF\"}"}}]},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":8,"total_tokens":18}}`, + `data: [DONE]`, + } + for _, c := range chunks { + fmt.Fprintln(w, c) + fmt.Fprintln(w) + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProvider("key", server.URL, "", WithStream(true)) + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "test", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if len(out.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls)) + } + tc := out.ToolCalls[0] + if tc.ID != "call_1" { + t.Fatalf("ToolCalls[0].ID = %q, want %q", tc.ID, "call_1") + } + if tc.Name != "get_weather" { + t.Fatalf("ToolCalls[0].Name = %q, want %q", tc.Name, "get_weather") + } + if tc.Arguments["city"] != "SF" { + t.Fatalf("ToolCalls[0].Arguments[city] = %v, want SF", tc.Arguments["city"]) + } +} + +func TestProviderChat_CustomEndpointPath(t *testing.T) { + var hitPath string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hitPath = r.URL.Path + resp := map[string]any{ + "choices": []map[string]any{ + {"message": map[string]any{"content": "ok"}, "finish_reason": "stop"}, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "", + WithEndpointPath("/text/chatcompletion_v2"), + ) + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if hitPath != "/text/chatcompletion_v2" { + t.Fatalf("endpoint path = %q, want %q", hitPath, "/text/chatcompletion_v2") + } +} + +func TestReadSSEIntoChannel_TextAndToolCalls(t *testing.T) { + sseData := strings.Join([]string{ + `data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{"content":" world"},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"greet","arguments":"{\"n"}}]},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ame\":\"Bob\"}"}}]},"finish_reason":""}]}`, + ``, + `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + + ch := make(chan protocoltypes.StreamEvent, 32) + go func() { + defer close(ch) + readSSEIntoChannel(context.Background(), strings.NewReader(sseData), ch) + }() + + var events []protocoltypes.StreamEvent + for ev := range ch { + events = append(events, ev) + } + + if len(events) < 3 { + t.Fatalf("got %d events, want at least 3", len(events)) + } + + if events[0].ContentDelta != "Hello" { + t.Errorf("events[0].ContentDelta = %q, want %q", events[0].ContentDelta, "Hello") + } + if events[1].ContentDelta != " world" { + t.Errorf("events[1].ContentDelta = %q, want %q", events[1].ContentDelta, " world") + } + + if len(events[2].ToolCallDeltas) != 1 || events[2].ToolCallDeltas[0].ID != "call_1" { + t.Errorf("events[2] should contain tool call with ID=call_1") + } + if events[2].ToolCallDeltas[0].Name != "greet" { + t.Errorf("events[2].ToolCallDeltas[0].Name = %q, want %q", events[2].ToolCallDeltas[0].Name, "greet") + } + + lastEv := events[len(events)-1] + if lastEv.FinishReason != "stop" { + t.Errorf("last event FinishReason = %q, want %q", lastEv.FinishReason, "stop") + } + if lastEv.Usage == nil || lastEv.Usage.TotalTokens != 7 { + t.Errorf("last event Usage.TotalTokens = %v, want 7", lastEv.Usage) + } +} + +func TestReadSSEIntoChannel_ContextCancel(t *testing.T) { + + ctx, cancel := context.WithCancel(context.Background()) + + sseData := `data: {"choices":[{"delta":{"content":"first"},"finish_reason":""}]}` + "\n\n" + + ch := make(chan protocoltypes.StreamEvent, 32) + go func() { + defer close(ch) + readSSEIntoChannel(ctx, strings.NewReader(sseData), ch) + }() + + ev := <-ch + if ev.ContentDelta != "first" { + t.Fatalf("ContentDelta = %q, want %q", ev.ContentDelta, "first") + } + + cancel() + _, ok := <-ch + if ok { + t.Fatal("expected channel to be closed after context cancel") + } +} + +func TestAccumulateStream_FullResponse(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "Hello"} + ch <- protocoltypes.StreamEvent{ContentDelta: " world"} + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ID: "call_1", Name: "test_tool", ArgumentsDelta: `{"key"`}, + }, + } + ch <- protocoltypes.StreamEvent{ + ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ + {Index: 0, ArgumentsDelta: `:"value"}`}, + }, + } + ch <- protocoltypes.StreamEvent{ + FinishReason: "stop", + Usage: &UsageInfo{PromptTokens: 5, CompletionTokens: 3, TotalTokens: 8}, + } + close(ch) + }() + + resp, err := AccumulateStream(ch) + if err != nil { + t.Fatalf("AccumulateStream() error = %v", err) + } + + if resp.Content != "Hello world" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world") + } + if resp.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") + } + if resp.Usage == nil || resp.Usage.TotalTokens != 8 { + t.Errorf("Usage.TotalTokens = %v, want 8", resp.Usage) + } + if len(resp.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls)) + } + if resp.ToolCalls[0].Name != "test_tool" { + t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_tool") + } + if resp.ToolCalls[0].Arguments["key"] != "value" { + t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "value") + } +} + +func TestAccumulateStream_Error(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 4) + + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "partial"} + ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("connection reset")} + close(ch) + }() + + _, err := AccumulateStream(ch) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "connection reset") { + t.Fatalf("error = %q, want to contain %q", err.Error(), "connection reset") + } +} + +func TestChatStream_EndToEnd(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + chunks := []string{ + `data: {"choices":[{"delta":{"content":"stream"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{"content":"ed"},"finish_reason":""}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}`, + `data: [DONE]`, + } + for _, c := range chunks { + fmt.Fprintln(w, c) + fmt.Fprintln(w) + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProvider("key", server.URL, "", WithStream(true)) + + ch, err := p.ChatStream(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) + if err != nil { + t.Fatalf("ChatStream() error = %v", err) + } + + resp, err := AccumulateStream(ch) + if err != nil { + t.Fatalf("AccumulateStream() error = %v", err) + } + + if resp.Content != "streamed" { + t.Errorf("Content = %q, want %q", resp.Content, "streamed") + } + if resp.FinishReason != "stop" { + t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") + } + if resp.Usage == nil || resp.Usage.TotalTokens != 3 { + t.Errorf("Usage.TotalTokens = %v, want 3", resp.Usage) + } +} + +func TestChatStream_EarlyCancel(t *testing.T) { + serverDone := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + defer close(serverDone) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + + for i := 0; i < 1000; i++ { + select { + case <-r.Context().Done(): + return + default: + } + fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"x\"},\"finish_reason\":\"\"}]}\n\n") + if flusher != nil { + flusher.Flush() + } + } + })) + defer server.Close() + + p := NewProvider("key", server.URL, "", WithStream(true)) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch, err := p.ChatStream(ctx, []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) + if err != nil { + t.Fatalf("ChatStream() error = %v", err) + } + + count := 0 + for ev := range ch { + if ev.Err != nil { + break + } + count++ + if count >= 5 { + cancel() + } + } + + if count < 5 { + t.Errorf("expected at least 5 events before cancel, got %d", count) + } + + <-serverDone +} + +func TestCanStream(t *testing.T) { + p1 := NewProvider("key", "https://example.com", "") + if p1.CanStream() { + t.Error("CanStream() = true for non-stream provider") + } + + p2 := NewProvider("key", "https://example.com", "", WithStream(true)) + if !p2.CanStream() { + t.Error("CanStream() = false for stream provider") + } +} + +func TestProvider_RequestTimeoutNonPositive(t *testing.T) { + p := NewProviderWithMaxTokensFieldAndTimeout("key", "https://example.com/v1", "", "", -1) + if p.httpClient.Timeout != defaultRequestTimeout { + t.Fatalf("http timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout) + } +} diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index e77d48725..41f278a1b 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -1,9 +1,10 @@ package openai_compat import ( - "context" + "bytes" "encoding/json" "fmt" + "io" "net/http" "net/http/httptest" "net/url" @@ -107,6 +108,55 @@ func TestProviderChat_ParsesToolCalls(t *testing.T) { } } +func TestProviderChat_ParsesToolCallsWithObjectArguments(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{ + "content": "", + "tool_calls": []map[string]any{ + { + "id": "call_1", + "type": "function", + "function": map[string]any{ + "name": "get_weather", + "arguments": map[string]any{ + "city": "SF", + "metric": true, + }, + }, + }, + }, + }, + "finish_reason": "tool_calls", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "gpt-4o", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if len(out.ToolCalls) != 1 { + t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls)) + } + if out.ToolCalls[0].Name != "get_weather" { + t.Fatalf("ToolCalls[0].Name = %q, want %q", out.ToolCalls[0].Name, "get_weather") + } + if out.ToolCalls[0].Arguments["city"] != "SF" { + t.Fatalf("ToolCalls[0].Arguments[city] = %v, want SF", out.ToolCalls[0].Arguments["city"]) + } + if out.ToolCalls[0].Arguments["metric"] != true { + t.Fatalf("ToolCalls[0].Arguments[metric] = %v, want true", out.ToolCalls[0].Arguments["metric"]) + } +} + func TestProviderChat_ParsesReasoningContent(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { resp := map[string]any{ @@ -151,6 +201,56 @@ func TestProviderChat_ParsesReasoningContent(t *testing.T) { } } +func TestProviderChat_PreservesReasoningContentInHistory(t *testing.T) { + var requestBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "ok"}, + "finish_reason": "stop", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + + // Simulate a multi-turn conversation where the assistant's previous + // reply included reasoning_content (e.g. from kimi-k2.5). + messages := []Message{ + {Role: "user", Content: "What is 1+1?"}, + {Role: "assistant", Content: "2", ReasoningContent: "Let me think... 1+1=2"}, + {Role: "user", Content: "What about 2+2?"}, + } + + _, err := p.Chat(t.Context(), messages, nil, "kimi-k2.5", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + + // Verify reasoning_content is preserved in the serialized request. + reqMessages, ok := requestBody["messages"].([]any) + if !ok { + t.Fatalf("messages is not []any: %T", requestBody["messages"]) + } + assistantMsg, ok := reqMessages[1].(map[string]any) + if !ok { + t.Fatalf("assistant message is not map[string]any: %T", reqMessages[1]) + } + if assistantMsg["reasoning_content"] != "Let me think... 1+1=2" { + t.Errorf("reasoning_content not preserved in request, got %v", assistantMsg["reasoning_content"]) + } +} + func TestProviderChat_HTTPError(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "bad request", http.StatusBadRequest) @@ -164,6 +264,132 @@ func TestProviderChat_HTTPError(t *testing.T) { } } +func TestProviderChat_JSONHTTPErrorDoesNotReportHTML(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"error":"bad request"}`)) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "gpt-4o", nil) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "Status: 400") { + t.Fatalf("expected status code in error, got %v", err) + } + if strings.Contains(err.Error(), "returned HTML instead of JSON") { + t.Fatalf("expected non-HTML http error, got %v", err) + } +} + +func TestProviderChat_HTMLResponsesReturnHelpfulError(t *testing.T) { + tests := []struct { + name string + contentType string + statusCode int + body string + }{ + { + name: "html success response", + contentType: "text/html; charset=utf-8", + statusCode: http.StatusOK, + body: "gateway login", + }, + { + name: "html error response", + contentType: "text/html; charset=utf-8", + statusCode: http.StatusBadGateway, + body: "bad gateway", + }, + { + name: "mislabeled html success response", + contentType: "application/json", + statusCode: http.StatusOK, + body: " \r\n\tgateway login", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", tt.contentType) + w.WriteHeader(tt.statusCode) + _, _ = w.Write([]byte(tt.body)) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "gpt-4o", nil) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), fmt.Sprintf("Status: %d", tt.statusCode)) { + t.Fatalf("expected status code in error, got %v", err) + } + if !strings.Contains(err.Error(), "returned HTML instead of JSON") { + t.Fatalf("expected helpful HTML error, got %v", err) + } + if !strings.Contains(err.Error(), "check api_base or proxy configuration") { + t.Fatalf("expected configuration hint, got %v", err) + } + }) + } +} + +func TestProviderChat_SuccessResponseUsesStreamingDecoder(t *testing.T) { + content := strings.Repeat("a", 1024) + body := `{"choices":[{"message":{"content":"` + content + `"},"finish_reason":"stop"}]}` + + p := NewProvider("key", "https://example.com/v1", "") + p.httpClient = &http.Client{ + Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &errAfterDataReadCloser{ + data: []byte(body), + chunkSize: 64, + }, + }, nil + }), + } + + out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "gpt-4o", nil) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + if out.Content != content { + t.Fatalf("Content = %q, want %q", out.Content, content) + } +} + +func TestProviderChat_LargeHTMLResponsePreviewIsTruncated(t *testing.T) { + body := append([]byte(""), bytes.Repeat([]byte("A"), 2048)...) + body = append(body, []byte("")...) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + w.WriteHeader(http.StatusBadGateway) + _, _ = w.Write(body) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "gpt-4o", nil) + if err == nil { + t.Fatal("expected error, got nil") + } + if !strings.Contains(err.Error(), "Body: ") { + t.Fatalf("expected html preview in error, got %v", err) + } + if !strings.Contains(err.Error(), "...") { + t.Fatalf("expected truncated preview, got %v", err) + } +} + func TestProviderChat_StripsMoonshotPrefixAndNormalizesKimiTemperature(t *testing.T) { var requestBody map[string]any @@ -205,12 +431,17 @@ func TestProviderChat_StripsMoonshotPrefixAndNormalizesKimiTemperature(t *testin } } -func TestProviderChat_StripsGroqAndOllamaPrefixes(t *testing.T) { +func TestProviderChat_StripsGroqOllamaDeepseekVivgridPrefixes(t *testing.T) { tests := []struct { name string input string wantModel string }{ + { + name: "strips litellm prefix and preserves proxy model name", + input: "litellm/my-proxy-alias", + wantModel: "my-proxy-alias", + }, { name: "strips groq prefix and keeps nested model", input: "groq/openai/gpt-oss-120b", @@ -226,6 +457,11 @@ func TestProviderChat_StripsGroqAndOllamaPrefixes(t *testing.T) { input: "deepseek/deepseek-chat", wantModel: "deepseek-chat", }, + { + name: "strips vivgrid prefix", + input: "vivgrid/auto", + wantModel: "auto", + }, } for _, tt := range tests { @@ -330,393 +566,11 @@ func TestNormalizeModel_UsesAPIBase(t *testing.T) { if got := normalizeModel("openrouter/auto", "https://openrouter.ai/api/v1"); got != "openrouter/auto" { t.Fatalf("normalizeModel(openrouter) = %q, want %q", got, "openrouter/auto") } -} - -func TestNormalizeModel_OpenAIPrefix(t *testing.T) { - if got := normalizeModel("openai/gpt-5.2", "https://api.openai.com/v1"); got != "gpt-5.2" { - t.Fatalf("normalizeModel(openai/gpt-5.2) = %q, want %q", got, "gpt-5.2") + if got := normalizeModel("vivgrid/managed", "https://api.vivgrid.com/v1"); got != "managed" { + t.Fatalf("normalizeModel(vivgrid) = %q, want %q", got, "managed") } -} - -func TestProviderChat_StreamingTextResponse(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/text/chatcompletion_v2" { - http.Error(w, "not found", http.StatusNotFound) - return - } - var body map[string]any - if err := json.NewDecoder(r.Body).Decode(&body); err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - if body["stream"] != true { - t.Error("expected stream=true in request body") - } - - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - - chunks := []string{ - `data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":""}]}`, - `data: {"choices":[{"delta":{"content":" world"},"finish_reason":""}]}`, - `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`, - `data: [DONE]`, - } - for _, c := range chunks { - fmt.Fprintln(w, c) - fmt.Fprintln(w) // blank line between events - if flusher != nil { - flusher.Flush() - } - } - })) - defer server.Close() - - p := NewProvider("key", server.URL, "", - WithEndpointPath("/text/chatcompletion_v2"), - WithStream(true), - ) - out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "MiniMax-M1", nil) - if err != nil { - t.Fatalf("Chat() error = %v", err) - } - if out.Content != "Hello world" { - t.Fatalf("Content = %q, want %q", out.Content, "Hello world") - } - if out.FinishReason != "stop" { - t.Fatalf("FinishReason = %q, want %q", out.FinishReason, "stop") - } - if out.Usage == nil || out.Usage.TotalTokens != 7 { - t.Fatalf("Usage.TotalTokens = %v, want 7", out.Usage) - } -} - -func TestProviderChat_StreamingToolCalls(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - - chunks := []string{ - `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":""}]}`, - `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"city\":"}}]},"finish_reason":""}]}`, - `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"SF\"}"}}]},"finish_reason":""}]}`, - `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":8,"total_tokens":18}}`, - `data: [DONE]`, - } - for _, c := range chunks { - fmt.Fprintln(w, c) - fmt.Fprintln(w) - if flusher != nil { - flusher.Flush() - } - } - })) - defer server.Close() - - p := NewProvider("key", server.URL, "", WithStream(true)) - out, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "weather?"}}, nil, "test", nil) - if err != nil { - t.Fatalf("Chat() error = %v", err) - } - if len(out.ToolCalls) != 1 { - t.Fatalf("len(ToolCalls) = %d, want 1", len(out.ToolCalls)) - } - tc := out.ToolCalls[0] - if tc.ID != "call_1" { - t.Fatalf("ToolCalls[0].ID = %q, want %q", tc.ID, "call_1") - } - if tc.Name != "get_weather" { - t.Fatalf("ToolCalls[0].Name = %q, want %q", tc.Name, "get_weather") - } - if tc.Arguments["city"] != "SF" { - t.Fatalf("ToolCalls[0].Arguments[city] = %v, want SF", tc.Arguments["city"]) - } -} - -func TestProviderChat_CustomEndpointPath(t *testing.T) { - var hitPath string - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - hitPath = r.URL.Path - resp := map[string]any{ - "choices": []map[string]any{ - {"message": map[string]any{"content": "ok"}, "finish_reason": "stop"}, - }, - } - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(resp) - })) - defer server.Close() - - p := NewProvider("key", server.URL, "", - WithEndpointPath("/text/chatcompletion_v2"), - ) - _, err := p.Chat(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) - if err != nil { - t.Fatalf("Chat() error = %v", err) - } - if hitPath != "/text/chatcompletion_v2" { - t.Fatalf("endpoint path = %q, want %q", hitPath, "/text/chatcompletion_v2") - } -} - -func TestReadSSEIntoChannel_TextAndToolCalls(t *testing.T) { - sseData := strings.Join([]string{ - `data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":""}]}`, - ``, - `data: {"choices":[{"delta":{"content":" world"},"finish_reason":""}]}`, - ``, - `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"greet","arguments":"{\"n"}}]},"finish_reason":""}]}`, - ``, - `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ame\":\"Bob\"}"}}]},"finish_reason":""}]}`, - ``, - `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":4,"total_tokens":7}}`, - ``, - `data: [DONE]`, - ``, - }, "\n") - - ch := make(chan protocoltypes.StreamEvent, 32) - go func() { - defer close(ch) - readSSEIntoChannel(context.Background(), strings.NewReader(sseData), ch) - }() - - var events []protocoltypes.StreamEvent - for ev := range ch { - events = append(events, ev) - } - - if len(events) < 3 { - t.Fatalf("got %d events, want at least 3", len(events)) - } - - // Check content deltas - if events[0].ContentDelta != "Hello" { - t.Errorf("events[0].ContentDelta = %q, want %q", events[0].ContentDelta, "Hello") - } - if events[1].ContentDelta != " world" { - t.Errorf("events[1].ContentDelta = %q, want %q", events[1].ContentDelta, " world") - } - - // Check tool call deltas - if len(events[2].ToolCallDeltas) != 1 || events[2].ToolCallDeltas[0].ID != "call_1" { - t.Errorf("events[2] should contain tool call with ID=call_1") - } - if events[2].ToolCallDeltas[0].Name != "greet" { - t.Errorf("events[2].ToolCallDeltas[0].Name = %q, want %q", events[2].ToolCallDeltas[0].Name, "greet") - } - - // Check finish event - lastEv := events[len(events)-1] - if lastEv.FinishReason != "stop" { - t.Errorf("last event FinishReason = %q, want %q", lastEv.FinishReason, "stop") - } - if lastEv.Usage == nil || lastEv.Usage.TotalTokens != 7 { - t.Errorf("last event Usage.TotalTokens = %v, want 7", lastEv.Usage) - } -} - -func TestReadSSEIntoChannel_ContextCancel(t *testing.T) { - // Simulate a slow SSE stream that gets canceled. - ctx, cancel := context.WithCancel(context.Background()) - - // Create a reader that blocks after sending one chunk. - sseData := `data: {"choices":[{"delta":{"content":"first"},"finish_reason":""}]}` + "\n\n" - - ch := make(chan protocoltypes.StreamEvent, 32) - go func() { - defer close(ch) - readSSEIntoChannel(ctx, strings.NewReader(sseData), ch) - }() - - // Read the first event. - ev := <-ch - if ev.ContentDelta != "first" { - t.Fatalf("ContentDelta = %q, want %q", ev.ContentDelta, "first") - } - - // Cancel the context; the channel should close. - cancel() - _, ok := <-ch - if ok { - t.Fatal("expected channel to be closed after context cancel") - } -} - -func TestAccumulateStream_FullResponse(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 8) - - go func() { - ch <- protocoltypes.StreamEvent{ContentDelta: "Hello"} - ch <- protocoltypes.StreamEvent{ContentDelta: " world"} - ch <- protocoltypes.StreamEvent{ - ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ - {Index: 0, ID: "call_1", Name: "test_tool", ArgumentsDelta: `{"key"`}, - }, - } - ch <- protocoltypes.StreamEvent{ - ToolCallDeltas: []protocoltypes.StreamToolCallDelta{ - {Index: 0, ArgumentsDelta: `:"value"}`}, - }, - } - ch <- protocoltypes.StreamEvent{ - FinishReason: "stop", - Usage: &UsageInfo{PromptTokens: 5, CompletionTokens: 3, TotalTokens: 8}, - } - close(ch) - }() - - resp, err := AccumulateStream(ch) - if err != nil { - t.Fatalf("AccumulateStream() error = %v", err) - } - - if resp.Content != "Hello world" { - t.Errorf("Content = %q, want %q", resp.Content, "Hello world") - } - if resp.FinishReason != "stop" { - t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") - } - if resp.Usage == nil || resp.Usage.TotalTokens != 8 { - t.Errorf("Usage.TotalTokens = %v, want 8", resp.Usage) - } - if len(resp.ToolCalls) != 1 { - t.Fatalf("len(ToolCalls) = %d, want 1", len(resp.ToolCalls)) - } - if resp.ToolCalls[0].Name != "test_tool" { - t.Errorf("ToolCalls[0].Name = %q, want %q", resp.ToolCalls[0].Name, "test_tool") - } - if resp.ToolCalls[0].Arguments["key"] != "value" { - t.Errorf("ToolCalls[0].Arguments[key] = %v, want %q", resp.ToolCalls[0].Arguments["key"], "value") - } -} - -func TestAccumulateStream_Error(t *testing.T) { - ch := make(chan protocoltypes.StreamEvent, 4) - - go func() { - ch <- protocoltypes.StreamEvent{ContentDelta: "partial"} - ch <- protocoltypes.StreamEvent{Err: fmt.Errorf("connection reset")} - close(ch) - }() - - _, err := AccumulateStream(ch) - if err == nil { - t.Fatal("expected error, got nil") - } - if !strings.Contains(err.Error(), "connection reset") { - t.Fatalf("error = %q, want to contain %q", err.Error(), "connection reset") - } -} - -func TestChatStream_EndToEnd(t *testing.T) { - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - - chunks := []string{ - `data: {"choices":[{"delta":{"content":"stream"},"finish_reason":""}]}`, - `data: {"choices":[{"delta":{"content":"ed"},"finish_reason":""}]}`, - `data: {"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}`, - `data: [DONE]`, - } - for _, c := range chunks { - fmt.Fprintln(w, c) - fmt.Fprintln(w) - if flusher != nil { - flusher.Flush() - } - } - })) - defer server.Close() - - p := NewProvider("key", server.URL, "", WithStream(true)) - - ch, err := p.ChatStream(t.Context(), []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) - if err != nil { - t.Fatalf("ChatStream() error = %v", err) - } - - resp, err := AccumulateStream(ch) - if err != nil { - t.Fatalf("AccumulateStream() error = %v", err) - } - - if resp.Content != "streamed" { - t.Errorf("Content = %q, want %q", resp.Content, "streamed") - } - if resp.FinishReason != "stop" { - t.Errorf("FinishReason = %q, want %q", resp.FinishReason, "stop") - } - if resp.Usage == nil || resp.Usage.TotalTokens != 3 { - t.Errorf("Usage.TotalTokens = %v, want 3", resp.Usage) - } -} - -func TestChatStream_EarlyCancel(t *testing.T) { - serverDone := make(chan struct{}) - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - defer close(serverDone) - w.Header().Set("Content-Type", "text/event-stream") - w.WriteHeader(http.StatusOK) - flusher, _ := w.(http.Flusher) - - // Send many chunks; expect the client to cancel early. - for i := 0; i < 1000; i++ { - select { - case <-r.Context().Done(): - return - default: - } - fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"x\"},\"finish_reason\":\"\"}]}\n\n") - if flusher != nil { - flusher.Flush() - } - } - })) - defer server.Close() - - p := NewProvider("key", server.URL, "", WithStream(true)) - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - ch, err := p.ChatStream(ctx, []Message{{Role: "user", Content: "hi"}}, nil, "test", nil) - if err != nil { - t.Fatalf("ChatStream() error = %v", err) - } - - // Read a few events, then cancel. - count := 0 - for ev := range ch { - if ev.Err != nil { - break - } - count++ - if count >= 5 { - cancel() - } - } - - if count < 5 { - t.Errorf("expected at least 5 events before cancel, got %d", count) - } - - // Server should have received the cancellation. - <-serverDone -} - -func TestCanStream(t *testing.T) { - p1 := NewProvider("key", "https://example.com", "") - if p1.CanStream() { - t.Error("CanStream() = true for non-stream provider") - } - - p2 := NewProvider("key", "https://example.com", "", WithStream(true)) - if !p2.CanStream() { - t.Error("CanStream() = false for stream provider") + if got := normalizeModel("vivgrid/auto", "https://api.vivgrid.com/v1"); got != "auto" { + t.Fatalf("normalizeModel(vivgrid auto) = %q, want %q", got, "auto") } } @@ -734,11 +588,38 @@ func TestProvider_RequestTimeoutOverride(t *testing.T) { } } -func TestProvider_RequestTimeoutNonPositive(t *testing.T) { - p := NewProviderWithMaxTokensFieldAndTimeout("key", "https://example.com/v1", "", "", -1) - if p.httpClient.Timeout != defaultRequestTimeout { - t.Fatalf("http timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout) +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { + return f(r) +} + +type errAfterDataReadCloser struct { + data []byte + chunkSize int + offset int +} + +func (r *errAfterDataReadCloser) Read(p []byte) (int, error) { + if r.offset >= len(r.data) { + return 0, io.ErrUnexpectedEOF } + + n := r.chunkSize + if n <= 0 || n > len(p) { + n = len(p) + } + remaining := len(r.data) - r.offset + if n > remaining { + n = remaining + } + copy(p, r.data[r.offset:r.offset+n]) + r.offset += n + return n, nil +} + +func (r *errAfterDataReadCloser) Close() error { + return nil } func TestProvider_FunctionalOptionMaxTokensField(t *testing.T) { @@ -761,3 +642,202 @@ func TestProvider_FunctionalOptionRequestTimeoutNonPositive(t *testing.T) { t.Fatalf("http timeout = %v, want %v", p.httpClient.Timeout, defaultRequestTimeout) } } + +func TestSerializeMessages_PlainText(t *testing.T) { + messages := []protocoltypes.Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi", ReasoningContent: "thinking..."}, + } + result := serializeMessages(messages) + + data, err := json.Marshal(result) + if err != nil { + t.Fatal(err) + } + + var msgs []map[string]any + json.Unmarshal(data, &msgs) + + if msgs[0]["content"] != "hello" { + t.Fatalf("expected plain string content, got %v", msgs[0]["content"]) + } + if msgs[1]["reasoning_content"] != "thinking..." { + t.Fatalf("reasoning_content not preserved, got %v", msgs[1]["reasoning_content"]) + } +} + +func TestSerializeMessages_WithMedia(t *testing.T) { + messages := []protocoltypes.Message{ + {Role: "user", Content: "describe this", Media: []string{"data:image/png;base64,abc123"}}, + } + result := serializeMessages(messages) + + data, _ := json.Marshal(result) + var msgs []map[string]any + json.Unmarshal(data, &msgs) + + content, ok := msgs[0]["content"].([]any) + if !ok { + t.Fatalf("expected array content for media message, got %T", msgs[0]["content"]) + } + if len(content) != 2 { + t.Fatalf("expected 2 content parts, got %d", len(content)) + } + + textPart := content[0].(map[string]any) + if textPart["type"] != "text" || textPart["text"] != "describe this" { + t.Fatalf("text part mismatch: %v", textPart) + } + + imgPart := content[1].(map[string]any) + if imgPart["type"] != "image_url" { + t.Fatalf("expected image_url type, got %v", imgPart["type"]) + } + imgURL := imgPart["image_url"].(map[string]any) + if imgURL["url"] != "data:image/png;base64,abc123" { + t.Fatalf("image url mismatch: %v", imgURL["url"]) + } +} + +func TestSerializeMessages_MediaWithToolCallID(t *testing.T) { + messages := []protocoltypes.Message{ + {Role: "tool", Content: "image result", Media: []string{"data:image/png;base64,xyz"}, ToolCallID: "call_1"}, + } + result := serializeMessages(messages) + + data, _ := json.Marshal(result) + var msgs []map[string]any + json.Unmarshal(data, &msgs) + + if msgs[0]["tool_call_id"] != "call_1" { + t.Fatalf("tool_call_id not preserved with media, got %v", msgs[0]["tool_call_id"]) + } + // Content should be multipart array + if _, ok := msgs[0]["content"].([]any); !ok { + t.Fatalf("expected array content, got %T", msgs[0]["content"]) + } +} + +// chatWithCacheKey sets up a test server, sends a Chat request with prompt_cache_key, +// and returns the decoded request body for assertion. +func chatWithCacheKey(t *testing.T, apiBase string) map[string]any { + t.Helper() + var requestBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + resp := map[string]any{ + "choices": []map[string]any{ + { + "message": map[string]any{"content": "ok"}, + "finish_reason": "stop", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + p.apiBase = apiBase + p.httpClient = &http.Client{ + Transport: roundTripperFunc(func(r *http.Request) (*http.Response, error) { + r.URL, _ = url.Parse(server.URL + r.URL.Path) + return http.DefaultTransport.RoundTrip(r) + }), + } + + _, err := p.Chat( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "test-model", + map[string]any{"prompt_cache_key": "agent-main"}, + ) + if err != nil { + t.Fatalf("Chat() error = %v", err) + } + return requestBody +} + +func TestProviderChat_PromptCacheKeySentToOpenAI(t *testing.T) { + body := chatWithCacheKey(t, "https://api.openai.com/v1") + if body["prompt_cache_key"] != "agent-main" { + t.Fatalf("prompt_cache_key = %v, want %q", body["prompt_cache_key"], "agent-main") + } +} + +func TestProviderChat_PromptCacheKeyOmittedForNonOpenAI(t *testing.T) { + tests := []struct { + name string + apiBase string + }{ + {"mistral", "https://api.mistral.ai/v1"}, + {"gemini", "https://generativelanguage.googleapis.com/v1beta"}, + {"deepseek", "https://api.deepseek.com/v1"}, + {"groq", "https://api.groq.com/openai/v1"}, + {"minimax", "https://api.minimaxi.com/v1"}, + {"ollama_local", "http://localhost:11434/v1"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := chatWithCacheKey(t, tt.apiBase) + if _, exists := body["prompt_cache_key"]; exists { + t.Fatalf("prompt_cache_key should NOT be sent to %s, but was included in request", tt.name) + } + }) + } +} + +func TestSupportsPromptCacheKey(t *testing.T) { + tests := []struct { + apiBase string + want bool + }{ + {"https://api.openai.com/v1", true}, + {"https://api.openai.com/v1/", true}, + {"https://myresource.openai.azure.com/openai/deployments/gpt-4", true}, + {"https://eastus.openai.azure.com/v1", true}, + {"https://api.mistral.ai/v1", false}, + {"https://generativelanguage.googleapis.com/v1beta", false}, + {"https://api.deepseek.com/v1", false}, + {"https://api.groq.com/openai/v1", false}, + {"http://localhost:11434/v1", false}, + {"https://openrouter.ai/api/v1", false}, + // Edge cases: proxy URLs with openai.com in path should NOT match + {"https://my-proxy.com/api.openai.com/v1", false}, + {"https://proxy.example.com/openai.azure.com/v1", false}, + // Malformed or empty + {"", false}, + {"not-a-url", false}, + } + for _, tt := range tests { + if got := supportsPromptCacheKey(tt.apiBase); got != tt.want { + t.Errorf("supportsPromptCacheKey(%q) = %v, want %v", tt.apiBase, got, tt.want) + } + } +} + +func TestSerializeMessages_StripsSystemParts(t *testing.T) { + messages := []protocoltypes.Message{ + { + Role: "system", + Content: "you are helpful", + SystemParts: []protocoltypes.ContentBlock{ + {Type: "text", Text: "you are helpful"}, + }, + }, + } + result := serializeMessages(messages) + + data, _ := json.Marshal(result) + raw := string(data) + if strings.Contains(raw, "system_parts") { + t.Fatal("system_parts should not appear in serialized output") + } +} diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 0b6b7b3b6..708f96a96 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -208,7 +208,10 @@ func (sm *SessionManager) TruncateHistory(key string, keepLast int) { // so loadSessions still maps back to the right in-memory key. func sanitizeFilename(key string) string { - return strings.ReplaceAll(key, ":", "_") + s := strings.ReplaceAll(key, ":", "_") + s = strings.ReplaceAll(s, "/", "_") + s = strings.ReplaceAll(s, "\\", "_") + return s } func (sm *SessionManager) Save(key string) error { diff --git a/pkg/session/manager_ext_test.go b/pkg/session/manager_ext_test.go new file mode 100644 index 000000000..5f7441af1 --- /dev/null +++ b/pkg/session/manager_ext_test.go @@ -0,0 +1,121 @@ +package session + +import ( + "github.com/sipeed/picoclaw/pkg/providers" + "testing" +) + +func TestSanitizeHistory_OrphanedToolCall(t *testing.T) { + history := []providers.Message{ + {Role: "user", Content: "hello"}, + + {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ + {ID: "call_1", Name: "exec"}, + + {ID: "call_2", Name: "list_dir"}, + }}, + + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, + } + + sanitized, removed := SanitizeHistory(history) + + if removed == 0 { + t.Fatal("expected orphaned messages to be removed") + } + + if len(sanitized) != 1 || sanitized[0].Role != "user" { + t.Errorf("expected [user], got %d messages", len(sanitized)) + } +} + +func TestSanitizeHistory_InterleavedMessages(t *testing.T) { + + history := []providers.Message{ + {Role: "user", Content: "first"}, + + {Role: "assistant", Content: "ok", ToolCalls: []providers.ToolCall{ + {ID: "call_1", Name: "exec"}, + }}, + + {Role: "user", Content: "collision!"}, + + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, + + {Role: "assistant", Content: "done"}, + } + + sanitized, removed := SanitizeHistory(history) + + if removed == 0 { + t.Fatal("expected interleaved messages to be removed") + } + + if len(sanitized) != 3 { + t.Errorf("expected 3 messages, got %d", len(sanitized)) + + for i, m := range sanitized { + t.Logf(" [%d] role=%s content=%q", i, m.Role, m.Content) + } + } +} + +func TestSanitizeHistory_CleanHistory(t *testing.T) { + history := []providers.Message{ + {Role: "user", Content: "hello"}, + + {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ + {ID: "call_1", Name: "exec"}, + }}, + + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, + + {Role: "assistant", Content: "done"}, + } + + sanitized, removed := SanitizeHistory(history) + + if removed != 0 { + t.Errorf("expected 0 removed, got %d", removed) + } + + if len(sanitized) != 4 { + t.Errorf("expected 4 messages, got %d", len(sanitized)) + } +} + +func TestSanitizeHistory_MultipleToolCalls(t *testing.T) { + history := []providers.Message{ + {Role: "user", Content: "hello"}, + + {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ + {ID: "call_1", Name: "exec"}, + + {ID: "call_2", Name: "read_file"}, + }}, + + {Role: "tool", Content: "ok", ToolCallID: "call_1"}, + + {Role: "tool", Content: "content", ToolCallID: "call_2"}, + + {Role: "assistant", Content: "all done"}, + } + + sanitized, removed := SanitizeHistory(history) + + if removed != 0 { + t.Errorf("expected 0 removed, got %d", removed) + } + + if len(sanitized) != 5 { + t.Errorf("expected 5 messages, got %d", len(sanitized)) + } +} + +func TestSanitizeHistory_Empty(t *testing.T) { + sanitized, removed := SanitizeHistory(nil) + + if removed != 0 || sanitized != nil { + t.Errorf("expected nil/0, got %v/%d", sanitized, removed) + } +} diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go index 05b150015..bc5615966 100644 --- a/pkg/session/manager_test.go +++ b/pkg/session/manager_test.go @@ -4,33 +4,25 @@ import ( "os" "path/filepath" "testing" - - "github.com/sipeed/picoclaw/pkg/providers" ) func TestSanitizeFilename(t *testing.T) { tests := []struct { - input string - + input string expected string }{ {"simple", "simple"}, - {"telegram:123456", "telegram_123456"}, - {"discord:987654321", "discord_987654321"}, - {"slack:C01234", "slack_C01234"}, - {"no-colons-here", "no-colons-here"}, - {"multiple:colons:here", "multiple_colons_here"}, + {"agent:main:telegram:group:-1003822706455/12", "agent_main_telegram_group_-1003822706455_12"}, } for _, tt := range tests { t.Run(tt.input, func(t *testing.T) { got := sanitizeFilename(tt.input) - if got != tt.expected { t.Errorf("sanitizeFilename(%q) = %q, want %q", tt.input, got, tt.expected) } @@ -40,185 +32,54 @@ func TestSanitizeFilename(t *testing.T) { func TestSave_WithColonInKey(t *testing.T) { tmpDir := t.TempDir() - sm := NewSessionManager(tmpDir) // Create a session with a key containing colon (typical channel session key). - key := "telegram:123456" - sm.GetOrCreate(key) - sm.AddMessage(key, "user", "hello") // Save should succeed even though the key contains ':' - if err := sm.Save(key); err != nil { t.Fatalf("Save(%q) failed: %v", key, err) } // The file on disk should use sanitized name. - expectedFile := filepath.Join(tmpDir, "telegram_123456.json") - if _, err := os.Stat(expectedFile); os.IsNotExist(err) { t.Fatalf("expected session file %s to exist", expectedFile) } // Load into a fresh manager and verify the session round-trips. - sm2 := NewSessionManager(tmpDir) - history := sm2.GetHistory(key) - if len(history) != 1 { t.Fatalf("expected 1 message after reload, got %d", len(history)) } - if history[0].Content != "hello" { t.Errorf("expected message content %q, got %q", "hello", history[0].Content) } } -func TestSanitizeHistory_OrphanedToolCall(t *testing.T) { - history := []providers.Message{ - {Role: "user", Content: "hello"}, - - {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ - {ID: "call_1", Name: "exec"}, - - {ID: "call_2", Name: "list_dir"}, - }}, - - {Role: "tool", Content: "ok", ToolCallID: "call_1"}, - - // Missing tool result for call_2 → orphaned - - } - - sanitized, removed := SanitizeHistory(history) - - if removed == 0 { - t.Fatal("expected orphaned messages to be removed") - } - - // After sanitization, only the user message should remain - - if len(sanitized) != 1 || sanitized[0].Role != "user" { - t.Errorf("expected [user], got %d messages", len(sanitized)) - } -} - -func TestSanitizeHistory_InterleavedMessages(t *testing.T) { - // Simulates session collision: a user message got interleaved between - - // an assistant tool call and its tool result - - history := []providers.Message{ - {Role: "user", Content: "first"}, - - {Role: "assistant", Content: "ok", ToolCalls: []providers.ToolCall{ - {ID: "call_1", Name: "exec"}, - }}, - - {Role: "user", Content: "collision!"}, // ← interleaved from other session - - {Role: "tool", Content: "ok", ToolCallID: "call_1"}, // ← out of order - - {Role: "assistant", Content: "done"}, - } - - sanitized, removed := SanitizeHistory(history) - - if removed == 0 { - t.Fatal("expected interleaved messages to be removed") - } - - // Should keep: user("first"), user("collision!"), assistant("done") - - // Should remove: assistant(call_1), tool(call_1) - - if len(sanitized) != 3 { - t.Errorf("expected 3 messages, got %d", len(sanitized)) - - for i, m := range sanitized { - t.Logf(" [%d] role=%s content=%q", i, m.Role, m.Content) - } - } -} - -func TestSanitizeHistory_CleanHistory(t *testing.T) { - history := []providers.Message{ - {Role: "user", Content: "hello"}, - - {Role: "assistant", Content: "sure", ToolCalls: []providers.ToolCall{ - {ID: "call_1", Name: "exec"}, - }}, - - {Role: "tool", Content: "ok", ToolCallID: "call_1"}, - - {Role: "assistant", Content: "done"}, - } - - sanitized, removed := SanitizeHistory(history) - - if removed != 0 { - t.Errorf("expected 0 removed, got %d", removed) - } - - if len(sanitized) != 4 { - t.Errorf("expected 4 messages, got %d", len(sanitized)) - } -} - -func TestSanitizeHistory_MultipleToolCalls(t *testing.T) { - history := []providers.Message{ - {Role: "user", Content: "hello"}, - - {Role: "assistant", Content: "", ToolCalls: []providers.ToolCall{ - {ID: "call_1", Name: "exec"}, - - {ID: "call_2", Name: "read_file"}, - }}, - - {Role: "tool", Content: "ok", ToolCallID: "call_1"}, - - {Role: "tool", Content: "content", ToolCallID: "call_2"}, - - {Role: "assistant", Content: "all done"}, - } - - sanitized, removed := SanitizeHistory(history) - - if removed != 0 { - t.Errorf("expected 0 removed, got %d", removed) - } - - if len(sanitized) != 5 { - t.Errorf("expected 5 messages, got %d", len(sanitized)) - } -} - -func TestSanitizeHistory_Empty(t *testing.T) { - sanitized, removed := SanitizeHistory(nil) - - if removed != 0 || sanitized != nil { - t.Errorf("expected nil/0, got %v/%d", sanitized, removed) - } -} - func TestSave_RejectsPathTraversal(t *testing.T) { tmpDir := t.TempDir() - sm := NewSessionManager(tmpDir) - badKeys := []string{"", ".", "..", "foo/bar", "foo\\bar"} - + // Invalid names that must still be rejected. + badKeys := []string{"", ".", ".."} for _, key := range badKeys { sm.GetOrCreate(key) - if err := sm.Save(key); err == nil { t.Errorf("Save(%q) should have failed but didn't", key) } } + + // Keys containing path separators are sanitized (no subdirs created). + sm.GetOrCreate("foo/bar") + if err := sm.Save("foo/bar"); err != nil { + t.Fatalf("Save(\"foo/bar\") after sanitize should succeed: %v", err) + } + if _, err := os.Stat(filepath.Join(tmpDir, "foo_bar.json")); os.IsNotExist(err) { + t.Errorf("expected foo_bar.json in storage (sanitized from foo/bar)") + } } diff --git a/pkg/skills/loader_test.go b/pkg/skills/loader_test.go index 31619f9c2..645d8b7ac 100644 --- a/pkg/skills/loader_test.go +++ b/pkg/skills/loader_test.go @@ -342,3 +342,78 @@ func TestSkillRootsTrimsWhitespaceAndDedups(t *testing.T) { builtin, }, roots) } + +func TestGetSkillMetadata_UsesMarkdownParagraphWhenNoFrontmatter(t *testing.T) { + tmp := t.TempDir() + skillDir := filepath.Join(tmp, "workspace", "skills", "plain-skill") + require.NoError(t, os.MkdirAll(skillDir, 0o755)) + + content := "# Plain Skill\n\nThis is parsed from markdown paragraph.\n" + require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644)) + + sl := &SkillsLoader{} + meta := sl.getSkillMetadata(filepath.Join(skillDir, "SKILL.md")) + require.NotNil(t, meta) + assert.Equal(t, "plain-skill", meta.Name) + assert.Equal(t, "This is parsed from markdown paragraph.", meta.Description) +} + +func TestGetSkillMetadata_FrontmatterOverridesMarkdown(t *testing.T) { + tmp := t.TempDir() + skillDir := filepath.Join(tmp, "workspace", "skills", "plain-skill") + require.NoError(t, os.MkdirAll(skillDir, 0o755)) + + content := "---\nname: frontmatter-skill\ndescription: frontmatter description\n---\n\n# Plain Skill\n\nBody description.\n" + require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644)) + + sl := &SkillsLoader{} + meta := sl.getSkillMetadata(filepath.Join(skillDir, "SKILL.md")) + require.NotNil(t, meta) + assert.Equal(t, "frontmatter-skill", meta.Name) + assert.Equal(t, "frontmatter description", meta.Description) +} + +func TestGetSkillMetadata_YAMLMultilineDescription(t *testing.T) { + tmp := t.TempDir() + skillDir := filepath.Join(tmp, "workspace", "skills", "plain-skill") + require.NoError(t, os.MkdirAll(skillDir, 0o755)) + + content := "---\nname: frontmatter-skill\ndescription: |\n line 1: with colon\n line 2\n---\n\n# Plain Skill\n\nBody description.\n" + require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644)) + + sl := &SkillsLoader{} + meta := sl.getSkillMetadata(filepath.Join(skillDir, "SKILL.md")) + require.NotNil(t, meta) + assert.Equal(t, "frontmatter-skill", meta.Name) + assert.Equal(t, "line 1: with colon\nline 2", meta.Description) +} + +func TestGetSkillMetadata_InvalidHeadingNameFallsBackToDirName(t *testing.T) { + tmp := t.TempDir() + skillDir := filepath.Join(tmp, "workspace", "skills", "valid-name") + require.NoError(t, os.MkdirAll(skillDir, 0o755)) + + content := "# Invalid Heading Name\n\nBody description.\n" + require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644)) + + sl := &SkillsLoader{} + meta := sl.getSkillMetadata(filepath.Join(skillDir, "SKILL.md")) + require.NotNil(t, meta) + assert.Equal(t, "valid-name", meta.Name) + assert.Equal(t, "Body description.", meta.Description) +} + +func TestGetSkillMetadata_IgnoresHTMLCommentBlocks(t *testing.T) { + tmp := t.TempDir() + skillDir := filepath.Join(tmp, "workspace", "skills", "biomed-skill") + require.NoError(t, os.MkdirAll(skillDir, 0o755)) + + content := "\n\n# Biomed Skill\n\nSummarize biomedical papers.\n" + require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644)) + + sl := &SkillsLoader{} + meta := sl.getSkillMetadata(filepath.Join(skillDir, "SKILL.md")) + require.NotNil(t, meta) + assert.Equal(t, "biomed-skill", meta.Name) + assert.Equal(t, "Summarize biomedical papers.", meta.Description) +} diff --git a/pkg/state/state.go b/pkg/state/state.go index 34d1576b6..60a204ab9 100644 --- a/pkg/state/state.go +++ b/pkg/state/state.go @@ -48,8 +48,8 @@ func NewManager(workspace string) *Manager { oldStateFile := filepath.Join(workspace, "state.json") // Create state directory if it doesn't exist - if err := os.MkdirAll(stateDir, 0o755); err != nil { - log.Fatalf("[FATAL] state: failed to create state directory: %v", err) + if err := os.MkdirAll(stateDir, 0o700); err != nil { + log.Printf("[WARN] state: failed to create state directory %s: %v", stateDir, err) } sm := &Manager{ diff --git a/pkg/state/state_ext_test.go b/pkg/state/state_ext_test.go new file mode 100644 index 000000000..ac7e83edb --- /dev/null +++ b/pkg/state/state_ext_test.go @@ -0,0 +1,38 @@ +package state + +import ( + "os" + "testing" +) + +func TestHeartbeatTargetsPersistence(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "state-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + sm := NewManager(tmpDir) + + if err := sm.SetLastHeartbeatTarget("telegram:-100123"); err != nil { + t.Fatalf("SetLastHeartbeatTarget failed: %v", err) + } + if err := sm.SetHeartbeatTarget("telegram:-100123/42"); err != nil { + t.Fatalf("SetHeartbeatTarget failed: %v", err) + } + + if got := sm.GetLastHeartbeatTarget(); got != "telegram:-100123" { + t.Fatalf("GetLastHeartbeatTarget = %q, want %q", got, "telegram:-100123") + } + if got := sm.GetHeartbeatTarget(); got != "telegram:-100123/42" { + t.Fatalf("GetHeartbeatTarget = %q, want %q", got, "telegram:-100123/42") + } + + sm2 := NewManager(tmpDir) + if got := sm2.GetLastHeartbeatTarget(); got != "telegram:-100123" { + t.Fatalf("persistent GetLastHeartbeatTarget = %q, want %q", got, "telegram:-100123") + } + if got := sm2.GetHeartbeatTarget(); got != "telegram:-100123/42" { + t.Fatalf("persistent GetHeartbeatTarget = %q, want %q", got, "telegram:-100123/42") + } +} diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go index 02d5c6227..3924e5533 100644 --- a/pkg/state/state_test.go +++ b/pkg/state/state_test.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "os" + "os/exec" "path/filepath" "testing" ) @@ -215,34 +216,31 @@ func TestNewManager_EmptyWorkspace(t *testing.T) { } } -func TestHeartbeatTargetsPersistence(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "state-test-*") +func TestNewManager_MkdirFailureDoesNotCrash(t *testing.T) { + if os.Getenv("BE_CRASHER") == "1" { + tmpDir := os.Getenv("CRASH_DIR") + + statePath := filepath.Join(tmpDir, "state") + if err := os.WriteFile(statePath, []byte("I'm a file, not a folder"), 0o644); err != nil { + fmt.Printf("setup failed: %v", err) + os.Exit(0) + } + + NewManager(tmpDir) + os.Exit(0) + } + + tmpDir, err := os.MkdirTemp("", "state-crash-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } defer os.RemoveAll(tmpDir) - sm := NewManager(tmpDir) + cmd := exec.Command(os.Args[0], "-test.run=TestNewManager_MkdirFailureDoesNotCrash") + cmd.Env = append(os.Environ(), "BE_CRASHER=1", "CRASH_DIR="+tmpDir) - if err := sm.SetLastHeartbeatTarget("telegram:-100123"); err != nil { - t.Fatalf("SetLastHeartbeatTarget failed: %v", err) - } - if err := sm.SetHeartbeatTarget("telegram:-100123/42"); err != nil { - t.Fatalf("SetHeartbeatTarget failed: %v", err) - } - - if got := sm.GetLastHeartbeatTarget(); got != "telegram:-100123" { - t.Fatalf("GetLastHeartbeatTarget = %q, want %q", got, "telegram:-100123") - } - if got := sm.GetHeartbeatTarget(); got != "telegram:-100123/42" { - t.Fatalf("GetHeartbeatTarget = %q, want %q", got, "telegram:-100123/42") - } - - sm2 := NewManager(tmpDir) - if got := sm2.GetLastHeartbeatTarget(); got != "telegram:-100123" { - t.Fatalf("persistent GetLastHeartbeatTarget = %q, want %q", got, "telegram:-100123") - } - if got := sm2.GetHeartbeatTarget(); got != "telegram:-100123/42" { - t.Fatalf("persistent GetHeartbeatTarget = %q, want %q", got, "telegram:-100123/42") + err = cmd.Run() + if err != nil { + t.Fatalf("NewManager should not crash when state dir creation fails, got: %v", err) } } diff --git a/pkg/tools/edit_test.go b/pkg/tools/edit_test.go index ccad894eb..83a7e778c 100644 --- a/pkg/tools/edit_test.go +++ b/pkg/tools/edit_test.go @@ -11,349 +11,261 @@ import ( ) // TestEditTool_EditFile_Success verifies successful file editing - func TestEditTool_EditFile_Success(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "test.txt") - os.WriteFile(testFile, []byte("Hello World\nThis is a test"), 0o644) tool := NewEditFileTool(tmpDir, true) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "old_text": "World", - "new_text": "Universe", } result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // Should return SilentResult - if !result.Silent { t.Errorf("Expected Silent=true for EditFile, got false") } // ForUser should be empty (silent result) - if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) } // Verify file was actually edited - content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read edited file: %v", err) } - contentStr := string(content) - if !strings.Contains(contentStr, "Hello Universe") { t.Errorf("Expected file to contain 'Hello Universe', got: %s", contentStr) } - if strings.Contains(contentStr, "Hello World") { t.Errorf("Expected 'Hello World' to be replaced, got: %s", contentStr) } } // TestEditTool_EditFile_NotFound verifies error handling for non-existent file - func TestEditTool_EditFile_NotFound(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "nonexistent.txt") tool := NewEditFileTool(tmpDir, true) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "old_text": "old", - "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error for non-existent file") } // Should mention file not found - if !strings.Contains(result.ForLLM, "not found") && !strings.Contains(result.ForUser, "not found") { t.Errorf("Expected 'file not found' message, got ForLLM: %s", result.ForLLM) } } // TestEditTool_EditFile_OldTextNotFound verifies error when old_text doesn't exist - func TestEditTool_EditFile_OldTextNotFound(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "test.txt") - os.WriteFile(testFile, []byte("Hello World"), 0o644) tool := NewEditFileTool(tmpDir, true) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "old_text": "Goodbye", - "new_text": "Hello", } result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when old_text not found") } // Should mention old_text not found - if !strings.Contains(result.ForLLM, "not found") && !strings.Contains(result.ForUser, "not found") { t.Errorf("Expected 'not found' message, got ForLLM: %s", result.ForLLM) } } // TestEditTool_EditFile_MultipleMatches verifies error when old_text appears multiple times - func TestEditTool_EditFile_MultipleMatches(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "test.txt") - os.WriteFile(testFile, []byte("test test test"), 0o644) tool := NewEditFileTool(tmpDir, true) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "old_text": "test", - "new_text": "done", } result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when old_text appears multiple times") } // Should mention multiple occurrences - if !strings.Contains(result.ForLLM, "times") && !strings.Contains(result.ForUser, "times") { t.Errorf("Expected 'multiple times' message, got ForLLM: %s", result.ForLLM) } } // TestEditTool_EditFile_OutsideAllowedDir verifies error when path is outside allowed directory - func TestEditTool_EditFile_OutsideAllowedDir(t *testing.T) { tmpDir := t.TempDir() - otherDir := t.TempDir() - testFile := filepath.Join(otherDir, "test.txt") - os.WriteFile(testFile, []byte("content"), 0o644) tool := NewEditFileTool(tmpDir, true) // Restrict to tmpDir - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "old_text": "content", - "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result - assert.True(t, result.IsError, "Expected error when path is outside allowed directory") // Should mention outside allowed directory - // Note: ErrorResult only sets ForLLM by default, so ForUser might be empty. - // We check ForLLM as it's the primary error channel. - assert.True( - t, - strings.Contains(result.ForLLM, "outside") || strings.Contains(result.ForLLM, "access denied") || - strings.Contains(result.ForLLM, "escapes"), - "Expected 'outside allowed' or 'access denied' message, got ForLLM: %s", - result.ForLLM, ) } // TestEditTool_EditFile_MissingPath verifies error handling for missing path - func TestEditTool_EditFile_MissingPath(t *testing.T) { tool := NewEditFileTool("", false) - ctx := context.Background() - args := map[string]any{ "old_text": "old", - "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when path is missing") } } // TestEditTool_EditFile_MissingOldText verifies error handling for missing old_text - func TestEditTool_EditFile_MissingOldText(t *testing.T) { tool := NewEditFileTool("", false) - ctx := context.Background() - args := map[string]any{ - "path": "/tmp/test.txt", - + "path": "/tmp/test.txt", "new_text": "new", } result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when old_text is missing") } } // TestEditTool_EditFile_MissingNewText verifies error handling for missing new_text - func TestEditTool_EditFile_MissingNewText(t *testing.T) { tool := NewEditFileTool("", false) - ctx := context.Background() - args := map[string]any{ - "path": "/tmp/test.txt", - + "path": "/tmp/test.txt", "old_text": "old", } result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when new_text is missing") } } // TestEditTool_AppendFile_Success verifies successful file appending - func TestEditTool_AppendFile_Success(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "test.txt") - os.WriteFile(testFile, []byte("Initial content"), 0o644) tool := NewAppendFileTool("", false) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "content": "\nAppended content", } result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // Should return SilentResult - if !result.Silent { t.Errorf("Expected Silent=true for AppendFile, got false") } // ForUser should be empty (silent result) - if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) } // Verify content was actually appended - content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read file: %v", err) } - contentStr := string(content) - if !strings.Contains(contentStr, "Initial content") { t.Errorf("Expected original content to remain, got: %s", contentStr) } - if !strings.Contains(contentStr, "Appended content") { t.Errorf("Expected appended content, got: %s", contentStr) } } // TestEditTool_AppendFile_MissingPath verifies error handling for missing path - func TestEditTool_AppendFile_MissingPath(t *testing.T) { tool := NewAppendFileTool("", false) - ctx := context.Background() - args := map[string]any{ "content": "test", } @@ -361,19 +273,15 @@ func TestEditTool_AppendFile_MissingPath(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when path is missing") } } // TestEditTool_AppendFile_MissingContent verifies error handling for missing content - func TestEditTool_AppendFile_MissingContent(t *testing.T) { tool := NewAppendFileTool("", false) - ctx := context.Background() - args := map[string]any{ "path": "/tmp/test.txt", } @@ -381,67 +289,43 @@ func TestEditTool_AppendFile_MissingContent(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when content is missing") } } // TestReplaceEditContent verifies the helper function replaceEditContent - func TestReplaceEditContent(t *testing.T) { tests := []struct { - name string - - content []byte - - oldText string - - newText string - - expected []byte - + name string + content []byte + oldText string + newText string + expected []byte expectError bool }{ { - name: "successful replacement", - - content: []byte("hello world"), - - oldText: "world", - - newText: "universe", - - expected: []byte("hello universe"), - + name: "successful replacement", + content: []byte("hello world"), + oldText: "world", + newText: "universe", + expected: []byte("hello universe"), expectError: false, }, - { - name: "old text not found", - - content: []byte("hello world"), - - oldText: "golang", - - newText: "rust", - - expected: nil, - + name: "old text not found", + content: []byte("hello world"), + oldText: "golang", + newText: "rust", + expected: nil, expectError: true, }, - { - name: "multiple matches found", - - content: []byte("test text test"), - - oldText: "test", - - newText: "done", - - expected: nil, - + name: "multiple matches found", + content: []byte("test text test"), + oldText: "test", + newText: "done", + expected: nil, expectError: true, }, } @@ -449,12 +333,10 @@ func TestReplaceEditContent(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result, err := replaceEditContent(tt.content, tt.oldText, tt.newText) - if tt.expectError { assert.Error(t, err) } else { assert.NoError(t, err) - assert.Equal(t, tt.expected, result) } }) @@ -462,142 +344,94 @@ func TestReplaceEditContent(t *testing.T) { } // TestAppendFileTool_AppendToNonExistent_Restricted verifies that AppendFileTool in restricted mode - // can append to a file that does not yet exist — it should silently create the file. - -// This exercises the errors.Is(err, fs.ErrNotExist) path in appendFile + sandboxFs. - +// This exercises the errors.Is(err, fs.ErrNotExist) path in appendFileWithRW + rootRW. func TestAppendFileTool_AppendToNonExistent_Restricted(t *testing.T) { workspace := t.TempDir() - tool := NewAppendFileTool(workspace, true) - ctx := context.Background() args := map[string]any{ - "path": "brand_new_file.txt", - + "path": "brand_new_file.txt", "content": "first content", } result := tool.Execute(ctx, args) - assert.False( - t, - result.IsError, - "Expected success when appending to non-existent file in restricted mode, got: %s", - result.ForLLM, ) // Verify the file was created with correct content - data, err := os.ReadFile(filepath.Join(workspace, "brand_new_file.txt")) - assert.NoError(t, err) - assert.Equal(t, "first content", string(data)) } // TestAppendFileTool_Restricted_Success verifies that AppendFileTool in restricted mode - // correctly appends to an existing file within the sandbox. - func TestAppendFileTool_Restricted_Success(t *testing.T) { workspace := t.TempDir() - testFile := "existing.txt" - err := os.WriteFile(filepath.Join(workspace, testFile), []byte("initial"), 0o644) - assert.NoError(t, err) tool := NewAppendFileTool(workspace, true) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "content": " appended", } result := tool.Execute(ctx, args) - assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) - assert.True(t, result.Silent) data, err := os.ReadFile(filepath.Join(workspace, testFile)) - assert.NoError(t, err) - assert.Equal(t, "initial appended", string(data)) } // TestEditFileTool_Restricted_InPlaceEdit verifies that EditFileTool in restricted mode - -// correctly edits a file using the sandboxFs path. - +// correctly edits a file using the single-open editFileInRoot path. func TestEditFileTool_Restricted_InPlaceEdit(t *testing.T) { workspace := t.TempDir() - testFile := "edit_target.txt" - err := os.WriteFile(filepath.Join(workspace, testFile), []byte("Hello World"), 0o644) - assert.NoError(t, err) tool := NewEditFileTool(workspace, true) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "old_text": "World", - "new_text": "Go", } result := tool.Execute(ctx, args) - assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) - assert.True(t, result.Silent) data, err := os.ReadFile(filepath.Join(workspace, testFile)) - assert.NoError(t, err) - assert.Equal(t, "Hello Go", string(data)) } -// TestEditFileTool_Restricted_FileNotFound verifies that editFile returns a proper - +// TestEditFileTool_Restricted_FileNotFound verifies that editFileInRoot returns a proper // error message when the target file does not exist. - func TestEditFileTool_Restricted_FileNotFound(t *testing.T) { workspace := t.TempDir() - tool := NewEditFileTool(workspace, true) - ctx := context.Background() - args := map[string]any{ - "path": "no_such_file.txt", - + "path": "no_such_file.txt", "old_text": "old", - "new_text": "new", } result := tool.Execute(ctx, args) - assert.True(t, result.IsError) - assert.Contains(t, result.ForLLM, "not found") } diff --git a/pkg/tools/filesystem_ext_test.go b/pkg/tools/filesystem_ext_test.go new file mode 100644 index 000000000..14429605e --- /dev/null +++ b/pkg/tools/filesystem_ext_test.go @@ -0,0 +1,183 @@ +package tools + +import ( + "github.com/stretchr/testify/assert" + "io" + "os" + "path/filepath" + "testing" +) + +func TestHostFs_Read_PermissionDenied(t *testing.T) { + if os.Getuid() == 0 { + t.Skip("skipping permission test: running as root") + } + + tmpDir := t.TempDir() + + protected := filepath.Join(tmpDir, "protected.txt") + + err := os.WriteFile(protected, []byte("secret"), 0o000) + + assert.NoError(t, err) + + defer os.Chmod(protected, 0o644) + + _, err = (&hostFs{}).ReadFile(protected) + + assert.Error(t, err) + + assert.Contains(t, err.Error(), "access denied") +} + +func TestHostFs_Read_Directory(t *testing.T) { + tmpDir := t.TempDir() + + _, err := (&hostFs{}).ReadFile(tmpDir) + + assert.Error(t, err, "expected error when reading a directory as a file") +} + +func TestSandboxFs_Read_Directory(t *testing.T) { + workspace := t.TempDir() + + root, err := os.OpenRoot(workspace) + + assert.NoError(t, err) + + defer root.Close() + + err = root.Mkdir("subdir", 0o755) + + assert.NoError(t, err) + + _, err = (&sandboxFs{workspace: workspace}).ReadFile("subdir") + + assert.Error(t, err, "expected error when reading a directory as a file") +} + +func TestHostFs_Write_ParentDirMissing(t *testing.T) { + tmpDir := t.TempDir() + + target := filepath.Join(tmpDir, "a", "b", "c", "file.txt") + + err := (&hostFs{}).WriteFile(target, []byte("hello")) + + assert.NoError(t, err) + + data, err := os.ReadFile(target) + + assert.NoError(t, err) + + assert.Equal(t, "hello", string(data)) +} + +func TestSandboxFs_Write_ParentDirMissing(t *testing.T) { + workspace := t.TempDir() + + relPath := "x/y/z/file.txt" + + err := (&sandboxFs{workspace: workspace}).WriteFile(relPath, []byte("nested")) + + assert.NoError(t, err) + + data, err := os.ReadFile(filepath.Join(workspace, relPath)) + + assert.NoError(t, err) + + assert.Equal(t, "nested", string(data)) +} + +func TestHostFs_Write(t *testing.T) { + tmpDir := t.TempDir() + + testFile := filepath.Join(tmpDir, "atomic_test.txt") + + testData := []byte("atomic test content") + + err := (&hostFs{}).WriteFile(testFile, testData) + + assert.NoError(t, err) + + content, err := os.ReadFile(testFile) + + assert.NoError(t, err) + + assert.Equal(t, testData, content) + + newData := []byte("new atomic content") + + err = (&hostFs{}).WriteFile(testFile, newData) + + assert.NoError(t, err) + + content, err = os.ReadFile(testFile) + + assert.NoError(t, err) + + assert.Equal(t, newData, content) +} + +func TestSandboxFs_Write(t *testing.T) { + tmpDir := t.TempDir() + + relPath := "atomic_root_test.txt" + + testData := []byte("atomic root test content") + + erw := &sandboxFs{workspace: tmpDir} + + err := erw.WriteFile(relPath, testData) + + assert.NoError(t, err) + + root, err := os.OpenRoot(tmpDir) + + assert.NoError(t, err) + + defer root.Close() + + f, err := root.Open(relPath) + + assert.NoError(t, err) + + defer f.Close() + + content, err := io.ReadAll(f) + + assert.NoError(t, err) + + assert.Equal(t, testData, content) + + newData := []byte("new root atomic content") + + err = erw.WriteFile(relPath, newData) + + assert.NoError(t, err) + + f2, err := root.Open(relPath) + + assert.NoError(t, err) + + defer f2.Close() + + content, err = io.ReadAll(f2) + + assert.NoError(t, err) + + assert.Equal(t, newData, content) +} + +func TestValidatePath_OutsideWorkspace_IncludesPath(t *testing.T) { + workspace := t.TempDir() + + outsidePath := filepath.Join(t.TempDir(), "secret.txt") + + _, err := validatePath(outsidePath, workspace, true) + + assert.Error(t, err) + + assert.Contains(t, err.Error(), "access denied") + + assert.Contains(t, err.Error(), workspace) +} diff --git a/pkg/tools/filesystem_test.go b/pkg/tools/filesystem_test.go index 456f8fbd3..0bbf6caf0 100644 --- a/pkg/tools/filesystem_test.go +++ b/pkg/tools/filesystem_test.go @@ -5,6 +5,7 @@ import ( "io" "os" "path/filepath" + "regexp" "strings" "testing" @@ -12,18 +13,13 @@ import ( ) // TestFilesystemTool_ReadFile_Success verifies successful file reading - func TestFilesystemTool_ReadFile_Success(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "test.txt") - os.WriteFile(testFile, []byte("test content"), 0o644) - tool := NewReadFileTool("", false) - + tool := NewReadFileTool("", false, MaxReadFileSize) ctx := context.Background() - args := map[string]any{ "path": testFile, } @@ -31,33 +27,26 @@ func TestFilesystemTool_ReadFile_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForLLM should contain file content - if !strings.Contains(result.ForLLM, "test content") { t.Errorf("Expected ForLLM to contain 'test content', got: %s", result.ForLLM) } // ReadFile returns NewToolResult which only sets ForLLM, not ForUser - // This is the expected behavior - file content goes to LLM, not directly to user - if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for NewToolResult, got: %s", result.ForUser) } } // TestFilesystemTool_ReadFile_NotFound verifies error handling for missing file - func TestFilesystemTool_ReadFile_NotFound(t *testing.T) { - tool := NewReadFileTool("", false) - + tool := NewReadFileTool("", false, MaxReadFileSize) ctx := context.Background() - args := map[string]any{ "path": "/nonexistent_file_12345.txt", } @@ -65,135 +54,107 @@ func TestFilesystemTool_ReadFile_NotFound(t *testing.T) { result := tool.Execute(ctx, args) // Failure should be marked as error - if !result.IsError { t.Errorf("Expected error for missing file, got IsError=false") } // Should contain error message - - if !strings.Contains(result.ForLLM, "failed to read") && !strings.Contains(result.ForUser, "failed to read") { + if !strings.Contains(result.ForLLM, "failed to open file") && !strings.Contains(result.ForUser, "failed to read") { t.Errorf("Expected error message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser) } } // TestFilesystemTool_ReadFile_MissingPath verifies error handling for missing path - func TestFilesystemTool_ReadFile_MissingPath(t *testing.T) { tool := &ReadFileTool{} - ctx := context.Background() - args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when path is missing") } // Should mention required parameter - if !strings.Contains(result.ForLLM, "path is required") && !strings.Contains(result.ForUser, "path is required") { t.Errorf("Expected 'path is required' message, got ForLLM: %s", result.ForLLM) } } // TestFilesystemTool_WriteFile_Success verifies successful file writing - func TestFilesystemTool_WriteFile_Success(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "newfile.txt") tool := NewWriteFileTool("", false) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "content": "hello world", } result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // WriteFile returns SilentResult - if !result.Silent { t.Errorf("Expected Silent=true for WriteFile, got false") } // ForUser should be empty (silent result) - if result.ForUser != "" { t.Errorf("Expected ForUser to be empty for SilentResult, got: %s", result.ForUser) } // Verify file was actually written - content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read written file: %v", err) } - if string(content) != "hello world" { t.Errorf("Expected file content 'hello world', got: %s", string(content)) } } // TestFilesystemTool_WriteFile_CreateDir verifies directory creation - func TestFilesystemTool_WriteFile_CreateDir(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "subdir", "newfile.txt") tool := NewWriteFileTool("", false) - ctx := context.Background() - args := map[string]any{ - "path": testFile, - + "path": testFile, "content": "test", } result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success with directory creation, got IsError=true: %s", result.ForLLM) } // Verify directory was created and file written - content, err := os.ReadFile(testFile) if err != nil { t.Fatalf("Failed to read written file: %v", err) } - if string(content) != "test" { t.Errorf("Expected file content 'test', got: %s", string(content)) } } // TestFilesystemTool_WriteFile_MissingPath verifies error handling for missing path - func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) { tool := NewWriteFileTool("", false) - ctx := context.Background() - args := map[string]any{ "content": "test", } @@ -201,19 +162,15 @@ func TestFilesystemTool_WriteFile_MissingPath(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when path is missing") } } // TestFilesystemTool_WriteFile_MissingContent verifies error handling for missing content - func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) { tool := NewWriteFileTool("", false) - ctx := context.Background() - args := map[string]any{ "path": "/tmp/test.txt", } @@ -221,35 +178,26 @@ func TestFilesystemTool_WriteFile_MissingContent(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when content is missing") } // Should mention required parameter - if !strings.Contains(result.ForLLM, "content is required") && - !strings.Contains(result.ForUser, "content is required") { t.Errorf("Expected 'content is required' message, got ForLLM: %s", result.ForLLM) } } // TestFilesystemTool_ListDir_Success verifies successful directory listing - func TestFilesystemTool_ListDir_Success(t *testing.T) { tmpDir := t.TempDir() - os.WriteFile(filepath.Join(tmpDir, "file1.txt"), []byte("content"), 0o644) - os.WriteFile(filepath.Join(tmpDir, "file2.txt"), []byte("content"), 0o644) - os.Mkdir(filepath.Join(tmpDir, "subdir"), 0o755) tool := NewListDirTool("", false) - ctx := context.Background() - args := map[string]any{ "path": tmpDir, } @@ -257,29 +205,23 @@ func TestFilesystemTool_ListDir_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // Should list files and directories - if !strings.Contains(result.ForLLM, "file1.txt") || !strings.Contains(result.ForLLM, "file2.txt") { t.Errorf("Expected files in listing, got: %s", result.ForLLM) } - if !strings.Contains(result.ForLLM, "subdir") { t.Errorf("Expected subdir in listing, got: %s", result.ForLLM) } } // TestFilesystemTool_ListDir_NotFound verifies error handling for non-existent directory - func TestFilesystemTool_ListDir_NotFound(t *testing.T) { tool := NewListDirTool("", false) - ctx := context.Background() - args := map[string]any{ "path": "/nonexistent_directory_12345", } @@ -287,61 +229,49 @@ func TestFilesystemTool_ListDir_NotFound(t *testing.T) { result := tool.Execute(ctx, args) // Failure should be marked as error - if !result.IsError { t.Errorf("Expected error for non-existent directory, got IsError=false") } // Should contain error message - if !strings.Contains(result.ForLLM, "failed to read") && !strings.Contains(result.ForUser, "failed to read") { t.Errorf("Expected error message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser) } } // TestFilesystemTool_ListDir_DefaultPath verifies default to current directory - func TestFilesystemTool_ListDir_DefaultPath(t *testing.T) { tool := NewListDirTool("", false) - ctx := context.Background() - args := map[string]any{} result := tool.Execute(ctx, args) // Should use "." as default path - if result.IsError { t.Errorf("Expected success with default path '.', got IsError=true: %s", result.ForLLM) } } // Block paths that look inside workspace but point outside via symlink. - func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) { root := t.TempDir() - workspace := filepath.Join(root, "workspace") - if err := os.MkdirAll(workspace, 0o755); err != nil { t.Fatalf("failed to create workspace: %v", err) } secret := filepath.Join(root, "secret.txt") - if err := os.WriteFile(secret, []byte("top secret"), 0o644); err != nil { t.Fatalf("failed to write secret file: %v", err) } link := filepath.Join(workspace, "leak.txt") - if err := os.Symlink(secret, link); err != nil { t.Skipf("symlink not supported in this environment: %v", err) } - tool := NewReadFileTool(workspace, true) - + tool := NewReadFileTool(workspace, true, MaxReadFileSize) result := tool.Execute(context.Background(), map[string]any{ "path": link, }) @@ -349,29 +279,21 @@ func TestFilesystemTool_ReadFile_RejectsSymlinkEscape(t *testing.T) { if !result.IsError { t.Fatalf("expected symlink escape to be blocked") } - // os.Root might return different errors depending on platform/implementation - // but it definitely should error. - // Our wrapper returns "access denied or file not found" - if !strings.Contains(result.ForLLM, "access denied") && !strings.Contains(result.ForLLM, "file not found") && - !strings.Contains(result.ForLLM, "no such file") { t.Fatalf("expected symlink escape error, got: %s", result.ForLLM) } } func TestFilesystemTool_EmptyWorkspace_AccessDenied(t *testing.T) { - tool := NewReadFileTool("", true) // restrict=true but workspace="" + tool := NewReadFileTool("", true, MaxReadFileSize) // restrict=true but workspace="" // Try to read a sensitive file (simulated by a temp file outside workspace) - tmpDir := t.TempDir() - secretFile := filepath.Join(tmpDir, "shadow") - os.WriteFile(secretFile, []byte("secret data"), 0o600) result := tool.Execute(context.Background(), map[string]any{ @@ -379,293 +301,346 @@ func TestFilesystemTool_EmptyWorkspace_AccessDenied(t *testing.T) { }) // We EXPECT IsError=true (access blocked due to empty workspace) - assert.True(t, result.IsError, "Security Regression: Empty workspace allowed access! content: %s", result.ForLLM) // Verify it failed for the right reason - assert.Contains(t, result.ForLLM, "workspace is not defined", "Expected 'workspace is not defined' error") } -// TestRootMkdirAll verifies that root.MkdirAll (used by sandboxFs.WriteFile) handles all cases: - +// TestRootMkdirAll verifies that root.MkdirAll (used by atomicWriteFileInRoot) handles all cases: // single dir, deeply nested dirs, already-existing dirs, and a file blocking a directory path. - func TestRootMkdirAll(t *testing.T) { workspace := t.TempDir() - root, err := os.OpenRoot(workspace) if err != nil { t.Fatalf("failed to open root: %v", err) } - defer root.Close() // Case 1: Single directory - err = root.MkdirAll("dir1", 0o755) - assert.NoError(t, err) - _, err = os.Stat(filepath.Join(workspace, "dir1")) - assert.NoError(t, err) // Case 2: Deeply nested directory - err = root.MkdirAll("a/b/c/d", 0o755) - assert.NoError(t, err) - _, err = os.Stat(filepath.Join(workspace, "a/b/c/d")) - assert.NoError(t, err) // Case 3: Already exists — must be idempotent - err = root.MkdirAll("a/b/c/d", 0o755) - assert.NoError(t, err) // Case 4: A regular file blocks directory creation — must error - err = os.WriteFile(filepath.Join(workspace, "file_exists"), []byte("data"), 0o644) - assert.NoError(t, err) - err = root.MkdirAll("file_exists", 0o755) - assert.Error(t, err, "expected error when a file exists at the directory path") } func TestFilesystemTool_WriteFile_Restricted_CreateDir(t *testing.T) { workspace := t.TempDir() - tool := NewWriteFileTool(workspace, true) - ctx := context.Background() testFile := "deep/nested/path/to/file.txt" - content := "deep content" - args := map[string]any{ - "path": testFile, - + "path": testFile, "content": content, } result := tool.Execute(ctx, args) - assert.False(t, result.IsError, "Expected success, got: %s", result.ForLLM) // Verify file content - actualPath := filepath.Join(workspace, testFile) - data, err := os.ReadFile(actualPath) - assert.NoError(t, err) - assert.Equal(t, content, string(data)) } -// TestHostFs_Read_PermissionDenied verifies that hostFs.ReadFile surfaces access denied errors. - -func TestHostFs_Read_PermissionDenied(t *testing.T) { +// TestHostRW_Read_PermissionDenied verifies that hostRW.Read surfaces access denied errors. +func TestHostRW_Read_PermissionDenied(t *testing.T) { if os.Getuid() == 0 { t.Skip("skipping permission test: running as root") } - tmpDir := t.TempDir() - protected := filepath.Join(tmpDir, "protected.txt") - err := os.WriteFile(protected, []byte("secret"), 0o000) - assert.NoError(t, err) - defer os.Chmod(protected, 0o644) // ensure cleanup _, err = (&hostFs{}).ReadFile(protected) - assert.Error(t, err) - assert.Contains(t, err.Error(), "access denied") } -// TestHostFs_Read_Directory verifies that hostFs.ReadFile returns an error when given a directory path. - -func TestHostFs_Read_Directory(t *testing.T) { +// TestHostRW_Read_Directory verifies that hostRW.Read returns an error when given a directory path. +func TestHostRW_Read_Directory(t *testing.T) { tmpDir := t.TempDir() _, err := (&hostFs{}).ReadFile(tmpDir) - assert.Error(t, err, "expected error when reading a directory as a file") } -// TestSandboxFs_Read_Directory verifies that sandboxFs.ReadFile returns an error when given a directory. - -func TestSandboxFs_Read_Directory(t *testing.T) { +// TestRootRW_Read_Directory verifies that rootRW.Read returns an error when given a directory. +func TestRootRW_Read_Directory(t *testing.T) { workspace := t.TempDir() - root, err := os.OpenRoot(workspace) - assert.NoError(t, err) - defer root.Close() // Create a subdirectory - err = root.Mkdir("subdir", 0o755) - assert.NoError(t, err) _, err = (&sandboxFs{workspace: workspace}).ReadFile("subdir") - assert.Error(t, err, "expected error when reading a directory as a file") } -// TestHostFs_Write_ParentDirMissing verifies that hostFs.WriteFile creates parent dirs automatically. - -func TestHostFs_Write_ParentDirMissing(t *testing.T) { +// TestHostRW_Write_ParentDirMissing verifies that hostRW.Write creates parent dirs automatically. +func TestHostRW_Write_ParentDirMissing(t *testing.T) { tmpDir := t.TempDir() - target := filepath.Join(tmpDir, "a", "b", "c", "file.txt") err := (&hostFs{}).WriteFile(target, []byte("hello")) - assert.NoError(t, err) data, err := os.ReadFile(target) - assert.NoError(t, err) - assert.Equal(t, "hello", string(data)) } -// TestSandboxFs_Write_ParentDirMissing verifies that sandboxFs.WriteFile creates - +// TestRootRW_Write_ParentDirMissing verifies that rootRW.Write creates // nested parent directories automatically within the sandbox. - -func TestSandboxFs_Write_ParentDirMissing(t *testing.T) { +func TestRootRW_Write_ParentDirMissing(t *testing.T) { workspace := t.TempDir() relPath := "x/y/z/file.txt" - err := (&sandboxFs{workspace: workspace}).WriteFile(relPath, []byte("nested")) - assert.NoError(t, err) data, err := os.ReadFile(filepath.Join(workspace, relPath)) - assert.NoError(t, err) - assert.Equal(t, "nested", string(data)) } -// TestHostFs_Write verifies the hostFs.WriteFile helper function - -func TestHostFs_Write(t *testing.T) { +// TestHostRW_Write verifies the hostRW.Write helper function +func TestHostRW_Write(t *testing.T) { tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "atomic_test.txt") - testData := []byte("atomic test content") err := (&hostFs{}).WriteFile(testFile, testData) - assert.NoError(t, err) content, err := os.ReadFile(testFile) - assert.NoError(t, err) - assert.Equal(t, testData, content) // Verify it overwrites correctly - newData := []byte("new atomic content") - err = (&hostFs{}).WriteFile(testFile, newData) - assert.NoError(t, err) content, err = os.ReadFile(testFile) - assert.NoError(t, err) - assert.Equal(t, newData, content) } -// TestSandboxFs_Write verifies the sandboxFs.WriteFile helper function - -func TestSandboxFs_Write(t *testing.T) { +// TestRootRW_Write verifies the rootRW.Write helper function +func TestRootRW_Write(t *testing.T) { tmpDir := t.TempDir() relPath := "atomic_root_test.txt" - testData := []byte("atomic root test content") erw := &sandboxFs{workspace: tmpDir} - err := erw.WriteFile(relPath, testData) - assert.NoError(t, err) root, err := os.OpenRoot(tmpDir) - assert.NoError(t, err) - defer root.Close() f, err := root.Open(relPath) - assert.NoError(t, err) - defer f.Close() content, err := io.ReadAll(f) - assert.NoError(t, err) - assert.Equal(t, testData, content) // Verify it overwrites correctly - newData := []byte("new root atomic content") - err = erw.WriteFile(relPath, newData) - assert.NoError(t, err) f2, err := root.Open(relPath) - assert.NoError(t, err) - defer f2.Close() content, err = io.ReadAll(f2) - assert.NoError(t, err) - assert.Equal(t, newData, content) } -// TestValidatePath_OutsideWorkspace_IncludesPath verifies that the access - -// denied error includes the workspace path so the caller knows the boundary. - -func TestValidatePath_OutsideWorkspace_IncludesPath(t *testing.T) { +// TestWhitelistFs_AllowsMatchingPaths verifies that whitelistFs allows access to +// paths matching the whitelist patterns while blocking non-matching paths. +func TestWhitelistFs_AllowsMatchingPaths(t *testing.T) { workspace := t.TempDir() + outsideDir := t.TempDir() + outsideFile := filepath.Join(outsideDir, "allowed.txt") + os.WriteFile(outsideFile, []byte("outside content"), 0o644) - outsidePath := filepath.Join(t.TempDir(), "secret.txt") + // Pattern allows access to the outsideDir. + patterns := []*regexp.Regexp{regexp.MustCompile(`^` + regexp.QuoteMeta(outsideDir))} - _, err := validatePath(outsidePath, workspace, true) + tool := NewReadFileTool(workspace, true, MaxReadFileSize, patterns) - assert.Error(t, err) + // Read from whitelisted path should succeed. + result := tool.Execute(context.Background(), map[string]any{"path": outsideFile}) + if result.IsError { + t.Errorf("expected whitelisted path to be readable, got: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "outside content") { + t.Errorf("expected file content, got: %s", result.ForLLM) + } - assert.Contains(t, err.Error(), "access denied") + // Read from non-whitelisted path outside workspace should fail. + otherDir := t.TempDir() + otherFile := filepath.Join(otherDir, "blocked.txt") + os.WriteFile(otherFile, []byte("blocked"), 0o644) - assert.Contains(t, err.Error(), workspace) + result = tool.Execute(context.Background(), map[string]any{"path": otherFile}) + if !result.IsError { + t.Errorf("expected non-whitelisted path to be blocked, got: %s", result.ForLLM) + } +} + +// TestReadFileTool_ChunkedReading verifies the pagination logic of the tool +// by reading a file in multiple chunks using 'offset' and 'length'. +func TestReadFileTool_ChunkedReading(t *testing.T) { + tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "pagination_test.txt") + + // Create a test file with exactly 26 bytes of content + fullContent := "abcdefghijklmnopqrstuvwxyz" + err := os.WriteFile(testFile, []byte(fullContent), 0o644) + if err != nil { + t.Fatalf("Failed to write test file: %v", err) + } + + tool := NewReadFileTool(tmpDir, false, MaxReadFileSize) + ctx := context.Background() + + // --- Step 1: Read the first chunk (10 bytes) --- + args1 := map[string]any{ + "path": testFile, + "offset": 0, + "length": 10, + } + result1 := tool.Execute(ctx, args1) + + if result1.IsError { + t.Fatalf("Chunk 1 failed: %s", result1.ForLLM) + } + + // Expect the first 10 characters + if !strings.Contains(result1.ForLLM, "abcdefghij") { + t.Errorf("Chunk 1 should contain 'abcdefghij', got: %s", result1.ForLLM) + } + // Expect the header to indicate the file is truncated + if !strings.Contains(result1.ForLLM, "[TRUNCATED") { + t.Errorf("Chunk 1 header should indicate truncation, got: %s", result1.ForLLM) + } + // Expect the header to suggest the next offset (10) + if !strings.Contains(result1.ForLLM, "offset=10") { + t.Errorf("Chunk 1 header should suggest next offset=10, got: %s", result1.ForLLM) + } + + // Step 2: Read the second chunk (10 bytes) --- + args2 := map[string]any{ + "path": testFile, + "offset": 10, + "length": 10, + } + result2 := tool.Execute(ctx, args2) + + if result2.IsError { + t.Fatalf("Chunk 2 failed: %s", result2.ForLLM) + } + + // Expect the next 10 characters + if !strings.Contains(result2.ForLLM, "klmnopqrst") { + t.Errorf("Chunk 2 should contain 'klmnopqrst', got: %s", result2.ForLLM) + } + // Expect the header to suggest the next offset (20) + if !strings.Contains(result2.ForLLM, "offset=20") { + t.Errorf("Chunk 2 header should suggest next offset=20, got: %s", result2.ForLLM) + } + + // Step 3: Read the final chunk (remaining 6 bytes) --- + // We ask for 10 bytes, but only 6 are left in the file + args3 := map[string]any{ + "path": testFile, + "offset": 20, + "length": 10, + } + result3 := tool.Execute(ctx, args3) + + if result3.IsError { + t.Fatalf("Chunk 3 failed: %s", result3.ForLLM) + } + + // Expect the last 6 characters + if !strings.Contains(result3.ForLLM, "uvwxyz") { + t.Errorf("Chunk 3 should contain 'uvwxyz', got: %s", result3.ForLLM) + } + // Expect the header to indicate the end of the file + if !strings.Contains(result3.ForLLM, "[END OF FILE") { + t.Errorf("Chunk 3 header should indicate end of file, got: %s", result3.ForLLM) + } + + // Ensure no TRUNCATED message is present in the final chunk + if strings.Contains(result3.ForLLM, "[TRUNCATED") { + t.Errorf("Chunk 3 header should NOT indicate truncation, got: %s", result3.ForLLM) + } +} + +// TestReadFileTool_OffsetBeyondEOF checks the behavior when requesting +// An offset that exceeds the total file size. +func TestReadFileTool_OffsetBeyondEOF(t *testing.T) { + tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "short.txt") + + // create a file of only 5 bytes + err := os.WriteFile(testFile, []byte("12345"), 0o644) + if err != nil { + t.Fatalf("Failed to write test file: %v", err) + } + + tool := NewReadFileTool(tmpDir, false, MaxReadFileSize) + ctx := context.Background() + + args := map[string]any{ + "path": testFile, + "offset": int64(100), // Offset beyond the end of the file + } + + result := tool.Execute(ctx, args) + + // It should not be classified as a tool execution error + if result.IsError { + t.Errorf("A mistake was not expected, obtained IsError=true: %s", result.ForLLM) + } + + // Must return EXACTLY the string provided in the code + expectedMsg := "[END OF FILE - no content at this offset]" + if result.ForLLM != expectedMsg { + t.Errorf("The message %q was expected, obtained: %q", expectedMsg, result.ForLLM) + } } diff --git a/pkg/tools/message_test.go b/pkg/tools/message_test.go index f49beab76..05630972e 100644 --- a/pkg/tools/message_test.go +++ b/pkg/tools/message_test.go @@ -9,22 +9,15 @@ import ( func TestMessageTool_Execute_Success(t *testing.T) { tool := NewMessageTool() - tool.SetContext("test-channel", "test-chat-id") - var sentChannel, sentChatID, sentContent string - tool.SetSendCallback(func(channel, chatID, content string) error { sentChannel = channel - sentChatID = chatID - sentContent = content - return nil }) - ctx := context.Background() - + ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id") args := map[string]any{ "content": "Hello, world!", } @@ -32,41 +25,33 @@ func TestMessageTool_Execute_Success(t *testing.T) { result := tool.Execute(ctx, args) // Verify message was sent with correct parameters - if sentChannel != "test-channel" { t.Errorf("Expected channel 'test-channel', got '%s'", sentChannel) } - if sentChatID != "test-chat-id" { t.Errorf("Expected chatID 'test-chat-id', got '%s'", sentChatID) } - if sentContent != "Hello, world!" { t.Errorf("Expected content 'Hello, world!', got '%s'", sentContent) } // Verify ToolResult meets US-011 criteria: - // - Send success returns SilentResult (Silent=true) - if !result.Silent { t.Error("Expected Silent=true for successful send") } // - ForLLM contains send status description - if result.ForLLM != "Message sent to test-channel:test-chat-id" { t.Errorf("Expected ForLLM 'Message sent to test-channel:test-chat-id', got '%s'", result.ForLLM) } // - ForUser is empty (user already received message directly) - if result.ForUser != "" { t.Errorf("Expected ForUser to be empty, got '%s'", result.ForUser) } // - IsError should be false - if result.IsError { t.Error("Expected IsError=false for successful send") } @@ -75,36 +60,26 @@ func TestMessageTool_Execute_Success(t *testing.T) { func TestMessageTool_Execute_WithCustomChannel(t *testing.T) { tool := NewMessageTool() - tool.SetContext("default-channel", "default-chat-id") - var sentChannel, sentChatID string - tool.SetSendCallback(func(channel, chatID, content string) error { sentChannel = channel - sentChatID = chatID - return nil }) - ctx := context.Background() - + ctx := WithToolContext(context.Background(), "default-channel", "default-chat-id") args := map[string]any{ "content": "Test message", - "channel": "custom-channel", - "chat_id": "custom-chat-id", } result := tool.Execute(ctx, args) // Verify custom channel/chatID were used instead of defaults - if sentChannel != "custom-channel" { t.Errorf("Expected channel 'custom-channel', got '%s'", sentChannel) } - if sentChatID != "custom-chat-id" { t.Errorf("Expected chatID 'custom-chat-id', got '%s'", sentChatID) } @@ -112,7 +87,6 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) { if !result.Silent { t.Error("Expected Silent=true") } - if result.ForLLM != "Message sent to custom-channel:custom-chat-id" { t.Errorf("Expected ForLLM 'Message sent to custom-channel:custom-chat-id', got '%s'", result.ForLLM) } @@ -121,16 +95,12 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) { func TestMessageTool_Execute_SendFailure(t *testing.T) { tool := NewMessageTool() - tool.SetContext("test-channel", "test-chat-id") - sendErr := errors.New("network error") - tool.SetSendCallback(func(channel, chatID, content string) error { return sendErr }) - ctx := context.Background() - + ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id") args := map[string]any{ "content": "Test message", } @@ -138,27 +108,21 @@ func TestMessageTool_Execute_SendFailure(t *testing.T) { result := tool.Execute(ctx, args) // Verify ToolResult for send failure: - // - Send failure returns ErrorResult (IsError=true) - if !result.IsError { t.Error("Expected IsError=true for failed send") } // - ForLLM contains error description - expectedErrMsg := "sending message: network error" - if result.ForLLM != expectedErrMsg { t.Errorf("Expected ForLLM '%s', got '%s'", expectedErrMsg, result.ForLLM) } // - Err field should contain original error - if result.Err == nil { t.Error("Expected Err to be set") } - if result.Err != sendErr { t.Errorf("Expected Err to be sendErr, got %v", result.Err) } @@ -167,20 +131,15 @@ func TestMessageTool_Execute_SendFailure(t *testing.T) { func TestMessageTool_Execute_MissingContent(t *testing.T) { tool := NewMessageTool() - tool.SetContext("test-channel", "test-chat-id") - - ctx := context.Background() - + ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id") args := map[string]any{} // content missing result := tool.Execute(ctx, args) // Verify error result for missing content - if !result.IsError { t.Error("Expected IsError=true for missing content") } - if result.ForLLM != "content is required" { t.Errorf("Expected ForLLM 'content is required', got '%s'", result.ForLLM) } @@ -188,15 +147,13 @@ func TestMessageTool_Execute_MissingContent(t *testing.T) { func TestMessageTool_Execute_NoTargetChannel(t *testing.T) { tool := NewMessageTool() - - // No SetContext called, so defaultChannel and defaultChatID are empty + // No WithToolContext — channel/chatID are empty tool.SetSendCallback(func(channel, chatID, content string) error { return nil }) ctx := context.Background() - args := map[string]any{ "content": "Test message", } @@ -204,11 +161,9 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) { result := tool.Execute(ctx, args) // Verify error when no target channel specified - if !result.IsError { t.Error("Expected IsError=true when no target channel") } - if result.ForLLM != "No target channel/chat specified" { t.Errorf("Expected ForLLM 'No target channel/chat specified', got '%s'", result.ForLLM) } @@ -216,13 +171,9 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) { func TestMessageTool_Execute_NotConfigured(t *testing.T) { tool := NewMessageTool() - - tool.SetContext("test-channel", "test-chat-id") - // No SetSendCallback called - ctx := context.Background() - + ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id") args := map[string]any{ "content": "Test message", } @@ -230,11 +181,9 @@ func TestMessageTool_Execute_NotConfigured(t *testing.T) { result := tool.Execute(ctx, args) // Verify error when send callback not configured - if !result.IsError { t.Error("Expected IsError=true when send callback not configured") } - if result.ForLLM != "Message sending not configured" { t.Errorf("Expected ForLLM 'Message sending not configured', got '%s'", result.ForLLM) } @@ -242,7 +191,6 @@ func TestMessageTool_Execute_NotConfigured(t *testing.T) { func TestMessageTool_Name(t *testing.T) { tool := NewMessageTool() - if tool.Name() != "message" { t.Errorf("Expected name 'message', got '%s'", tool.Name()) } @@ -250,9 +198,7 @@ func TestMessageTool_Name(t *testing.T) { func TestMessageTool_Description(t *testing.T) { tool := NewMessageTool() - desc := tool.Description() - if desc == "" { t.Error("Description should not be empty") } @@ -260,63 +206,48 @@ func TestMessageTool_Description(t *testing.T) { func TestMessageTool_Parameters(t *testing.T) { tool := NewMessageTool() - params := tool.Parameters() // Verify parameters structure - typ, ok := params["type"].(string) - if !ok || typ != "object" { t.Error("Expected type 'object'") } props, ok := params["properties"].(map[string]any) - if !ok { t.Fatal("Expected properties to be a map") } // Check required properties - required, ok := params["required"].([]string) - if !ok || len(required) != 1 || required[0] != "content" { t.Error("Expected 'content' to be required") } // Check content property - contentProp, ok := props["content"].(map[string]any) - if !ok { t.Error("Expected 'content' property") } - if contentProp["type"] != "string" { t.Error("Expected content type to be 'string'") } // Check channel property (optional) - channelProp, ok := props["channel"].(map[string]any) - if !ok { t.Error("Expected 'channel' property") } - if channelProp["type"] != "string" { t.Error("Expected channel type to be 'string'") } // Check chat_id property (optional) - chatIDProp, ok := props["chat_id"].(map[string]any) - if !ok { t.Error("Expected 'chat_id' property") } - if chatIDProp["type"] != "string" { t.Error("Expected chat_id type to be 'string'") } diff --git a/pkg/tools/registry_ext_test.go b/pkg/tools/registry_ext_test.go new file mode 100644 index 000000000..cc42d2e4c --- /dev/null +++ b/pkg/tools/registry_ext_test.go @@ -0,0 +1,269 @@ +package tools + +import ( + "context" + "strings" + "testing" +) + +type mockCtxTool struct { + mockRegistryTool + + channel string + + chatID string +} + +func (m *mockCtxTool) SetContext(channel, chatID string) { + m.channel = channel + + m.chatID = chatID +} + +func (m *mockAsyncRegistryTool) SetCallback(cb AsyncCallback) { + m.cb = cb +} + +func TestNormalizeToolName(t *testing.T) { + tests := []struct { + input, want string + }{ + {"read_file", "readfile"}, + + {"readfile", "readfile"}, + + {"ReadFile", "readfile"}, + + {"read-file", "readfile"}, + + {"edit_file", "editfile"}, + + {"web_search", "websearch"}, + + {"EXEC", "exec"}, + } + + for _, tt := range tests { + got := NormalizeToolName(tt.input) + + if got != tt.want { + t.Errorf("NormalizeToolName(%q) = %q, want %q", tt.input, got, tt.want) + } + } +} + +func TestToolRegistry_Get_FuzzyMatch(t *testing.T) { + r := NewToolRegistry() + + r.Register(newMockTool("read_file", "reads a file")) + + r.Register(newMockTool("edit_file", "edits a file")) + + r.Register(newMockTool("web_search", "searches the web")) + + tests := []struct { + query string + + wantName string + }{ + {"readfile", "read_file"}, + + {"ReadFile", "read_file"}, + + {"read-file", "read_file"}, + + {"editfile", "edit_file"}, + + {"EditFile", "edit_file"}, + + {"websearch", "web_search"}, + + {"WebSearch", "web_search"}, + } + + for _, tt := range tests { + tool, ok := r.Get(tt.query) + + if !ok { + t.Errorf("Get(%q) not found, want %q", tt.query, tt.wantName) + + continue + } + + if tool.Name() != tt.wantName { + t.Errorf("Get(%q).Name() = %q, want %q", tt.query, tool.Name(), tt.wantName) + } + } +} + +func TestToolRegistry_ExecuteWithContext_ContextualTool(t *testing.T) { + r := NewToolRegistry() + + ct := &mockCtxTool{ + mockRegistryTool: *newMockTool("ctx_tool", "needs context"), + } + + r.Register(ct) + + r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "telegram", "chat-42", nil) + + if ct.channel != "telegram" { + t.Errorf("expected channel 'telegram', got %q", ct.channel) + } + + if ct.chatID != "chat-42" { + t.Errorf("expected chatID 'chat-42', got %q", ct.chatID) + } +} + +func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) { + r := NewToolRegistry() + + ct := &mockCtxTool{ + mockRegistryTool: *newMockTool("ctx_tool", "needs context"), + } + + r.Register(ct) + + r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "", "", nil) + + if ct.channel != "" || ct.chatID != "" { + t.Error("SetContext should not be called with empty channel/chatID") + } +} + +func TestBuildParamHint(t *testing.T) { + tests := []struct { + name string + + schema map[string]any + + want string + }{ + { + name: "required and optional", + + schema: map[string]any{ + "type": "object", + + "properties": map[string]any{ + "task": map[string]any{"type": "string"}, + + "label": map[string]any{"type": "string"}, + }, + + "required": []string{"task"}, + }, + + want: "(task, label?)", + }, + + { + name: "all required", + + schema: map[string]any{ + "type": "object", + + "properties": map[string]any{ + "command": map[string]any{"type": "string"}, + }, + + "required": []string{"command"}, + }, + + want: "(command)", + }, + + { + name: "no properties", + + schema: map[string]any{ + "type": "object", + }, + + want: "", + }, + + { + name: "empty schema", + + schema: map[string]any{}, + + want: "", + }, + + { + name: "nil schema", + + schema: nil, + + want: "", + }, + + { + name: "multiple optional sorted", + + schema: map[string]any{ + "type": "object", + + "properties": map[string]any{ + "task": map[string]any{"type": "string"}, + + "preset": map[string]any{"type": "string"}, + + "label": map[string]any{"type": "string"}, + + "agent_id": map[string]any{"type": "string"}, + }, + + "required": []string{"task"}, + }, + + want: "(task, agent_id?, label?, preset?)", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := buildParamHint(tt.schema) + + if got != tt.want { + t.Errorf("buildParamHint() = %q, want %q", got, tt.want) + } + }) + } +} + +func TestToolRegistry_GetSummaries_WithParamHint(t *testing.T) { + r := NewToolRegistry() + + r.Register(&mockRegistryTool{ + name: "spawn", + + desc: "Spawn a subagent", + + params: map[string]any{ + "type": "object", + + "properties": map[string]any{ + "task": map[string]any{"type": "string"}, + + "preset": map[string]any{"type": "string"}, + }, + + "required": []string{"task"}, + }, + + result: SilentResult("ok"), + }) + + summaries := r.GetSummaries() + + if len(summaries) != 1 { + t.Fatalf("expected 1 summary, got %d", len(summaries)) + } + + if !strings.Contains(summaries[0], "(task, preset?)") { + t.Errorf("expected param hint in summary, got %q", summaries[0]) + } +} diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go index cff843c3e..92d7d5abd 100644 --- a/pkg/tools/registry_test.go +++ b/pkg/tools/registry_test.go @@ -12,100 +12,57 @@ import ( // --- mock types --- type mockRegistryTool struct { - name string - - desc string - + name string + desc string params map[string]any - result *ToolResult } -func (m *mockRegistryTool) Name() string { return m.name } - -func (m *mockRegistryTool) Description() string { return m.desc } - +func (m *mockRegistryTool) Name() string { return m.name } +func (m *mockRegistryTool) Description() string { return m.desc } func (m *mockRegistryTool) Parameters() map[string]any { return m.params } - func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]any) *ToolResult { return m.result } -type mockCtxTool struct { +type mockContextAwareTool struct { mockRegistryTool - - channel string - - chatID string + lastCtx context.Context } -func (m *mockCtxTool) SetContext(channel, chatID string) { - m.channel = channel - - m.chatID = chatID +func (m *mockContextAwareTool) Execute(ctx context.Context, _ map[string]any) *ToolResult { + m.lastCtx = ctx + return m.result } type mockAsyncRegistryTool struct { mockRegistryTool - - cb AsyncCallback + lastCB AsyncCallback } -func (m *mockAsyncRegistryTool) SetCallback(cb AsyncCallback) { - m.cb = cb +func (m *mockAsyncRegistryTool) ExecuteAsync(_ context.Context, args map[string]any, cb AsyncCallback) *ToolResult { + m.lastCB = cb + return m.result } // --- helpers --- func newMockTool(name, desc string) *mockRegistryTool { return &mockRegistryTool{ - name: name, - - desc: desc, - + name: name, + desc: desc, params: map[string]any{"type": "object"}, - result: SilentResult("ok"), } } // --- tests --- -func TestNormalizeToolName(t *testing.T) { - tests := []struct { - input, want string - }{ - {"read_file", "readfile"}, - - {"readfile", "readfile"}, - - {"ReadFile", "readfile"}, - - {"read-file", "readfile"}, - - {"edit_file", "editfile"}, - - {"web_search", "websearch"}, - - {"EXEC", "exec"}, - } - - for _, tt := range tests { - got := NormalizeToolName(tt.input) - - if got != tt.want { - t.Errorf("NormalizeToolName(%q) = %q, want %q", tt.input, got, tt.want) - } - } -} - func TestNewToolRegistry(t *testing.T) { r := NewToolRegistry() - if r.Count() != 0 { t.Errorf("expected empty registry, got count %d", r.Count()) } - if len(r.List()) != 0 { t.Errorf("expected empty list, got %v", r.List()) } @@ -113,17 +70,13 @@ func TestNewToolRegistry(t *testing.T) { func TestToolRegistry_RegisterAndGet(t *testing.T) { r := NewToolRegistry() - tool := newMockTool("echo", "echoes input") - r.Register(tool) got, ok := r.Get("echo") - if !ok { t.Fatal("expected to find registered tool") } - if got.Name() != "echo" { t.Errorf("expected name 'echo', got %q", got.Name()) } @@ -131,71 +84,21 @@ func TestToolRegistry_RegisterAndGet(t *testing.T) { func TestToolRegistry_Get_NotFound(t *testing.T) { r := NewToolRegistry() - _, ok := r.Get("nonexistent") - if ok { t.Error("expected ok=false for unregistered tool") } } -func TestToolRegistry_Get_FuzzyMatch(t *testing.T) { - r := NewToolRegistry() - - r.Register(newMockTool("read_file", "reads a file")) - - r.Register(newMockTool("edit_file", "edits a file")) - - r.Register(newMockTool("web_search", "searches the web")) - - tests := []struct { - query string - - wantName string - }{ - {"readfile", "read_file"}, - - {"ReadFile", "read_file"}, - - {"read-file", "read_file"}, - - {"editfile", "edit_file"}, - - {"EditFile", "edit_file"}, - - {"websearch", "web_search"}, - - {"WebSearch", "web_search"}, - } - - for _, tt := range tests { - tool, ok := r.Get(tt.query) - - if !ok { - t.Errorf("Get(%q) not found, want %q", tt.query, tt.wantName) - - continue - } - - if tool.Name() != tt.wantName { - t.Errorf("Get(%q).Name() = %q, want %q", tt.query, tool.Name(), tt.wantName) - } - } -} - func TestToolRegistry_RegisterOverwrite(t *testing.T) { r := NewToolRegistry() - r.Register(newMockTool("dup", "first")) - r.Register(newMockTool("dup", "second")) if r.Count() != 1 { t.Errorf("expected count 1 after overwrite, got %d", r.Count()) } - tool, _ := r.Get("dup") - if tool.Description() != "second" { t.Errorf("expected overwritten description 'second', got %q", tool.Description()) } @@ -203,23 +106,17 @@ func TestToolRegistry_RegisterOverwrite(t *testing.T) { func TestToolRegistry_Execute_Success(t *testing.T) { r := NewToolRegistry() - r.Register(&mockRegistryTool{ - name: "greet", - - desc: "says hello", - + name: "greet", + desc: "says hello", params: map[string]any{}, - result: SilentResult("hello"), }) result := r.Execute(context.Background(), "greet", nil) - if result.IsError { t.Errorf("expected success, got error: %s", result.ForLLM) } - if result.ForLLM != "hello" { t.Errorf("expected ForLLM 'hello', got %q", result.ForLLM) } @@ -227,85 +124,79 @@ func TestToolRegistry_Execute_Success(t *testing.T) { func TestToolRegistry_Execute_NotFound(t *testing.T) { r := NewToolRegistry() - result := r.Execute(context.Background(), "missing", nil) - if !result.IsError { t.Error("expected error for missing tool") } - if !strings.Contains(result.ForLLM, "not found") { t.Errorf("expected 'not found' in error, got %q", result.ForLLM) } - if result.Err == nil { t.Error("expected Err to be set via WithError") } } -func TestToolRegistry_ExecuteWithContext_ContextualTool(t *testing.T) { +func TestToolRegistry_ExecuteWithContext_InjectsToolContext(t *testing.T) { r := NewToolRegistry() - - ct := &mockCtxTool{ + ct := &mockContextAwareTool{ mockRegistryTool: *newMockTool("ctx_tool", "needs context"), } - r.Register(ct) r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "telegram", "chat-42", nil) - if ct.channel != "telegram" { - t.Errorf("expected channel 'telegram', got %q", ct.channel) + if ct.lastCtx == nil { + t.Fatal("expected Execute to be called") } - - if ct.chatID != "chat-42" { - t.Errorf("expected chatID 'chat-42', got %q", ct.chatID) + if got := ToolChannel(ct.lastCtx); got != "telegram" { + t.Errorf("expected channel 'telegram', got %q", got) + } + if got := ToolChatID(ct.lastCtx); got != "chat-42" { + t.Errorf("expected chatID 'chat-42', got %q", got) } } -func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) { +func TestToolRegistry_ExecuteWithContext_EmptyContext(t *testing.T) { r := NewToolRegistry() - - ct := &mockCtxTool{ + ct := &mockContextAwareTool{ mockRegistryTool: *newMockTool("ctx_tool", "needs context"), } - r.Register(ct) r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "", "", nil) - if ct.channel != "" || ct.chatID != "" { - t.Error("SetContext should not be called with empty channel/chatID") + if ct.lastCtx == nil { + t.Fatal("expected Execute to be called") + } + // Empty values are still injected; tools decide what to do with them. + if got := ToolChannel(ct.lastCtx); got != "" { + t.Errorf("expected empty channel, got %q", got) + } + if got := ToolChatID(ct.lastCtx); got != "" { + t.Errorf("expected empty chatID, got %q", got) } } func TestToolRegistry_ExecuteWithContext_AsyncCallback(t *testing.T) { r := NewToolRegistry() - at := &mockAsyncRegistryTool{ mockRegistryTool: *newMockTool("async_tool", "async work"), } - at.result = AsyncResult("started") - r.Register(at) called := false - cb := func(_ context.Context, _ *ToolResult) { called = true } result := r.ExecuteWithContext(context.Background(), "async_tool", nil, "", "", cb) - - if at.cb == nil { - t.Error("expected SetCallback to have been called") + if at.lastCB == nil { + t.Error("expected ExecuteAsync to have received a callback") } - if !result.Async { t.Error("expected async result") } - at.cb(context.Background(), SilentResult("done")) - + at.lastCB(context.Background(), SilentResult("done")) if !called { t.Error("expected callback to be invoked") } @@ -313,29 +204,22 @@ func TestToolRegistry_ExecuteWithContext_AsyncCallback(t *testing.T) { func TestToolRegistry_GetDefinitions(t *testing.T) { r := NewToolRegistry() - r.Register(newMockTool("alpha", "tool A")) defs := r.GetDefinitions() - if len(defs) != 1 { t.Fatalf("expected 1 definition, got %d", len(defs)) } - if defs[0]["type"] != "function" { t.Errorf("expected type 'function', got %v", defs[0]["type"]) } - fn, ok := defs[0]["function"].(map[string]any) - if !ok { t.Fatal("expected 'function' key to be a map") } - if fn["name"] != "alpha" { t.Errorf("expected name 'alpha', got %v", fn["name"]) } - if fn["description"] != "tool A" { t.Errorf("expected description 'tool A', got %v", fn["description"]) } @@ -343,47 +227,34 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { func TestToolRegistry_ToProviderDefs(t *testing.T) { r := NewToolRegistry() - params := map[string]any{"type": "object", "properties": map[string]any{}} - r.Register(&mockRegistryTool{ - name: "beta", - - desc: "tool B", - + name: "beta", + desc: "tool B", params: params, - result: SilentResult("ok"), }) defs := r.ToProviderDefs() - if len(defs) != 1 { t.Fatalf("expected 1 provider def, got %d", len(defs)) } want := providers.ToolDefinition{ Type: "function", - Function: providers.ToolFunctionDefinition{ - Name: "beta", - + Name: "beta", Description: "tool B", - - Parameters: providers.MustMarshalParameters(params), + Parameters: params, }, } - got := defs[0] - if got.Type != want.Type { t.Errorf("Type: want %q, got %q", want.Type, got.Type) } - if got.Function.Name != want.Function.Name { t.Errorf("Name: want %q, got %q", want.Function.Name, got.Function.Name) } - if got.Function.Description != want.Function.Description { t.Errorf("Description: want %q, got %q", want.Function.Description, got.Function.Description) } @@ -391,23 +262,18 @@ func TestToolRegistry_ToProviderDefs(t *testing.T) { func TestToolRegistry_List(t *testing.T) { r := NewToolRegistry() - r.Register(newMockTool("x", "")) - r.Register(newMockTool("y", "")) names := r.List() - if len(names) != 2 { t.Fatalf("expected 2 names, got %d", len(names)) } nameSet := map[string]bool{} - for _, n := range names { nameSet[n] = true } - if !nameSet["x"] || !nameSet["y"] { t.Errorf("expected names {x, y}, got %v", names) } @@ -415,207 +281,55 @@ func TestToolRegistry_List(t *testing.T) { func TestToolRegistry_Count(t *testing.T) { r := NewToolRegistry() - if r.Count() != 0 { t.Errorf("expected 0, got %d", r.Count()) } r.Register(newMockTool("a", "")) - r.Register(newMockTool("b", "")) - if r.Count() != 2 { t.Errorf("expected 2, got %d", r.Count()) } r.Register(newMockTool("a", "replaced")) - if r.Count() != 2 { t.Errorf("expected 2 after overwrite, got %d", r.Count()) } } -func TestBuildParamHint(t *testing.T) { - tests := []struct { - name string - - schema map[string]any - - want string - }{ - { - name: "required and optional", - - schema: map[string]any{ - "type": "object", - - "properties": map[string]any{ - "task": map[string]any{"type": "string"}, - - "label": map[string]any{"type": "string"}, - }, - - "required": []string{"task"}, - }, - - want: "(task, label?)", - }, - - { - name: "all required", - - schema: map[string]any{ - "type": "object", - - "properties": map[string]any{ - "command": map[string]any{"type": "string"}, - }, - - "required": []string{"command"}, - }, - - want: "(command)", - }, - - { - name: "no properties", - - schema: map[string]any{ - "type": "object", - }, - - want: "", - }, - - { - name: "empty schema", - - schema: map[string]any{}, - - want: "", - }, - - { - name: "nil schema", - - schema: nil, - - want: "", - }, - - { - name: "multiple optional sorted", - - schema: map[string]any{ - "type": "object", - - "properties": map[string]any{ - "task": map[string]any{"type": "string"}, - - "preset": map[string]any{"type": "string"}, - - "label": map[string]any{"type": "string"}, - - "agent_id": map[string]any{"type": "string"}, - }, - - "required": []string{"task"}, - }, - - want: "(task, agent_id?, label?, preset?)", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := buildParamHint(tt.schema) - - if got != tt.want { - t.Errorf("buildParamHint() = %q, want %q", got, tt.want) - } - }) - } -} - func TestToolRegistry_GetSummaries(t *testing.T) { r := NewToolRegistry() - r.Register(newMockTool("read_file", "Reads a file")) summaries := r.GetSummaries() - if len(summaries) != 1 { t.Fatalf("expected 1 summary, got %d", len(summaries)) } - if !strings.Contains(summaries[0], "`read_file`") { t.Errorf("expected backtick-quoted name in summary, got %q", summaries[0]) } - if !strings.Contains(summaries[0], "Reads a file") { t.Errorf("expected description in summary, got %q", summaries[0]) } } -func TestToolRegistry_GetSummaries_WithParamHint(t *testing.T) { - r := NewToolRegistry() - - r.Register(&mockRegistryTool{ - name: "spawn", - - desc: "Spawn a subagent", - - params: map[string]any{ - "type": "object", - - "properties": map[string]any{ - "task": map[string]any{"type": "string"}, - - "preset": map[string]any{"type": "string"}, - }, - - "required": []string{"task"}, - }, - - result: SilentResult("ok"), - }) - - summaries := r.GetSummaries() - - if len(summaries) != 1 { - t.Fatalf("expected 1 summary, got %d", len(summaries)) - } - - // Should contain param hint - - if !strings.Contains(summaries[0], "(task, preset?)") { - t.Errorf("expected param hint in summary, got %q", summaries[0]) - } -} - func TestToolToSchema(t *testing.T) { tool := newMockTool("demo", "demo tool") - schema := ToolToSchema(tool) if schema["type"] != "function" { t.Errorf("expected type 'function', got %v", schema["type"]) } - fn, ok := schema["function"].(map[string]any) - if !ok { t.Fatal("expected 'function' to be a map") } - if fn["name"] != "demo" { t.Errorf("expected name 'demo', got %v", fn["name"]) } - if fn["description"] != "demo tool" { t.Errorf("expected description 'demo tool', got %v", fn["description"]) } - if fn["parameters"] == nil { t.Error("expected parameters to be set") } @@ -623,25 +337,17 @@ func TestToolToSchema(t *testing.T) { func TestToolRegistry_ConcurrentAccess(t *testing.T) { r := NewToolRegistry() - var wg sync.WaitGroup for i := range 50 { wg.Add(1) - go func(n int) { defer wg.Done() - name := string(rune('A' + n%26)) - r.Register(newMockTool(name, "concurrent")) - r.Get(name) - r.Count() - r.List() - r.GetDefinitions() }(i) } diff --git a/pkg/tools/result_test.go b/pkg/tools/result_test.go index ac7d1bfda..a234e33f3 100644 --- a/pkg/tools/result_test.go +++ b/pkg/tools/result_test.go @@ -12,15 +12,12 @@ func TestNewToolResult(t *testing.T) { if result.ForLLM != "test content" { t.Errorf("Expected ForLLM 'test content', got '%s'", result.ForLLM) } - if result.Silent { t.Error("Expected Silent to be false") } - if result.IsError { t.Error("Expected IsError to be false") } - if result.Async { t.Error("Expected Async to be false") } @@ -32,15 +29,12 @@ func TestSilentResult(t *testing.T) { if result.ForLLM != "silent operation" { t.Errorf("Expected ForLLM 'silent operation', got '%s'", result.ForLLM) } - if !result.Silent { t.Error("Expected Silent to be true") } - if result.IsError { t.Error("Expected IsError to be false") } - if result.Async { t.Error("Expected Async to be false") } @@ -52,15 +46,12 @@ func TestAsyncResult(t *testing.T) { if result.ForLLM != "async task started" { t.Errorf("Expected ForLLM 'async task started', got '%s'", result.ForLLM) } - if result.Silent { t.Error("Expected Silent to be false") } - if result.IsError { t.Error("Expected IsError to be false") } - if !result.Async { t.Error("Expected Async to be true") } @@ -72,15 +63,12 @@ func TestErrorResult(t *testing.T) { if result.ForLLM != "operation failed" { t.Errorf("Expected ForLLM 'operation failed', got '%s'", result.ForLLM) } - if result.Silent { t.Error("Expected Silent to be false") } - if !result.IsError { t.Error("Expected IsError to be true") } - if result.Async { t.Error("Expected Async to be false") } @@ -88,25 +76,20 @@ func TestErrorResult(t *testing.T) { func TestUserResult(t *testing.T) { content := "user visible message" - result := UserResult(content) if result.ForLLM != content { t.Errorf("Expected ForLLM '%s', got '%s'", content, result.ForLLM) } - if result.ForUser != content { t.Errorf("Expected ForUser '%s', got '%s'", content, result.ForUser) } - if result.Silent { t.Error("Expected Silent to be false") } - if result.IsError { t.Error("Expected IsError to be false") } - if result.Async { t.Error("Expected Async to be false") } @@ -114,37 +97,27 @@ func TestUserResult(t *testing.T) { func TestToolResultJSONSerialization(t *testing.T) { tests := []struct { - name string - + name string result *ToolResult }{ { - name: "basic result", - + name: "basic result", result: NewToolResult("basic content"), }, - { - name: "silent result", - + name: "silent result", result: SilentResult("silent content"), }, - { - name: "async result", - + name: "async result", result: AsyncResult("async content"), }, - { - name: "error result", - + name: "error result", result: ErrorResult("error content"), }, - { - name: "user result", - + name: "user result", result: UserResult("user content"), }, } @@ -152,38 +125,30 @@ func TestToolResultJSONSerialization(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Marshal to JSON - data, err := json.Marshal(tt.result) if err != nil { t.Fatalf("Failed to marshal: %v", err) } // Unmarshal back - var decoded ToolResult - if err := json.Unmarshal(data, &decoded); err != nil { t.Fatalf("Failed to unmarshal: %v", err) } // Verify fields match (Err should be excluded) - if decoded.ForLLM != tt.result.ForLLM { t.Errorf("ForLLM mismatch: got '%s', want '%s'", decoded.ForLLM, tt.result.ForLLM) } - if decoded.ForUser != tt.result.ForUser { t.Errorf("ForUser mismatch: got '%s', want '%s'", decoded.ForUser, tt.result.ForUser) } - if decoded.Silent != tt.result.Silent { t.Errorf("Silent mismatch: got %v, want %v", decoded.Silent, tt.result.Silent) } - if decoded.IsError != tt.result.IsError { t.Errorf("IsError mismatch: got %v, want %v", decoded.IsError, tt.result.IsError) } - if decoded.Async != tt.result.Async { t.Errorf("Async mismatch: got %v, want %v", decoded.Async, tt.result.Async) } @@ -193,27 +158,22 @@ func TestToolResultJSONSerialization(t *testing.T) { func TestToolResultWithErrors(t *testing.T) { err := errors.New("underlying error") - result := ErrorResult("error message").WithError(err) if result.Err == nil { t.Error("Expected Err to be set") } - if result.Err.Error() != "underlying error" { t.Errorf("Expected Err message 'underlying error', got '%s'", result.Err.Error()) } // Verify Err is not serialized - data, marshalErr := json.Marshal(result) - if marshalErr != nil { t.Fatalf("Failed to marshal: %v", marshalErr) } var decoded ToolResult - if unmarshalErr := json.Unmarshal(data, &decoded); unmarshalErr != nil { t.Fatalf("Failed to unmarshal: %v", unmarshalErr) } @@ -232,47 +192,37 @@ func TestToolResultJSONStructure(t *testing.T) { } // Verify JSON structure - var parsed map[string]any - if err := json.Unmarshal(data, &parsed); err != nil { t.Fatalf("Failed to parse JSON: %v", err) } // Check expected keys exist - if _, ok := parsed["for_llm"]; !ok { t.Error("Expected 'for_llm' key in JSON") } - if _, ok := parsed["for_user"]; !ok { t.Error("Expected 'for_user' key in JSON") } - if _, ok := parsed["silent"]; !ok { t.Error("Expected 'silent' key in JSON") } - if _, ok := parsed["is_error"]; !ok { t.Error("Expected 'is_error' key in JSON") } - if _, ok := parsed["async"]; !ok { t.Error("Expected 'async' key in JSON") } // Check that 'err' is NOT present (it should have json:"-" tag) - if _, ok := parsed["err"]; ok { t.Error("Expected 'err' key to be excluded from JSON") } // Verify values - if parsed["for_llm"] != "test content" { t.Errorf("Expected for_llm 'test content', got %v", parsed["for_llm"]) } - if parsed["silent"] != false { t.Errorf("Expected silent false, got %v", parsed["silent"]) } diff --git a/pkg/tools/shell_ext_test.go b/pkg/tools/shell_ext_test.go new file mode 100644 index 000000000..0f8352de9 --- /dev/null +++ b/pkg/tools/shell_ext_test.go @@ -0,0 +1,896 @@ +package tools + +import ( + "context" + "os" + "path/filepath" + "regexp" + "runtime" + "strings" + "testing" + "time" +) + +func TestGuardCommand_RelativePathWithSlashes(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + cmds := []string{ + "pytest tests/cold/test_solver.py -v --tb=short", + + "cd projects/terra-py-form && pytest", + + "uv run pytest tests/cold/test_solver.py -v --tb=short", + + "cat src/terra_py_form/cold/parser.py", + + "python src/main.py --config config/dev.json", + } + + for _, cmd := range cmds { + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("Relative path should not be blocked: %q → %s", cmd, result) + } + } +} + +func TestGuardCommand_VenvBinary(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + cmds := []string{ + ".venv/bin/python -m pytest", + + ".venv/bin/pytest tests/ -v", + + ".venv/bin/pip install -e .", + } + + for _, cmd := range cmds { + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("Venv relative path should not be blocked: %q → %s", cmd, result) + } + } +} + +func TestGuardCommand_ExecutableBinaryAllowed(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix executable permission test not applicable on Windows") + } + + workspace := t.TempDir() + + externalDir := t.TempDir() + + execPath := filepath.Join(externalDir, "mybin") + + os.WriteFile(execPath, []byte("#!/bin/sh\necho ok"), 0o755) + + tool, _ := NewExecTool(workspace, true) + + cmd := execPath + " --help" + + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("Executable binary outside workspace should be allowed: %q → %s", cmd, result) + } +} + +func TestGuardCommand_ExecutableBinaryAllowed_Windows(t *testing.T) { + if runtime.GOOS != "windows" { + t.Skip("Windows-specific test") + } + + workspace := t.TempDir() + + externalDir := t.TempDir() + + execPath := filepath.Join(externalDir, "tool.exe") + + os.WriteFile(execPath, []byte("MZ"), 0o644) + + tool, _ := NewExecTool(workspace, true) + + cmd := execPath + " --version" + + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("Windows .exe outside workspace should be allowed: %q → %s", cmd, result) + } +} + +func TestGuardCommand_NonExecutableOutsideBlocked(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix permission test not applicable on Windows") + } + + workspace := t.TempDir() + + externalDir := t.TempDir() + + dataFile := filepath.Join(externalDir, "secret.txt") + + os.WriteFile(dataFile, []byte("secret data"), 0o644) + + tool, _ := NewExecTool(workspace, true) + + cmd := "cat " + dataFile + + result := tool.guardCommand(cmd, workspace) + + if result == "" { + t.Errorf("Non-executable file outside workspace should be blocked: %q", cmd) + } + + if !strings.Contains(result, "path outside working dir") { + t.Errorf("Expected 'path outside working dir' message, got: %s", result) + } +} + +func TestGuardCommand_NonExistentAbsolutePathBlocked(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "echo hello > C:\\nonexistent_picoclaw_test_output" + } else { + cmd = "echo hello > /tmp/nonexistent_picoclaw_test_output" + } + + result := tool.guardCommand(cmd, workspace) + + if result == "" { + t.Errorf("Non-existent absolute path outside workspace should be blocked: %q", cmd) + } +} + +func TestGuardCommand_FlagEmbeddedPathSkipped(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + cmds := []string{ + "gcc -I/usr/local/include -L/usr/lib main.c", + + "g++ -std=c++17 -I/opt/include file.cpp", + + "python --prefix=/usr/local script.py", + } + + for _, cmd := range cmds { + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("Flag-embedded path should not be blocked: %q → %s", cmd, result) + } + } +} + +func TestGuardCommand_AbsolutePathInsideWorkspace(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + innerDir := filepath.Join(workspace, "projects", "myapp") + + os.MkdirAll(innerDir, 0o755) + + cmd := "ls " + innerDir + + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("Absolute path inside workspace should be allowed: %q → %s", cmd, result) + } +} + +func TestGuardCommand_PathTraversal(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + cmds := []string{ + "cat ../../etc/passwd", + + "cat ../../../etc/shadow", + + "ls projects/../../../../etc", + } + + for _, cmd := range cmds { + result := tool.guardCommand(cmd, workspace) + + if result == "" { + t.Errorf("Path traversal should be blocked: %q", cmd) + } + + if !strings.Contains(result, "path traversal") { + t.Errorf("Expected 'path traversal' message, got: %s", result) + } + } +} + +func TestGuardCommand_CdWithAbsoluteWorkspacePath(t *testing.T) { + workspace := t.TempDir() + + innerDir := filepath.Join(workspace, "projects", "foo") + + os.MkdirAll(innerDir, 0o755) + + tool, _ := NewExecTool(workspace, true) + + cmd := "cd " + innerDir + " && ls -la" + + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("cd to workspace subdir should be allowed: %q → %s", cmd, result) + } +} + +func TestGuardCommand_AgentCLISlashCommand(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + cmds := []string{ + `codex exec --yolo "/review skip-git-repo-check"`, + + `claude "/review"`, + + `gemini "/help"`, + } + + for _, cmd := range cmds { + result := tool.guardCommand(cmd, workspace) + + if result != "" { + t.Errorf("Agent CLI slash command should not be blocked: %q → %s", cmd, result) + } + } + + if runtime.GOOS != "windows" { + blocked := `cat /etc/hosts` + + result := tool.guardCommand(blocked, workspace) + + if result == "" { + t.Errorf("Non-agent command with absolute path should be blocked: %q", blocked) + } + } +} + +func TestGuardCommand_DenyPattern_IncludesPattern(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + tool.denyPatterns = append(tool.denyPatterns, regexp.MustCompile(`\bdangerous_cmd\b`)) + + result := tool.guardCommand("dangerous_cmd --force", workspace) + + if result == "" { + t.Fatal("expected deny pattern to block the command") + } + + if !strings.Contains(result, "deny pattern") { + t.Errorf("expected 'deny pattern' in message, got: %s", result) + } + + if !strings.Contains(result, `\bdangerous_cmd\b`) { + t.Errorf("expected pattern string in message, got: %s", result) + } +} + +func TestGuardCommand_Allowlist_ShowsRules(t *testing.T) { + workspace := t.TempDir() + + tool, _ := NewExecTool(workspace, true) + + tool.SetAllowRules([]string{"go test", "git"}) + + result := tool.guardCommand("curl http://example.com", workspace) + + if result == "" { + t.Fatal("expected allowlist to block the command") + } + + if !strings.Contains(result, "not in allowlist") { + t.Errorf("expected 'not in allowlist' in message, got: %s", result) + } + + if !strings.Contains(result, "go test") || !strings.Contains(result, "git") { + t.Errorf("expected allowlist rules in message, got: %s", result) + } +} + +func TestGuardCommand_PathOutside_IncludesPath(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Unix absolute path test not applicable on Windows") + } + + workspace := t.TempDir() + + externalDir := t.TempDir() + + dataFile := filepath.Join(externalDir, "secret.txt") + + os.WriteFile(dataFile, []byte("secret"), 0o644) + + tool, _ := NewExecTool(workspace, true) + + result := tool.guardCommand("cat "+dataFile, workspace) + + if result == "" { + t.Fatal("expected path outside workspace to be blocked") + } + + if !strings.Contains(result, "path outside working dir") { + t.Errorf("expected 'path outside working dir' in message, got: %s", result) + } + + if !strings.Contains(result, dataFile) { + t.Errorf("expected offending path %q in message, got: %s", dataFile, result) + } +} + +func TestExecTool_Bg_StartAndOutput(t *testing.T) { + tool, _ := NewExecTool("", false) + + defer tool.Shutdown() + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "Write-Output 'hello from bg'; Start-Sleep -Seconds 30" + } else { + cmd = "echo 'hello from bg'; sleep 30" + } + + result := tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + if result.IsError { + t.Fatalf("failed to start bg process: %s", result.ForLLM) + } + + if !strings.Contains(result.ForLLM, "bg-1") { + t.Errorf("expected bg-1 in result, got: %s", result.ForLLM) + } + + if !strings.Contains(result.ForLLM, "Background process started") { + t.Errorf("expected start message, got: %s", result.ForLLM) + } + + outputResult := tool.Execute(context.Background(), map[string]any{ + "bg_action": "output", + + "bg_id": "bg-1", + }) + + if outputResult.IsError { + t.Fatalf("failed to get output: %s", outputResult.ForLLM) + } + + if !strings.Contains(outputResult.ForLLM, "hello from bg") { + t.Errorf("expected 'hello from bg' in output, got: %s", outputResult.ForLLM) + } + + if !strings.Contains(outputResult.ForLLM, "running") { + t.Errorf("expected 'running' status, got: %s", outputResult.ForLLM) + } +} + +func TestExecTool_Bg_Kill(t *testing.T) { + tool, _ := NewExecTool("", false) + + defer tool.Shutdown() + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "Start-Sleep -Seconds 60" + } else { + cmd = "sleep 60" + } + + result := tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + if result.IsError { + t.Fatalf("failed to start bg process: %s", result.ForLLM) + } + + killResult := tool.Execute(context.Background(), map[string]any{ + "bg_action": "kill", + + "bg_id": "bg-1", + }) + + if killResult.IsError { + t.Fatalf("failed to kill: %s", killResult.ForLLM) + } + + if !strings.Contains(killResult.ForLLM, "terminated") { + t.Errorf("expected 'terminated' message, got: %s", killResult.ForLLM) + } + + procs := tool.BgProcesses() + + if _, ok := procs["bg-1"]; ok { + t.Errorf("expected bg-1 to be removed after kill") + } +} + +func TestExecTool_Bg_ExitedProcess(t *testing.T) { + tool, _ := NewExecTool("", false) + + defer tool.Shutdown() + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "Write-Output 'quick exit'" + } else { + cmd = "echo 'quick exit'" + } + + result := tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + if result.IsError { + t.Fatalf("failed to start bg process: %s", result.ForLLM) + } + + time.Sleep(4 * time.Second) + + outputResult := tool.Execute(context.Background(), map[string]any{ + "bg_action": "output", + + "bg_id": "bg-1", + }) + + if outputResult.IsError { + t.Fatalf("failed to get output: %s", outputResult.ForLLM) + } + + if !strings.Contains(outputResult.ForLLM, "exited") { + t.Errorf("expected 'exited' in output, got: %s", outputResult.ForLLM) + } + + if !strings.Contains(outputResult.ForLLM, "quick exit") { + t.Errorf("expected 'quick exit' in output, got: %s", outputResult.ForLLM) + } +} + +func TestExecTool_Bg_InvalidID(t *testing.T) { + tool, _ := NewExecTool("", false) + + defer tool.Shutdown() + + result := tool.Execute(context.Background(), map[string]any{ + "bg_action": "output", + + "bg_id": "bg-999", + }) + + if !result.IsError { + t.Fatalf("expected error for invalid bg_id") + } + + if !strings.Contains(result.ForLLM, "not found") { + t.Errorf("expected 'not found' message, got: %s", result.ForLLM) + } + + result = tool.Execute(context.Background(), map[string]any{ + "bg_action": "kill", + + "bg_id": "bg-999", + }) + + if !result.IsError { + t.Fatalf("expected error for invalid bg_id") + } +} + +func TestExecTool_Bg_InitialOutputCapture(t *testing.T) { + tool, _ := NewExecTool("", false) + + defer tool.Shutdown() + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "Write-Output 'initial line 1'; Write-Output 'initial line 2'; Start-Sleep -Seconds 30" + } else { + cmd = "echo 'initial line 1'; echo 'initial line 2'; sleep 30" + } + + result := tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + if result.IsError { + t.Fatalf("failed to start bg process: %s", result.ForLLM) + } + + if !strings.Contains(result.ForLLM, "initial line 1") { + t.Errorf("expected 'initial line 1' in initial output, got: %s", result.ForLLM) + } + + if !strings.Contains(result.ForLLM, "initial line 2") { + t.Errorf("expected 'initial line 2' in initial output, got: %s", result.ForLLM) + } +} + +func TestExecTool_Bg_RuntimeStatus(t *testing.T) { + tool, _ := NewExecTool("", false) + + defer tool.Shutdown() + + if s := tool.RuntimeStatus(); s != "" { + t.Errorf("expected empty runtime status with no bg processes, got: %s", s) + } + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "Start-Sleep -Seconds 30" + } else { + cmd = "sleep 30" + } + + tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + status := tool.RuntimeStatus() + + if !strings.Contains(status, "Background Processes") { + t.Errorf("expected 'Background Processes' section, got: %s", status) + } + + if !strings.Contains(status, "bg-1") { + t.Errorf("expected 'bg-1' in status, got: %s", status) + } + + if !strings.Contains(status, "running") { + t.Errorf("expected 'running' in status, got: %s", status) + } +} + +func TestExecTool_Bg_Shutdown(t *testing.T) { + tool, _ := NewExecTool("", false) + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "Start-Sleep -Seconds 60" + } else { + cmd = "sleep 60" + } + + tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + procs := tool.BgProcesses() + + for _, bp := range procs { + if !bp.isRunning() { + t.Errorf("expected process to be running before shutdown") + } + } + + tool.Shutdown() + + procs = tool.BgProcesses() + + for _, bp := range procs { + if bp.isRunning() { + t.Errorf("expected process to be stopped after shutdown") + } + } +} + +func TestRingBuffer(t *testing.T) { + t.Run("Write and String", func(t *testing.T) { + rb := newRingBuffer(100) + + rb.Write([]byte("hello ")) + + rb.Write([]byte("world")) + + if got := rb.String(); got != "hello world" { + t.Errorf("expected 'hello world', got %q", got) + } + }) + + t.Run("Lines", func(t *testing.T) { + rb := newRingBuffer(100) + + rb.Write([]byte("line1\nline2\nline3\nline4\nline5\n")) + + lines := rb.Lines(3) + + if len(lines) != 3 { + t.Fatalf("expected 3 lines, got %d", len(lines)) + } + + if lines[0] != "line3" || lines[1] != "line4" || lines[2] != "line5" { + t.Errorf("unexpected lines: %v", lines) + } + }) + + t.Run("Match", func(t *testing.T) { + rb := newRingBuffer(100) + + rb.Write([]byte("starting...\nServer ready on port 3000\nwaiting...\n")) + + re := regexp.MustCompile(`ready.*port`) + + match := rb.Match(re) + + if match == "" { + t.Fatal("expected match but got empty string") + } + + if !strings.Contains(match, "ready") { + t.Errorf("expected match to contain 'ready', got: %s", match) + } + + re2 := regexp.MustCompile(`never_match`) + + match2 := rb.Match(re2) + + if match2 != "" { + t.Errorf("expected no match, got: %s", match2) + } + }) + + t.Run("Overflow", func(t *testing.T) { + rb := newRingBuffer(10) + + rb.Write([]byte("1234567890ABCDEF")) + + got := rb.String() + + if len(got) != 10 { + t.Errorf("expected buffer to be 10 bytes, got %d", len(got)) + } + + if got != "7890ABCDEF" { + t.Errorf("expected '7890ABCDEF', got %q", got) + } + }) + + t.Run("Len", func(t *testing.T) { + rb := newRingBuffer(100) + + if rb.Len() != 0 { + t.Errorf("expected 0 length initially") + } + + rb.Write([]byte("hello")) + + if rb.Len() != 5 { + t.Errorf("expected 5, got %d", rb.Len()) + } + }) + + t.Run("Empty Lines", func(t *testing.T) { + rb := newRingBuffer(100) + + lines := rb.Lines(5) + + if lines != nil { + t.Errorf("expected nil for empty buffer, got: %v", lines) + } + }) +} + +func TestExecTool_Bg_RingBufferOverflow(t *testing.T) { + tool, _ := NewExecTool("", false) + + defer tool.Shutdown() + + var cmd string + + if runtime.GOOS == "windows" { + cmd = "1..2000 | ForEach-Object { Write-Output ('x' * 50) }; Start-Sleep -Seconds 30" + } else { + cmd = "yes 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx' | head -n 2000; sleep 30" + } + + result := tool.Execute(context.Background(), map[string]any{ + "command": cmd, + + "background": true, + }) + + if result.IsError { + t.Fatalf("failed to start bg process: %s", result.ForLLM) + } + + time.Sleep(5 * time.Second) + + outputResult := tool.Execute(context.Background(), map[string]any{ + "bg_action": "output", + + "bg_id": "bg-1", + }) + + if outputResult.IsError { + t.Fatalf("failed to get output: %s", outputResult.ForLLM) + } + + procs := tool.BgProcesses() + + bp := procs["bg-1"] + + if bp == nil { + t.Fatal("bg-1 not found") + } + + bufLen := bp.output.Len() + + if bufLen > bgRingBufSize { + t.Errorf("ring buffer exceeded max size: %d > %d", bufLen, bgRingBufSize) + } +} + +func TestIsLocalHost(t *testing.T) { + tests := []struct { + host string + + want bool + }{ + + {"localhost", true}, + + {"LOCALHOST", true}, + + {"127.0.0.1", true}, + + {"127.0.0.2", true}, + + {"::1", true}, + + {"10.0.0.1", true}, + + {"10.255.255.255", true}, + + {"172.16.0.1", true}, + + {"172.31.255.255", true}, + + {"192.168.0.1", true}, + + {"192.168.1.100", true}, + + {"8.8.8.8", false}, + + {"1.1.1.1", false}, + + {"example.com", false}, + + {"api.github.com", false}, + + {"172.15.255.255", false}, + + {"172.32.0.0", false}, + } + + for _, tt := range tests { + got := isLocalHost(tt.host) + + if got != tt.want { + t.Errorf("isLocalHost(%q) = %v, want %v", tt.host, got, tt.want) + } + } +} + +func TestCheckCurlLocalNet(t *testing.T) { + tests := []struct { + cmd string + + wantErr bool + }{ + + {"curl http://localhost:3000/health", false}, + + {"curl -v http://127.0.0.1:8080/api/status", false}, + + {"wget http://192.168.1.10/file.bin", false}, + + {"curl -X POST http://10.0.0.5:9000/webhook", false}, + + {"curl http://example.com", true}, + + {"wget https://releases.github.com/v1.tar.gz", true}, + + {"curl http://8.8.8.8/data", true}, + + {"curl --help", false}, + + {"curl --version", false}, + + {"wget --help", false}, + } + + for _, tt := range tests { + errMsg := checkCurlLocalNet(tt.cmd) + + gotErr := errMsg != "" + + if gotErr != tt.wantErr { + t.Errorf("checkCurlLocalNet(%q): gotErr=%v wantErr=%v (msg: %q)", + + tt.cmd, gotErr, tt.wantErr, errMsg) + } + } +} + +func TestExecTool_LocalNetOnly(t *testing.T) { + tool, _ := NewExecTool("", false) + + tool.SetLocalNetOnly(true) + + tests := []struct { + cmd string + + wantErr bool + }{ + {"curl http://localhost:3000", false}, + + {"curl http://example.com", true}, + + {"echo hello", false}, + } + + ctx := context.Background() + + for _, tt := range tests { + result := tool.Execute(ctx, map[string]any{"command": tt.cmd}) + + if tt.wantErr && !result.IsError { + t.Errorf("cmd %q: expected blocked, but succeeded", tt.cmd) + } + + if !tt.wantErr && result.IsError && strings.Contains(result.ForLLM, "safety guard") { + t.Errorf("cmd %q: expected allowed, but safety guard blocked: %s", tt.cmd, result.ForLLM) + } + } +} diff --git a/pkg/tools/shell_test.go b/pkg/tools/shell_test.go index c4f4530be..90265e5bd 100644 --- a/pkg/tools/shell_test.go +++ b/pkg/tools/shell_test.go @@ -4,15 +4,14 @@ import ( "context" "os" "path/filepath" - "regexp" - "runtime" "strings" "testing" "time" + + "github.com/sipeed/picoclaw/pkg/config" ) // TestShellTool_Success verifies successful command execution - func TestShellTool_Success(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -20,7 +19,6 @@ func TestShellTool_Success(t *testing.T) { } ctx := context.Background() - args := map[string]any{ "command": "echo 'hello world'", } @@ -28,26 +26,22 @@ func TestShellTool_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain command output - if !strings.Contains(result.ForUser, "hello world") { t.Errorf("Expected ForUser to contain 'hello world', got: %s", result.ForUser) } // ForLLM should contain full output - if !strings.Contains(result.ForLLM, "hello world") { t.Errorf("Expected ForLLM to contain 'hello world', got: %s", result.ForLLM) } } // TestShellTool_Failure verifies failed command execution - func TestShellTool_Failure(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -55,7 +49,6 @@ func TestShellTool_Failure(t *testing.T) { } ctx := context.Background() - args := map[string]any{ "command": "ls /nonexistent_directory_12345", } @@ -63,26 +56,22 @@ func TestShellTool_Failure(t *testing.T) { result := tool.Execute(ctx, args) // Failure should be marked as error - if !result.IsError { t.Errorf("Expected error for failed command, got IsError=false") } // ForUser should contain error information - if result.ForUser == "" { t.Errorf("Expected ForUser to contain error info, got empty string") } // ForLLM should contain exit code or error - if !strings.Contains(result.ForLLM, "Exit code") && result.ForUser == "" { t.Errorf("Expected ForLLM to contain exit code or error, got: %s", result.ForLLM) } } // TestShellTool_Timeout verifies command timeout handling - func TestShellTool_Timeout(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -92,7 +81,6 @@ func TestShellTool_Timeout(t *testing.T) { tool.SetTimeout(100 * time.Millisecond) ctx := context.Background() - args := map[string]any{ "command": "sleep 10", } @@ -100,27 +88,21 @@ func TestShellTool_Timeout(t *testing.T) { result := tool.Execute(ctx, args) // Timeout should be marked as error - if !result.IsError { t.Errorf("Expected error for timeout, got IsError=false") } // Should mention timeout - if !strings.Contains(result.ForLLM, "timed out") && !strings.Contains(result.ForUser, "timed out") { t.Errorf("Expected timeout message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser) } } // TestShellTool_WorkingDir verifies custom working directory - func TestShellTool_WorkingDir(t *testing.T) { // Create temp directory - tmpDir := t.TempDir() - testFile := filepath.Join(tmpDir, "test.txt") - os.WriteFile(testFile, []byte("test content"), 0o644) tool, err := NewExecTool("", false) @@ -129,10 +111,8 @@ func TestShellTool_WorkingDir(t *testing.T) { } ctx := context.Background() - args := map[string]any{ - "command": "cat test.txt", - + "command": "cat test.txt", "working_dir": tmpDir, } @@ -148,7 +128,6 @@ func TestShellTool_WorkingDir(t *testing.T) { } // TestShellTool_DangerousCommand verifies safety guard blocks dangerous commands - func TestShellTool_DangerousCommand(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -156,7 +135,6 @@ func TestShellTool_DangerousCommand(t *testing.T) { } ctx := context.Background() - args := map[string]any{ "command": "rm -rf /", } @@ -164,7 +142,6 @@ func TestShellTool_DangerousCommand(t *testing.T) { result := tool.Execute(ctx, args) // Dangerous command should be blocked - if !result.IsError { t.Errorf("Expected dangerous command to be blocked (IsError=true)") } @@ -174,8 +151,27 @@ func TestShellTool_DangerousCommand(t *testing.T) { } } -// TestShellTool_MissingCommand verifies error handling for missing command +func TestShellTool_DangerousCommand_KillBlocked(t *testing.T) { + tool, err := NewExecTool("", false) + if err != nil { + t.Errorf("unable to configure exec tool: %s", err) + } + ctx := context.Background() + args := map[string]any{ + "command": "kill 12345", + } + + result := tool.Execute(ctx, args) + if !result.IsError { + t.Errorf("Expected kill command to be blocked") + } + if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") { + t.Errorf("Expected blocked message, got ForLLM: %s, ForUser: %s", result.ForLLM, result.ForUser) + } +} + +// TestShellTool_MissingCommand verifies error handling for missing command func TestShellTool_MissingCommand(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -183,20 +179,17 @@ func TestShellTool_MissingCommand(t *testing.T) { } ctx := context.Background() - args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when command is missing") } } // TestShellTool_StderrCapture verifies stderr is captured and included - func TestShellTool_StderrCapture(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -204,7 +197,6 @@ func TestShellTool_StderrCapture(t *testing.T) { } ctx := context.Background() - args := map[string]any{ "command": "sh -c 'echo stdout; echo stderr >&2'", } @@ -212,18 +204,15 @@ func TestShellTool_StderrCapture(t *testing.T) { result := tool.Execute(ctx, args) // Both stdout and stderr should be in output - if !strings.Contains(result.ForLLM, "stdout") { t.Errorf("Expected stdout in output, got: %s", result.ForLLM) } - if !strings.Contains(result.ForLLM, "stderr") { t.Errorf("Expected stderr in output, got: %s", result.ForLLM) } } // TestShellTool_OutputTruncation verifies long output is truncated - func TestShellTool_OutputTruncation(t *testing.T) { tool, err := NewExecTool("", false) if err != nil { @@ -231,9 +220,7 @@ func TestShellTool_OutputTruncation(t *testing.T) { } ctx := context.Background() - // Generate long output (>10000 chars) - args := map[string]any{ "command": "python3 -c \"print('x' * 20000)\" || echo " + strings.Repeat("x", 20000), } @@ -241,25 +228,19 @@ func TestShellTool_OutputTruncation(t *testing.T) { result := tool.Execute(ctx, args) // Should have truncation message or be truncated - if len(result.ForLLM) > 15000 { t.Errorf("Expected output to be truncated, got length: %d", len(result.ForLLM)) } } // TestShellTool_WorkingDir_OutsideWorkspace verifies that working_dir cannot escape the workspace directly - func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) { root := t.TempDir() - workspace := filepath.Join(root, "workspace") - outsideDir := filepath.Join(root, "outside") - if err := os.MkdirAll(workspace, 0o755); err != nil { t.Fatalf("failed to create workspace: %v", err) } - if err := os.MkdirAll(outsideDir, 0o755); err != nil { t.Fatalf("failed to create outside dir: %v", err) } @@ -270,45 +251,34 @@ func TestShellTool_WorkingDir_OutsideWorkspace(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": "pwd", - + "command": "pwd", "working_dir": outsideDir, }) if !result.IsError { t.Fatalf("expected working_dir outside workspace to be blocked, got output: %s", result.ForLLM) } - if !strings.Contains(result.ForLLM, "blocked") { t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM) } } // TestShellTool_WorkingDir_SymlinkEscape verifies that a symlink inside the workspace - // pointing outside cannot be used as working_dir to escape the sandbox. - func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) { root := t.TempDir() - workspace := filepath.Join(root, "workspace") - secretDir := filepath.Join(root, "secret") - if err := os.MkdirAll(workspace, 0o755); err != nil { t.Fatalf("failed to create workspace: %v", err) } - if err := os.MkdirAll(secretDir, 0o755); err != nil { t.Fatalf("failed to create secret dir: %v", err) } - os.WriteFile(filepath.Join(secretDir, "secret.txt"), []byte("top secret"), 0o644) // symlink lives inside the workspace but resolves to secretDir outside it - link := filepath.Join(workspace, "escape") - if err := os.Symlink(secretDir, link); err != nil { t.Skipf("symlinks not supported in this environment: %v", err) } @@ -319,25 +289,100 @@ func TestShellTool_WorkingDir_SymlinkEscape(t *testing.T) { } result := tool.Execute(context.Background(), map[string]any{ - "command": "cat secret.txt", - + "command": "cat secret.txt", "working_dir": link, }) if !result.IsError { t.Fatalf("expected symlink working_dir escape to be blocked, got output: %s", result.ForLLM) } - if !strings.Contains(result.ForLLM, "blocked") { t.Errorf("expected 'blocked' in error, got: %s", result.ForLLM) } } -// TestShellTool_RestrictToWorkspace verifies workspace restriction +// TestShellTool_RemoteChannelBlockedByDefault verifies exec is blocked for remote channels +func TestShellTool_RemoteChannelBlockedByDefault(t *testing.T) { + cfg := &config.Config{} + cfg.Tools.Exec.EnableDenyPatterns = true + cfg.Tools.Exec.AllowRemote = false + tool, err := NewExecToolWithConfig("", false, cfg) + if err != nil { + t.Fatalf("NewExecToolWithConfig() error: %v", err) + } + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + result := tool.Execute(ctx, map[string]any{"command": "echo hi"}) + + if !result.IsError { + t.Fatal("expected remote-channel exec to be blocked") + } + if !strings.Contains(result.ForLLM, "restricted to internal channels") { + t.Errorf("expected 'restricted to internal channels' message, got: %s", result.ForLLM) + } +} + +// TestShellTool_InternalChannelAllowed verifies exec is allowed for internal channels +func TestShellTool_InternalChannelAllowed(t *testing.T) { + cfg := &config.Config{} + cfg.Tools.Exec.EnableDenyPatterns = true + cfg.Tools.Exec.AllowRemote = false + + tool, err := NewExecToolWithConfig("", false, cfg) + if err != nil { + t.Fatalf("NewExecToolWithConfig() error: %v", err) + } + ctx := WithToolContext(context.Background(), "cli", "direct") + result := tool.Execute(ctx, map[string]any{"command": "echo hi"}) + + if result.IsError { + t.Fatalf("expected internal channel exec to succeed, got: %s", result.ForLLM) + } + if !strings.Contains(result.ForLLM, "hi") { + t.Errorf("expected output to contain 'hi', got: %s", result.ForLLM) + } +} + +// TestShellTool_EmptyChannelBlockedWhenNotAllowRemote verifies fail-closed when no channel context +func TestShellTool_EmptyChannelBlockedWhenNotAllowRemote(t *testing.T) { + cfg := &config.Config{} + cfg.Tools.Exec.EnableDenyPatterns = true + cfg.Tools.Exec.AllowRemote = false + + tool, err := NewExecToolWithConfig("", false, cfg) + if err != nil { + t.Fatalf("NewExecToolWithConfig() error: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "command": "echo hi", + }) + + if !result.IsError { + t.Fatal("expected exec with empty channel to be blocked when allowRemote=false") + } +} + +// TestShellTool_AllowRemoteBypassesChannelCheck verifies allowRemote=true permits any channel +func TestShellTool_AllowRemoteBypassesChannelCheck(t *testing.T) { + cfg := &config.Config{} + cfg.Tools.Exec.EnableDenyPatterns = true + cfg.Tools.Exec.AllowRemote = true + + tool, err := NewExecToolWithConfig("", false, cfg) + if err != nil { + t.Fatalf("NewExecToolWithConfig() error: %v", err) + } + ctx := WithToolContext(context.Background(), "telegram", "chat-1") + result := tool.Execute(ctx, map[string]any{"command": "echo hi"}) + + if result.IsError { + t.Fatalf("expected allowRemote=true to permit remote channel, got: %s", result.ForLLM) + } +} + +// TestShellTool_RestrictToWorkspace verifies workspace restriction func TestShellTool_RestrictToWorkspace(t *testing.T) { tmpDir := t.TempDir() - tool, err := NewExecTool(tmpDir, false) if err != nil { t.Errorf("unable to configure exec tool: %s", err) @@ -346,7 +391,6 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) { tool.SetRestrictToWorkspace(true) ctx := context.Background() - args := map[string]any{ "command": "cat ../../etc/passwd", } @@ -354,1035 +398,127 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) { result := tool.Execute(ctx, args) // Path traversal should be blocked - if !result.IsError { t.Errorf("Expected path traversal to be blocked with restrictToWorkspace=true") } if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") { t.Errorf( - "Expected 'blocked' message for path traversal, got ForLLM: %s, ForUser: %s", - result.ForLLM, - result.ForUser, ) } } -// --- guardCommand unit tests --- - -// TestGuardCommand_RelativePathWithSlashes verifies that relative paths - -// containing slashes (e.g., tests/cold/test.py, projects/terra-py-form) - -// are NOT falsely blocked. This was a regression caused by the old regex - -// matching "/cold/test.py" from "tests/cold/test.py" as an absolute path. - -func TestGuardCommand_RelativePathWithSlashes(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - cmds := []string{ - "pytest tests/cold/test_solver.py -v --tb=short", - - "cd projects/terra-py-form && pytest", - - "uv run pytest tests/cold/test_solver.py -v --tb=short", - - "cat src/terra_py_form/cold/parser.py", - - "python src/main.py --config config/dev.json", +// TestShellTool_DevNullAllowed verifies that /dev/null redirections are not blocked (issue #964). +func TestShellTool_DevNullAllowed(t *testing.T) { + tmpDir := t.TempDir() + tool, err := NewExecTool(tmpDir, true) + if err != nil { + t.Fatalf("unable to configure exec tool: %s", err) } - for _, cmd := range cmds { - result := tool.guardCommand(cmd, workspace) + commands := []string{ + "echo hello 2>/dev/null", + "echo hello >/dev/null", + "echo hello > /dev/null", + "echo hello 2> /dev/null", + "echo hello >/dev/null 2>&1", + "find " + tmpDir + " -name '*.go' 2>/dev/null", + } - if result != "" { - t.Errorf("Relative path should not be blocked: %q → %s", cmd, result) + for _, cmd := range commands { + result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + if result.IsError && strings.Contains(result.ForLLM, "blocked") { + t.Errorf("command should not be blocked: %s\n error: %s", cmd, result.ForLLM) } } } -// TestGuardCommand_VenvBinary verifies that .venv/bin/... paths are allowed - -// (they are relative paths, not absolute). - -func TestGuardCommand_VenvBinary(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - cmds := []string{ - ".venv/bin/python -m pytest", - - ".venv/bin/pytest tests/ -v", - - ".venv/bin/pip install -e .", +// TestShellTool_BlockDevices verifies that writes to block devices are blocked (issue #965). +func TestShellTool_BlockDevices(t *testing.T) { + tool, err := NewExecTool("", false) + if err != nil { + t.Fatalf("unable to configure exec tool: %s", err) } - for _, cmd := range cmds { - result := tool.guardCommand(cmd, workspace) + blocked := []string{ + "echo x > /dev/sda", + "echo x > /dev/hda", + "echo x > /dev/vda", + "echo x > /dev/xvda", + "echo x > /dev/nvme0n1", + "echo x > /dev/mmcblk0", + "echo x > /dev/loop0", + "echo x > /dev/dm-0", + "echo x > /dev/md0", + "echo x > /dev/sr0", + "echo x > /dev/nbd0", + } - if result != "" { - t.Errorf("Venv relative path should not be blocked: %q → %s", cmd, result) + for _, cmd := range blocked { + result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + if !result.IsError { + t.Errorf("expected block device write to be blocked: %s", cmd) } } } -// TestGuardCommand_ExecutableBinaryAllowed verifies that absolute paths - -// to executable files outside the workspace are allowed (system binaries). - -func TestGuardCommand_ExecutableBinaryAllowed(t *testing.T) { - if runtime.GOOS == "windows" { - t.Skip("Unix executable permission test not applicable on Windows") +// TestShellTool_SafePathsInWorkspaceRestriction verifies that safe kernel pseudo-devices +// are allowed even when workspace restriction is active. +func TestShellTool_SafePathsInWorkspaceRestriction(t *testing.T) { + tmpDir := t.TempDir() + tool, err := NewExecTool(tmpDir, true) + if err != nil { + t.Fatalf("unable to configure exec tool: %s", err) } - workspace := t.TempDir() - - externalDir := t.TempDir() - - // Create a fake executable outside the workspace - - execPath := filepath.Join(externalDir, "mybin") - - os.WriteFile(execPath, []byte("#!/bin/sh\necho ok"), 0o755) - - tool, _ := NewExecTool(workspace, true) - - cmd := execPath + " --help" - - result := tool.guardCommand(cmd, workspace) - - if result != "" { - t.Errorf("Executable binary outside workspace should be allowed: %q → %s", cmd, result) - } -} - -// TestGuardCommand_ExecutableBinaryAllowed_Windows verifies that .exe files - -// outside the workspace are allowed on Windows. - -func TestGuardCommand_ExecutableBinaryAllowed_Windows(t *testing.T) { - if runtime.GOOS != "windows" { - t.Skip("Windows-specific test") + // These reference paths outside workspace but should be allowed via safePaths. + commands := []string{ + "cat /dev/urandom | head -c 16 | od", + "echo test > /dev/null", + "dd if=/dev/zero bs=1 count=1", } - workspace := t.TempDir() - - externalDir := t.TempDir() - - // Create a fake .exe outside the workspace - - execPath := filepath.Join(externalDir, "tool.exe") - - os.WriteFile(execPath, []byte("MZ"), 0o644) - - tool, _ := NewExecTool(workspace, true) - - cmd := execPath + " --version" - - result := tool.guardCommand(cmd, workspace) - - if result != "" { - t.Errorf("Windows .exe outside workspace should be allowed: %q → %s", cmd, result) - } -} - -// TestGuardCommand_NonExecutableOutsideBlocked verifies that non-executable - -// files outside the workspace are blocked (e.g., reading /etc/shadow). - -func TestGuardCommand_NonExecutableOutsideBlocked(t *testing.T) { - if runtime.GOOS == "windows" { - t.Skip("Unix permission test not applicable on Windows") - } - - workspace := t.TempDir() - - externalDir := t.TempDir() - - // Create a regular (non-executable) file outside workspace - - dataFile := filepath.Join(externalDir, "secret.txt") - - os.WriteFile(dataFile, []byte("secret data"), 0o644) - - tool, _ := NewExecTool(workspace, true) - - cmd := "cat " + dataFile - - result := tool.guardCommand(cmd, workspace) - - if result == "" { - t.Errorf("Non-executable file outside workspace should be blocked: %q", cmd) - } - - if !strings.Contains(result, "path outside working dir") { - t.Errorf("Expected 'path outside working dir' message, got: %s", result) - } -} - -// TestGuardCommand_NonExistentAbsolutePathBlocked verifies that absolute - -// paths that don't exist are blocked (could be file creation outside workspace). - -func TestGuardCommand_NonExistentAbsolutePathBlocked(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - // Use platform-appropriate absolute path - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "echo hello > C:\\nonexistent_picoclaw_test_output" - } else { - cmd = "echo hello > /tmp/nonexistent_picoclaw_test_output" - } - - result := tool.guardCommand(cmd, workspace) - - if result == "" { - t.Errorf("Non-existent absolute path outside workspace should be blocked: %q", cmd) - } -} - -// TestGuardCommand_FlagEmbeddedPathSkipped verifies that paths embedded in - -// flags (e.g., -I/usr/local/include) are NOT extracted as absolute paths - -// because the token starts with "-", not "/". - -func TestGuardCommand_FlagEmbeddedPathSkipped(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - cmds := []string{ - "gcc -I/usr/local/include -L/usr/lib main.c", - - "g++ -std=c++17 -I/opt/include file.cpp", - - "python --prefix=/usr/local script.py", - } - - for _, cmd := range cmds { - result := tool.guardCommand(cmd, workspace) - - if result != "" { - t.Errorf("Flag-embedded path should not be blocked: %q → %s", cmd, result) + for _, cmd := range commands { + result := tool.Execute(context.Background(), map[string]any{"command": cmd}) + if result.IsError && strings.Contains(result.ForLLM, "path outside working dir") { + t.Errorf("safe path should not be blocked by workspace check: %s\n error: %s", cmd, result.ForLLM) } } } -// TestGuardCommand_AbsolutePathInsideWorkspace verifies that absolute paths - -// within the workspace are always allowed. - -func TestGuardCommand_AbsolutePathInsideWorkspace(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - innerDir := filepath.Join(workspace, "projects", "myapp") - - os.MkdirAll(innerDir, 0o755) - - cmd := "ls " + innerDir - - result := tool.guardCommand(cmd, workspace) - - if result != "" { - t.Errorf("Absolute path inside workspace should be allowed: %q → %s", cmd, result) - } -} - -// TestGuardCommand_PathTraversal verifies that various path traversal - -// patterns are blocked. - -func TestGuardCommand_PathTraversal(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - cmds := []string{ - "cat ../../etc/passwd", - - "cat ../../../etc/shadow", - - "ls projects/../../../../etc", +// TestShellTool_CustomAllowPatterns verifies that custom allow patterns exempt +// commands from deny pattern checks. +func TestShellTool_CustomAllowPatterns(t *testing.T) { + cfg := &config.Config{ + Tools: config.ToolsConfig{ + Exec: config.ExecConfig{ + EnableDenyPatterns: true, + CustomAllowPatterns: []string{`\bgit\s+push\s+origin\b`}, + }, + }, } - for _, cmd := range cmds { - result := tool.guardCommand(cmd, workspace) - - if result == "" { - t.Errorf("Path traversal should be blocked: %q", cmd) - } - - if !strings.Contains(result, "path traversal") { - t.Errorf("Expected 'path traversal' message, got: %s", result) - } - } -} - -// TestGuardCommand_CdWithAbsoluteWorkspacePath verifies that cd to an - -// absolute path within the workspace followed by other commands is allowed. - -func TestGuardCommand_CdWithAbsoluteWorkspacePath(t *testing.T) { - workspace := t.TempDir() - - innerDir := filepath.Join(workspace, "projects", "foo") - - os.MkdirAll(innerDir, 0o755) - - tool, _ := NewExecTool(workspace, true) - - cmd := "cd " + innerDir + " && ls -la" - - result := tool.guardCommand(cmd, workspace) - - if result != "" { - t.Errorf("cd to workspace subdir should be allowed: %q → %s", cmd, result) - } -} - -func TestGuardCommand_AgentCLISlashCommand(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - // Agent CLI slash commands (e.g., "/review") are not file paths. - - // They should be allowed because they don't exist on disk. - - cmds := []string{ - `codex exec --yolo "/review skip-git-repo-check"`, - - `claude "/review"`, - - `gemini "/help"`, - } - - for _, cmd := range cmds { - result := tool.guardCommand(cmd, workspace) - - if result != "" { - t.Errorf("Agent CLI slash command should not be blocked: %q → %s", cmd, result) - } - } - - // Non-agent commands with absolute paths should still be blocked. - - if runtime.GOOS != "windows" { - blocked := `cat /etc/hosts` - - result := tool.guardCommand(blocked, workspace) - - if result == "" { - t.Errorf("Non-agent command with absolute path should be blocked: %q", blocked) - } - } -} - -// TestGuardCommand_DenyPattern_IncludesPattern verifies that deny-match - -// error messages include the matched pattern string. - -func TestGuardCommand_DenyPattern_IncludesPattern(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - // Also add a custom deny pattern for precise matching. - - tool.denyPatterns = append(tool.denyPatterns, regexp.MustCompile(`\bdangerous_cmd\b`)) - - result := tool.guardCommand("dangerous_cmd --force", workspace) - - if result == "" { - t.Fatal("expected deny pattern to block the command") - } - - if !strings.Contains(result, "deny pattern") { - t.Errorf("expected 'deny pattern' in message, got: %s", result) - } - - if !strings.Contains(result, `\bdangerous_cmd\b`) { - t.Errorf("expected pattern string in message, got: %s", result) - } -} - -// TestGuardCommand_Allowlist_ShowsRules verifies that allowlist violation - -// messages include all configured rules. - -func TestGuardCommand_Allowlist_ShowsRules(t *testing.T) { - workspace := t.TempDir() - - tool, _ := NewExecTool(workspace, true) - - tool.SetAllowRules([]string{"go test", "git"}) - - result := tool.guardCommand("curl http://example.com", workspace) - - if result == "" { - t.Fatal("expected allowlist to block the command") - } - - if !strings.Contains(result, "not in allowlist") { - t.Errorf("expected 'not in allowlist' in message, got: %s", result) - } - - if !strings.Contains(result, "go test") || !strings.Contains(result, "git") { - t.Errorf("expected allowlist rules in message, got: %s", result) - } -} - -// TestGuardCommand_PathOutside_IncludesPath verifies that workspace-escape - -// messages include the offending path token. - -func TestGuardCommand_PathOutside_IncludesPath(t *testing.T) { - if runtime.GOOS == "windows" { - t.Skip("Unix absolute path test not applicable on Windows") - } - - workspace := t.TempDir() - - externalDir := t.TempDir() - - dataFile := filepath.Join(externalDir, "secret.txt") - - os.WriteFile(dataFile, []byte("secret"), 0o644) - - tool, _ := NewExecTool(workspace, true) - - result := tool.guardCommand("cat "+dataFile, workspace) - - if result == "" { - t.Fatal("expected path outside workspace to be blocked") - } - - if !strings.Contains(result, "path outside working dir") { - t.Errorf("expected 'path outside working dir' in message, got: %s", result) - } - - if !strings.Contains(result, dataFile) { - t.Errorf("expected offending path %q in message, got: %s", dataFile, result) - } -} - -// --- Background process tests --- - -func TestExecTool_Bg_StartAndOutput(t *testing.T) { - tool, _ := NewExecTool("", false) - - defer tool.Shutdown() - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "Write-Output 'hello from bg'; Start-Sleep -Seconds 30" - } else { - cmd = "echo 'hello from bg'; sleep 30" + tool, err := NewExecToolWithConfig("", false, cfg) + if err != nil { + t.Fatalf("unable to configure exec tool: %s", err) } + // "git push origin main" should be allowed by custom allow pattern. result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, + "command": "git push origin main", }) - - if result.IsError { - t.Fatalf("failed to start bg process: %s", result.ForLLM) + if result.IsError && strings.Contains(result.ForLLM, "blocked") { + t.Errorf("custom allow pattern should exempt 'git push origin main', got: %s", result.ForLLM) } - if !strings.Contains(result.ForLLM, "bg-1") { - t.Errorf("expected bg-1 in result, got: %s", result.ForLLM) - } - - if !strings.Contains(result.ForLLM, "Background process started") { - t.Errorf("expected start message, got: %s", result.ForLLM) - } - - // Get output - - outputResult := tool.Execute(context.Background(), map[string]any{ - "bg_action": "output", - - "bg_id": "bg-1", - }) - - if outputResult.IsError { - t.Fatalf("failed to get output: %s", outputResult.ForLLM) - } - - if !strings.Contains(outputResult.ForLLM, "hello from bg") { - t.Errorf("expected 'hello from bg' in output, got: %s", outputResult.ForLLM) - } - - if !strings.Contains(outputResult.ForLLM, "running") { - t.Errorf("expected 'running' status, got: %s", outputResult.ForLLM) - } -} - -func TestExecTool_Bg_Kill(t *testing.T) { - tool, _ := NewExecTool("", false) - - defer tool.Shutdown() - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "Start-Sleep -Seconds 60" - } else { - cmd = "sleep 60" - } - - result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, - }) - - if result.IsError { - t.Fatalf("failed to start bg process: %s", result.ForLLM) - } - - // Kill it - - killResult := tool.Execute(context.Background(), map[string]any{ - "bg_action": "kill", - - "bg_id": "bg-1", - }) - - if killResult.IsError { - t.Fatalf("failed to kill: %s", killResult.ForLLM) - } - - if !strings.Contains(killResult.ForLLM, "terminated") { - t.Errorf("expected 'terminated' message, got: %s", killResult.ForLLM) - } - - // Process should no longer be in the map - - procs := tool.BgProcesses() - - if _, ok := procs["bg-1"]; ok { - t.Errorf("expected bg-1 to be removed after kill") - } -} - -func TestExecTool_Bg_ExitedProcess(t *testing.T) { - tool, _ := NewExecTool("", false) - - defer tool.Shutdown() - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "Write-Output 'quick exit'" - } else { - cmd = "echo 'quick exit'" - } - - result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, - }) - - if result.IsError { - t.Fatalf("failed to start bg process: %s", result.ForLLM) - } - - // Wait for process to exit (initial capture is 3s, so after that it should be done) - - time.Sleep(4 * time.Second) - - // Get output — should show exited - - outputResult := tool.Execute(context.Background(), map[string]any{ - "bg_action": "output", - - "bg_id": "bg-1", - }) - - if outputResult.IsError { - t.Fatalf("failed to get output: %s", outputResult.ForLLM) - } - - if !strings.Contains(outputResult.ForLLM, "exited") { - t.Errorf("expected 'exited' in output, got: %s", outputResult.ForLLM) - } - - if !strings.Contains(outputResult.ForLLM, "quick exit") { - t.Errorf("expected 'quick exit' in output, got: %s", outputResult.ForLLM) - } -} - -func TestExecTool_Bg_InvalidID(t *testing.T) { - tool, _ := NewExecTool("", false) - - defer tool.Shutdown() - - // Output for non-existent ID - - result := tool.Execute(context.Background(), map[string]any{ - "bg_action": "output", - - "bg_id": "bg-999", - }) - - if !result.IsError { - t.Fatalf("expected error for invalid bg_id") - } - - if !strings.Contains(result.ForLLM, "not found") { - t.Errorf("expected 'not found' message, got: %s", result.ForLLM) - } - - // Kill for non-existent ID - + // "git push upstream main" should still be blocked (does not match allow pattern). result = tool.Execute(context.Background(), map[string]any{ - "bg_action": "kill", - - "bg_id": "bg-999", + "command": "git push upstream main", }) - if !result.IsError { - t.Fatalf("expected error for invalid bg_id") - } -} - -func TestExecTool_Bg_InitialOutputCapture(t *testing.T) { - tool, _ := NewExecTool("", false) - - defer tool.Shutdown() - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "Write-Output 'initial line 1'; Write-Output 'initial line 2'; Start-Sleep -Seconds 30" - } else { - cmd = "echo 'initial line 1'; echo 'initial line 2'; sleep 30" - } - - result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, - }) - - if result.IsError { - t.Fatalf("failed to start bg process: %s", result.ForLLM) - } - - if !strings.Contains(result.ForLLM, "initial line 1") { - t.Errorf("expected 'initial line 1' in initial output, got: %s", result.ForLLM) - } - - if !strings.Contains(result.ForLLM, "initial line 2") { - t.Errorf("expected 'initial line 2' in initial output, got: %s", result.ForLLM) - } -} - -func TestExecTool_Bg_RuntimeStatus(t *testing.T) { - tool, _ := NewExecTool("", false) - - defer tool.Shutdown() - - // No bg processes — should return empty - - if s := tool.RuntimeStatus(); s != "" { - t.Errorf("expected empty runtime status with no bg processes, got: %s", s) - } - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "Start-Sleep -Seconds 30" - } else { - cmd = "sleep 30" - } - - tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, - }) - - status := tool.RuntimeStatus() - - if !strings.Contains(status, "Background Processes") { - t.Errorf("expected 'Background Processes' section, got: %s", status) - } - - if !strings.Contains(status, "bg-1") { - t.Errorf("expected 'bg-1' in status, got: %s", status) - } - - if !strings.Contains(status, "running") { - t.Errorf("expected 'running' in status, got: %s", status) - } -} - -func TestExecTool_Bg_Shutdown(t *testing.T) { - tool, _ := NewExecTool("", false) - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "Start-Sleep -Seconds 60" - } else { - cmd = "sleep 60" - } - - tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, - }) - - tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, - }) - - // Both should be running - - procs := tool.BgProcesses() - - for _, bp := range procs { - if !bp.isRunning() { - t.Errorf("expected process to be running before shutdown") - } - } - - // Shutdown - - tool.Shutdown() - - // All should be done - - procs = tool.BgProcesses() - - for _, bp := range procs { - if bp.isRunning() { - t.Errorf("expected process to be stopped after shutdown") - } - } -} - -func TestRingBuffer(t *testing.T) { - t.Run("Write and String", func(t *testing.T) { - rb := newRingBuffer(100) - - rb.Write([]byte("hello ")) - - rb.Write([]byte("world")) - - if got := rb.String(); got != "hello world" { - t.Errorf("expected 'hello world', got %q", got) - } - }) - - t.Run("Lines", func(t *testing.T) { - rb := newRingBuffer(100) - - rb.Write([]byte("line1\nline2\nline3\nline4\nline5\n")) - - lines := rb.Lines(3) - - if len(lines) != 3 { - t.Fatalf("expected 3 lines, got %d", len(lines)) - } - - if lines[0] != "line3" || lines[1] != "line4" || lines[2] != "line5" { - t.Errorf("unexpected lines: %v", lines) - } - }) - - t.Run("Match", func(t *testing.T) { - rb := newRingBuffer(100) - - rb.Write([]byte("starting...\nServer ready on port 3000\nwaiting...\n")) - - re := regexp.MustCompile(`ready.*port`) - - match := rb.Match(re) - - if match == "" { - t.Fatal("expected match but got empty string") - } - - if !strings.Contains(match, "ready") { - t.Errorf("expected match to contain 'ready', got: %s", match) - } - - // Non-matching pattern - - re2 := regexp.MustCompile(`never_match`) - - match2 := rb.Match(re2) - - if match2 != "" { - t.Errorf("expected no match, got: %s", match2) - } - }) - - t.Run("Overflow", func(t *testing.T) { - rb := newRingBuffer(10) // small buffer - - rb.Write([]byte("1234567890ABCDEF")) - - got := rb.String() - - if len(got) != 10 { - t.Errorf("expected buffer to be 10 bytes, got %d", len(got)) - } - - // Should keep the last 10 bytes - - if got != "7890ABCDEF" { - t.Errorf("expected '7890ABCDEF', got %q", got) - } - }) - - t.Run("Len", func(t *testing.T) { - rb := newRingBuffer(100) - - if rb.Len() != 0 { - t.Errorf("expected 0 length initially") - } - - rb.Write([]byte("hello")) - - if rb.Len() != 5 { - t.Errorf("expected 5, got %d", rb.Len()) - } - }) - - t.Run("Empty Lines", func(t *testing.T) { - rb := newRingBuffer(100) - - lines := rb.Lines(5) - - if lines != nil { - t.Errorf("expected nil for empty buffer, got: %v", lines) - } - }) -} - -func TestExecTool_Bg_RingBufferOverflow(t *testing.T) { - tool, _ := NewExecTool("", false) - - defer tool.Shutdown() - - // Generate output larger than 32KB ring buffer - - var cmd string - - if runtime.GOOS == "windows" { - cmd = "1..2000 | ForEach-Object { Write-Output ('x' * 50) }; Start-Sleep -Seconds 30" - } else { - cmd = "yes 'xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx' | head -n 2000; sleep 30" - } - - result := tool.Execute(context.Background(), map[string]any{ - "command": cmd, - - "background": true, - }) - - if result.IsError { - t.Fatalf("failed to start bg process: %s", result.ForLLM) - } - - // Wait for output to accumulate - - time.Sleep(5 * time.Second) - - // Get output — ring buffer should have truncated old data - - outputResult := tool.Execute(context.Background(), map[string]any{ - "bg_action": "output", - - "bg_id": "bg-1", - }) - - if outputResult.IsError { - t.Fatalf("failed to get output: %s", outputResult.ForLLM) - } - - // The output should contain data but be bounded by the ring buffer size - - procs := tool.BgProcesses() - - bp := procs["bg-1"] - - if bp == nil { - t.Fatal("bg-1 not found") - } - - bufLen := bp.output.Len() - - if bufLen > bgRingBufSize { - t.Errorf("ring buffer exceeded max size: %d > %d", bufLen, bgRingBufSize) - } -} - -// TestIsLocalHost verifies localhost and RFC 1918 detection using net package. - -func TestIsLocalHost(t *testing.T) { - tests := []struct { - host string - - want bool - }{ - // Loopback / localhost - - {"localhost", true}, - - {"LOCALHOST", true}, - - {"127.0.0.1", true}, - - {"127.0.0.2", true}, - - {"::1", true}, - - // RFC 1918 private ranges - - {"10.0.0.1", true}, - - {"10.255.255.255", true}, - - {"172.16.0.1", true}, - - {"172.31.255.255", true}, - - {"192.168.0.1", true}, - - {"192.168.1.100", true}, - - // Public addresses - - {"8.8.8.8", false}, - - {"1.1.1.1", false}, - - {"example.com", false}, - - {"api.github.com", false}, - - // Edge: non-private but routable private-looking address - - {"172.15.255.255", false}, // just below 172.16/12 - - {"172.32.0.0", false}, // just above 172.31/12 - - } - - for _, tt := range tests { - got := isLocalHost(tt.host) - - if got != tt.want { - t.Errorf("isLocalHost(%q) = %v, want %v", tt.host, got, tt.want) - } - } -} - -// TestCheckCurlLocalNet verifies URL-level enforcement for curl/wget commands. - -func TestCheckCurlLocalNet(t *testing.T) { - tests := []struct { - cmd string - - wantErr bool - }{ - // Allowed: localhost and private IPs - - {"curl http://localhost:3000/health", false}, - - {"curl -v http://127.0.0.1:8080/api/status", false}, - - {"wget http://192.168.1.10/file.bin", false}, - - {"curl -X POST http://10.0.0.5:9000/webhook", false}, - - // Blocked: public addresses - - {"curl http://example.com", true}, - - {"wget https://releases.github.com/v1.tar.gz", true}, - - {"curl http://8.8.8.8/data", true}, - - // Allowed: no http URL (e.g. --help, --version — no network access) - - {"curl --help", false}, - - {"curl --version", false}, - - {"wget --help", false}, - } - - for _, tt := range tests { - errMsg := checkCurlLocalNet(tt.cmd) - - gotErr := errMsg != "" - - if gotErr != tt.wantErr { - t.Errorf("checkCurlLocalNet(%q): gotErr=%v wantErr=%v (msg: %q)", - - tt.cmd, gotErr, tt.wantErr, errMsg) - } - } -} - -// TestExecTool_LocalNetOnly verifies curl/wget blocking via SetLocalNetOnly. - -func TestExecTool_LocalNetOnly(t *testing.T) { - tool, _ := NewExecTool("", false) - - tool.SetLocalNetOnly(true) - - tests := []struct { - cmd string - - wantErr bool - }{ - {"curl http://localhost:3000", false}, - - {"curl http://example.com", true}, - - {"echo hello", false}, // non-curl not affected - - } - - ctx := context.Background() - - for _, tt := range tests { - result := tool.Execute(ctx, map[string]any{"command": tt.cmd}) - - if tt.wantErr && !result.IsError { - t.Errorf("cmd %q: expected blocked, but succeeded", tt.cmd) - } - - if !tt.wantErr && result.IsError && strings.Contains(result.ForLLM, "safety guard") { - t.Errorf("cmd %q: expected allowed, but safety guard blocked: %s", tt.cmd, result.ForLLM) - } + t.Errorf("'git push upstream main' should still be blocked by deny pattern") } } diff --git a/pkg/tools/shell_timeout_unix_ext_test.go b/pkg/tools/shell_timeout_unix_ext_test.go new file mode 100644 index 000000000..f8bcfca1a --- /dev/null +++ b/pkg/tools/shell_timeout_unix_ext_test.go @@ -0,0 +1,35 @@ +package tools + +import ( + "os" + "strconv" + "strings" + "syscall" +) + +func processRunning(pid int) bool { + if pid <= 0 { + return false + } + + err := syscall.Kill(pid, 0) + if err != nil && err != syscall.EPERM { + return false + } + + data, readErr := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat") + if readErr != nil { + return false + } + raw := string(data) + end := strings.LastIndex(raw, ")") + if end == -1 || end+2 >= len(raw) { + return true + } + fields := strings.Fields(raw[end+2:]) + if len(fields) == 0 { + return true + } + state := fields[0] + return state != "Z" +} diff --git a/pkg/tools/shell_timeout_unix_test.go b/pkg/tools/shell_timeout_unix_test.go index 748a055e8..357e1276e 100644 --- a/pkg/tools/shell_timeout_unix_test.go +++ b/pkg/tools/shell_timeout_unix_test.go @@ -13,33 +13,12 @@ import ( "time" ) -func processRunning(pid int) bool { +func processExists(pid int) bool { if pid <= 0 { return false } - - // kill(0) can return success for zombie processes too, so inspect /proc - // state and treat zombies as not-running for timeout cleanup assertions. err := syscall.Kill(pid, 0) - if err != nil && err != syscall.EPERM { - return false - } - - data, readErr := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat") - if readErr != nil { - return false - } - raw := string(data) - end := strings.LastIndex(raw, ")") - if end == -1 || end+2 >= len(raw) { - return true // best effort fallback - } - fields := strings.Fields(raw[end+2:]) - if len(fields) == 0 { - return true // best effort fallback - } - state := fields[0] - return state != "Z" + return err == nil || err == syscall.EPERM } func TestShellTool_TimeoutKillsChildProcess(t *testing.T) { @@ -47,12 +26,14 @@ func TestShellTool_TimeoutKillsChildProcess(t *testing.T) { if err != nil { t.Errorf("unable to configure exec tool: %s", err) } + tool.SetTimeout(500 * time.Millisecond) args := map[string]any{ // Spawn a child process that would outlive the shell unless process-group kill is used. "command": "sleep 60 & echo $! > child.pid; wait", } + result := tool.Execute(context.Background(), args) if !result.IsError { t.Fatalf("expected timeout error, got success: %s", result.ForLLM) @@ -66,6 +47,7 @@ func TestShellTool_TimeoutKillsChildProcess(t *testing.T) { if err != nil { t.Fatalf("failed to read child pid file: %v", err) } + childPID, err := strconv.Atoi(strings.TrimSpace(string(data))) if err != nil { t.Fatalf("failed to parse child pid: %v", err) @@ -73,10 +55,11 @@ func TestShellTool_TimeoutKillsChildProcess(t *testing.T) { deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { - if !processRunning(childPID) { + if !processExists(childPID) { return } time.Sleep(50 * time.Millisecond) } + t.Fatalf("child process %d is still running after timeout", childPID) } diff --git a/pkg/tools/skills_install_test.go b/pkg/tools/skills_install_test.go index 882e446c6..676fcecc0 100644 --- a/pkg/tools/skills_install_test.go +++ b/pkg/tools/skills_install_test.go @@ -14,29 +14,22 @@ import ( func TestInstallSkillToolName(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) - assert.Equal(t, "install_skill", tool.Name()) } func TestInstallSkillToolMissingSlug(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) - result := tool.Execute(context.Background(), map[string]any{}) - assert.True(t, result.IsError) - assert.Contains(t, result.ForLLM, "identifier is required and must be a non-empty string") } func TestInstallSkillToolEmptySlug(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) - result := tool.Execute(context.Background(), map[string]any{ "slug": " ", }) - assert.True(t, result.IsError) - assert.Contains(t, result.ForLLM, "identifier is required and must be a non-empty string") } @@ -45,9 +38,7 @@ func TestInstallSkillToolUnsafeSlug(t *testing.T) { cases := []string{ "../etc/passwd", - "path/traversal", - "path\\traversal", } @@ -55,85 +46,59 @@ func TestInstallSkillToolUnsafeSlug(t *testing.T) { result := tool.Execute(context.Background(), map[string]any{ "slug": slug, }) - assert.True(t, result.IsError, "slug %q should be rejected", slug) - assert.Contains(t, result.ForLLM, "invalid slug") } } func TestInstallSkillToolAlreadyExists(t *testing.T) { workspace := t.TempDir() - skillDir := filepath.Join(workspace, "skills", "existing-skill") - require.NoError(t, os.MkdirAll(skillDir, 0o755)) tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace) - result := tool.Execute(context.Background(), map[string]any{ - "slug": "existing-skill", - + "slug": "existing-skill", "registry": "clawhub", }) - assert.True(t, result.IsError) - assert.Contains(t, result.ForLLM, "already installed") } func TestInstallSkillToolRegistryNotFound(t *testing.T) { workspace := t.TempDir() - tool := NewInstallSkillTool(skills.NewRegistryManager(), workspace) - result := tool.Execute(context.Background(), map[string]any{ - "slug": "some-skill", - + "slug": "some-skill", "registry": "nonexistent", }) - assert.True(t, result.IsError) - assert.Contains(t, result.ForLLM, "registry") - assert.Contains(t, result.ForLLM, "not found") } func TestInstallSkillToolParameters(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) - params := tool.Parameters() props, ok := params["properties"].(map[string]any) - assert.True(t, ok) - assert.Contains(t, props, "slug") - assert.Contains(t, props, "version") - assert.Contains(t, props, "registry") - assert.Contains(t, props, "force") required, ok := params["required"].([]string) - assert.True(t, ok) - assert.Contains(t, required, "slug") - assert.Contains(t, required, "registry") } func TestInstallSkillToolMissingRegistry(t *testing.T) { tool := NewInstallSkillTool(skills.NewRegistryManager(), t.TempDir()) - result := tool.Execute(context.Background(), map[string]any{ "slug": "some-skill", }) - assert.True(t, result.IsError) - assert.Contains(t, result.ForLLM, "invalid registry") } diff --git a/pkg/tools/skills_search_test.go b/pkg/tools/skills_search_test.go index cb6a0e104..0e5387cf5 100644 --- a/pkg/tools/skills_search_test.go +++ b/pkg/tools/skills_search_test.go @@ -11,110 +11,80 @@ import ( func TestFindSkillsToolName(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) - assert.Equal(t, "find_skills", tool.Name()) } func TestFindSkillsToolMissingQuery(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) - result := tool.Execute(context.Background(), map[string]any{}) - assert.True(t, result.IsError) - assert.Contains(t, result.ForLLM, "query is required") } func TestFindSkillsToolEmptyQuery(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) - result := tool.Execute(context.Background(), map[string]any{ "query": " ", }) - assert.True(t, result.IsError) } func TestFindSkillsToolCacheHit(t *testing.T) { cache := skills.NewSearchCache(10, 5*60*1000*1000*1000) // 5 min - cache.Put("github", []skills.SearchResult{ {Slug: "github", Score: 0.9, RegistryName: "clawhub"}, }) tool := NewFindSkillsTool(skills.NewRegistryManager(), cache) - result := tool.Execute(context.Background(), map[string]any{ "query": "github", }) assert.False(t, result.IsError) - assert.Contains(t, result.ForLLM, "github") - assert.Contains(t, result.ForLLM, "cached") } func TestFindSkillsToolParameters(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) - params := tool.Parameters() props, ok := params["properties"].(map[string]any) - assert.True(t, ok) - assert.Contains(t, props, "query") - assert.Contains(t, props, "limit") required, ok := params["required"].([]string) - assert.True(t, ok) - assert.Contains(t, required, "query") } func TestFindSkillsToolDescription(t *testing.T) { tool := NewFindSkillsTool(skills.NewRegistryManager(), nil) - assert.NotEmpty(t, tool.Description()) - assert.Contains(t, tool.Description(), "skill") } func TestFormatSearchResultsEmpty(t *testing.T) { result := formatSearchResults("test query", nil, false) - assert.Contains(t, result, "No skills found") } func TestFormatSearchResultsWithData(t *testing.T) { results := []skills.SearchResult{ { - Slug: "github", - - Score: 0.95, - - DisplayName: "GitHub", - - Summary: "GitHub API integration", - - Version: "1.0.0", - + Slug: "github", + Score: 0.95, + DisplayName: "GitHub", + Summary: "GitHub API integration", + Version: "1.0.0", RegistryName: "clawhub", }, } - output := formatSearchResults("github", results, false) - assert.Contains(t, output, "github") - assert.Contains(t, output, "v1.0.0") - assert.Contains(t, output, "0.950") - assert.Contains(t, output, "clawhub") - assert.Contains(t, output, "install_skill") } diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index 35eb2c0a9..43223b8db 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -8,43 +8,33 @@ import ( func TestSpawnTool_Execute_EmptyTask(t *testing.T) { provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSpawnTool(manager) ctx := context.Background() tests := []struct { name string - args map[string]any }{ {"empty string", map[string]any{"task": ""}}, - {"whitespace only", map[string]any{"task": " "}}, - {"tabs and newlines", map[string]any{"task": "\t\n "}}, - {"missing task key", map[string]any{"label": "test"}}, - {"wrong type", map[string]any{"task": 123}}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := tool.Execute(ctx, tt.args) - if result == nil { t.Fatal("Result should not be nil") } - if !result.IsError { t.Error("Expected error for invalid task parameter") } - - if !strings.Contains(result.ForLLM, `"task"`) { - t.Errorf("Error message should mention '\"task\"', got: %s", result.ForLLM) + if !strings.Contains(result.ForLLM, "task is required") { + t.Errorf("Error message should mention 'task is required', got: %s", result.ForLLM) } }) } @@ -52,29 +42,22 @@ func TestSpawnTool_Execute_EmptyTask(t *testing.T) { func TestSpawnTool_Execute_ValidTask(t *testing.T) { provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, nil, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSpawnTool(manager) ctx := context.Background() - args := map[string]any{ - "task": "Write a haiku about coding", - + "task": "Write a haiku about coding", "label": "haiku-task", } result := tool.Execute(ctx, args) - if result == nil { t.Fatal("Result should not be nil") } - if result.IsError { t.Errorf("Expected success for valid task, got error: %s", result.ForLLM) } - if !result.Async { t.Error("SpawnTool should return async result") } @@ -84,16 +67,13 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) { tool := NewSpawnTool(nil) ctx := context.Background() - args := map[string]any{"task": "test task"} result := tool.Execute(ctx, args) - if !result.IsError { t.Error("Expected error for nil manager") } - - if !strings.Contains(result.ForLLM, "spawn tool is not available") { - t.Errorf("Error message should mention spawn tool not available, got: %s", result.ForLLM) + if !strings.Contains(result.ForLLM, "Subagent manager not configured") { + t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM) } } diff --git a/pkg/tools/subagent_tool_ext_test.go b/pkg/tools/subagent_tool_ext_test.go new file mode 100644 index 000000000..60cd52ee8 --- /dev/null +++ b/pkg/tools/subagent_tool_ext_test.go @@ -0,0 +1,102 @@ +package tools + +import ( + "context" + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/orch" + "testing" + "time" +) + +func TestSubagentTool_SetContext(t *testing.T) { + provider := &MockLLMProvider{} + + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) + + tool := NewSubagentTool(manager) + + tool.SetContext("test-channel", "test-chat") + +} + +func TestFormatToolStats(t *testing.T) { + tests := []struct { + name string + + stats map[string]int + + want string + }{ + {"empty", map[string]int{}, ""}, + + {"single", map[string]int{"exec": 3}, "exec:3"}, + + { + "multiple sorted", + + map[string]int{"read_file": 5, "exec": 3, "write_file": 1}, + + "exec:3,read_file:5,write_file:1", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := formatToolStats(tt.stats) + + if got != tt.want { + t.Errorf("formatToolStats(%v) = %q, want %q", tt.stats, got, tt.want) + } + }) + } +} + +func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) { + provider := &MockLLMProvider{} + + msgBus := bus.NewMessageBus() + + mgr := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) + + _, err := mgr.Spawn( + + context.Background(), + + "say hello", "meta-test", "", "cli", "direct", "", + + nil, + ) + if err != nil { + t.Fatalf("Spawn() error: %v", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + + defer cancel() + + received, ok := msgBus.ConsumeInbound(ctx) + + if !ok { + t.Fatal("timed out waiting for bus message") + } + + if received.Channel != "system" { + t.Fatalf("expected channel 'system', got %q", received.Channel) + } + + if received.Metadata == nil { + t.Fatal("Metadata should not be nil") + } + + if received.Metadata["iterations"] != "1" { + t.Errorf("iterations = %q, want %q", received.Metadata["iterations"], "1") + } + + if received.Metadata["tool_calls"] != "0" { + t.Errorf("tool_calls = %q, want %q", received.Metadata["tool_calls"], "0") + } + + if received.Metadata["duration_ms"] == "" { + t.Error("duration_ms should be present") + } +} diff --git a/pkg/tools/subagent_tool_test.go b/pkg/tools/subagent_tool_test.go index f9e1f988a..4b6f130a5 100644 --- a/pkg/tools/subagent_tool_test.go +++ b/pkg/tools/subagent_tool_test.go @@ -4,34 +4,24 @@ import ( "context" "strings" "testing" - "time" - "github.com/sipeed/picoclaw/pkg/bus" - "github.com/sipeed/picoclaw/pkg/orch" "github.com/sipeed/picoclaw/pkg/providers" ) // MockLLMProvider is a test implementation of LLMProvider - type MockLLMProvider struct { lastOptions map[string]any } func (m *MockLLMProvider) Chat( ctx context.Context, - messages []providers.Message, - tools []providers.ToolDefinition, - model string, - options map[string]any, ) (*providers.LLMResponse, error) { m.lastOptions = options - // Find the last user message to generate a response - for i := len(messages) - 1; i >= 0; i-- { if messages[i].Role == "user" { return &providers.LLMResponse{ @@ -39,7 +29,6 @@ func (m *MockLLMProvider) Chat( }, nil } } - return &providers.LLMResponse{Content: "No task provided"}, nil } @@ -57,19 +46,12 @@ func (m *MockLLMProvider) GetContextWindow() int { func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") manager.SetLLMOptions(2048, 0.6) - tool := NewSubagentTool(manager) - tool.SetContext("cli", "direct") - - ctx := context.Background() - + ctx := WithToolContext(context.Background(), "cli", "direct") args := map[string]any{"task": "Do something"} - result := tool.Execute(ctx, args) if result == nil || result.IsError { @@ -79,23 +61,18 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { if provider.lastOptions == nil { t.Fatal("Expected LLM options to be passed, got nil") } - if provider.lastOptions["max_tokens"] != 2048 { t.Fatalf("max_tokens = %v, want %d", provider.lastOptions["max_tokens"], 2048) } - if provider.lastOptions["temperature"] != 0.6 { t.Fatalf("temperature = %v, want %v", provider.lastOptions["temperature"], 0.6) } } // TestSubagentTool_Name verifies tool name - func TestSubagentTool_Name(t *testing.T) { provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) if tool.Name() != "subagent" { @@ -104,198 +81,131 @@ func TestSubagentTool_Name(t *testing.T) { } // TestSubagentTool_Description verifies tool description - func TestSubagentTool_Description(t *testing.T) { provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) desc := tool.Description() - if desc == "" { t.Error("Description should not be empty") } - - if !strings.Contains(desc, "BLOCK") { - t.Errorf("Description should mention 'BLOCK', got: %s", desc) - } - - if !strings.Contains(desc, "spawn") { - t.Errorf("Description should contrast with spawn, got: %s", desc) + if !strings.Contains(desc, "subagent") { + t.Errorf("Description should mention 'subagent', got: %s", desc) } } // TestSubagentTool_Parameters verifies tool parameters schema - func TestSubagentTool_Parameters(t *testing.T) { provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) params := tool.Parameters() - if params == nil { t.Error("Parameters should not be nil") } // Check type - if params["type"] != "object" { t.Errorf("Expected type 'object', got: %v", params["type"]) } // Check properties - props, ok := params["properties"].(map[string]any) - if !ok { t.Fatal("Properties should be a map") } // Verify task parameter - task, ok := props["task"].(map[string]any) - if !ok { t.Fatal("Task parameter should exist") } - if task["type"] != "string" { t.Errorf("Task type should be 'string', got: %v", task["type"]) } // Verify label parameter - label, ok := props["label"].(map[string]any) - if !ok { t.Fatal("Label parameter should exist") } - if label["type"] != "string" { t.Errorf("Label type should be 'string', got: %v", label["type"]) } // Check required fields - required, ok := params["required"].([]string) - if !ok { t.Fatal("Required should be a string array") } - if len(required) != 1 || required[0] != "task" { t.Errorf("Required should be ['task'], got: %v", required) } } -// TestSubagentTool_SetContext verifies context setting - -func TestSubagentTool_SetContext(t *testing.T) { - provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) - - tool := NewSubagentTool(manager) - - tool.SetContext("test-channel", "test-chat") - - // Verify context is set (we can't directly access private fields, - - // but we can verify it doesn't crash) - - // The actual context usage is tested in Execute tests -} - // TestSubagentTool_Execute_Success tests successful execution - func TestSubagentTool_Execute_Success(t *testing.T) { provider := &MockLLMProvider{} - - msgBus := bus.NewMessageBus() - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) - tool.SetContext("telegram", "chat-123") - - ctx := context.Background() - + ctx := WithToolContext(context.Background(), "telegram", "chat-123") args := map[string]any{ - "task": "Write a haiku about coding", - + "task": "Write a haiku about coding", "label": "haiku-task", } result := tool.Execute(ctx, args) // Verify basic ToolResult structure - if result == nil { t.Fatal("Result should not be nil") } // Verify no error - if result.IsError { t.Errorf("Expected success, got error: %s", result.ForLLM) } // Verify not async - if result.Async { t.Error("SubagentTool should be synchronous, not async") } // Verify not silent - if result.Silent { t.Error("SubagentTool should not be silent") } // Verify ForUser contains brief summary (not empty) - if result.ForUser == "" { t.Error("ForUser should contain result summary") } - if !strings.Contains(result.ForUser, "Task completed") { t.Errorf("ForUser should contain task completion, got: %s", result.ForUser) } // Verify ForLLM contains full details - if result.ForLLM == "" { t.Error("ForLLM should contain full details") } - if !strings.Contains(result.ForLLM, "haiku-task") { t.Errorf("ForLLM should contain label 'haiku-task', got: %s", result.ForLLM) } - if !strings.Contains(result.ForLLM, "Task completed:") { t.Errorf("ForLLM should contain task result, got: %s", result.ForLLM) } } // TestSubagentTool_Execute_NoLabel tests execution without label - func TestSubagentTool_Execute_NoLabel(t *testing.T) { provider := &MockLLMProvider{} - - msgBus := bus.NewMessageBus() - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) ctx := context.Background() - args := map[string]any{ "task": "Test task without label", } @@ -307,23 +217,18 @@ func TestSubagentTool_Execute_NoLabel(t *testing.T) { } // ForLLM should show (unnamed) for missing label - if !strings.Contains(result.ForLLM, "(unnamed)") { t.Errorf("ForLLM should show '(unnamed)' for missing label, got: %s", result.ForLLM) } } // TestSubagentTool_Execute_MissingTask tests error handling for missing task - func TestSubagentTool_Execute_MissingTask(t *testing.T) { provider := &MockLLMProvider{} - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) ctx := context.Background() - args := map[string]any{ "label": "test", } @@ -331,35 +236,26 @@ func TestSubagentTool_Execute_MissingTask(t *testing.T) { result := tool.Execute(ctx, args) // Should return error - if !result.IsError { t.Error("Expected error for missing task parameter") } - // ForLLM should contain helpful error with example - - if !strings.Contains(result.ForLLM, `"task"`) { - t.Errorf("Error message should mention '\"task\"', got: %s", result.ForLLM) - } - - if !strings.Contains(result.ForLLM, "Example") { - t.Errorf("Error message should include usage example, got: %s", result.ForLLM) + // ForLLM should contain error message + if !strings.Contains(result.ForLLM, "task is required") { + t.Errorf("Error message should mention 'task is required', got: %s", result.ForLLM) } // Err should be set - if result.Err == nil { t.Error("Err should be set for validation failure") } } // TestSubagentTool_Execute_NilManager tests error handling for nil manager - func TestSubagentTool_Execute_NilManager(t *testing.T) { tool := NewSubagentTool(nil) ctx := context.Background() - args := map[string]any{ "task": "test task", } @@ -367,37 +263,24 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) { result := tool.Execute(ctx, args) // Should return error - if !result.IsError { t.Error("Expected error for nil manager") } - if !strings.Contains(result.ForLLM, "not available in this session") { - t.Errorf("Error message should mention 'not available in this session', got: %s", result.ForLLM) + if !strings.Contains(result.ForLLM, "Subagent manager not configured") { + t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM) } } // TestSubagentTool_Execute_ContextPassing verifies context is properly used - func TestSubagentTool_Execute_ContextPassing(t *testing.T) { provider := &MockLLMProvider{} - - msgBus := bus.NewMessageBus() - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) - // Set context - channel := "test-channel" - chatID := "test-chat" - - tool.SetContext(channel, chatID) - - ctx := context.Background() - + ctx := WithToolContext(context.Background(), channel, chatID) args := map[string]any{ "task": "Test context passing", } @@ -405,144 +288,40 @@ func TestSubagentTool_Execute_ContextPassing(t *testing.T) { result := tool.Execute(ctx, args) // Should succeed - if result.IsError { t.Errorf("Expected success with context, got error: %s", result.ForLLM) } // The context is used internally; we can't directly test it - // but execution success indicates context was handled properly } // TestSubagentTool_ForUserTruncation verifies long content is truncated for user - func TestSubagentTool_ForUserTruncation(t *testing.T) { // Create a mock provider that returns very long content - provider := &MockLLMProvider{} - - msgBus := bus.NewMessageBus() - - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) - + manager := NewSubagentManager(provider, "test-model", "/tmp/test") tool := NewSubagentTool(manager) ctx := context.Background() // Create a task that will generate long response - longTask := strings.Repeat("This is a very long task description. ", 100) - args := map[string]any{ - "task": longTask, - + "task": longTask, "label": "long-test", } result := tool.Execute(ctx, args) // ForUser should be truncated to 500 chars + "..." - maxUserLen := 500 - if len(result.ForUser) > maxUserLen+3 { // +3 for "..." t.Errorf("ForUser should be truncated to ~%d chars, got: %d", maxUserLen, len(result.ForUser)) } // ForLLM should have full content - if !strings.Contains(result.ForLLM, longTask[:50]) { t.Error("ForLLM should contain reference to original task") } } - -func TestFormatToolStats(t *testing.T) { - tests := []struct { - name string - - stats map[string]int - - want string - }{ - {"empty", map[string]int{}, ""}, - - {"single", map[string]int{"exec": 3}, "exec:3"}, - - { - "multiple sorted", - - map[string]int{"read_file": 5, "exec": 3, "write_file": 1}, - - "exec:3,read_file:5,write_file:1", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got := formatToolStats(tt.stats) - - if got != tt.want { - t.Errorf("formatToolStats(%v) = %q, want %q", tt.stats, got, tt.want) - } - }) - } -} - -// TestSubagentManager_Spawn_SetsMetadata verifies that the bus message from a - -// completed spawn includes execution statistics in Metadata. - -func TestSubagentManager_Spawn_SetsMetadata(t *testing.T) { - provider := &MockLLMProvider{} - - msgBus := bus.NewMessageBus() - - mgr := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus, orch.Noop, WebSearchToolOptions{}) - - _, err := mgr.Spawn( - - context.Background(), - - "say hello", "meta-test", "", "cli", "direct", "", - - nil, - ) - if err != nil { - t.Fatalf("Spawn() error: %v", err) - } - - // Consume the inbound message from the bus - - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) - - defer cancel() - - received, ok := msgBus.ConsumeInbound(ctx) - - if !ok { - t.Fatal("timed out waiting for bus message") - } - - if received.Channel != "system" { - t.Fatalf("expected channel 'system', got %q", received.Channel) - } - - if received.Metadata == nil { - t.Fatal("Metadata should not be nil") - } - - if received.Metadata["iterations"] != "1" { - t.Errorf("iterations = %q, want %q", received.Metadata["iterations"], "1") - } - - if received.Metadata["tool_calls"] != "0" { - t.Errorf("tool_calls = %q, want %q", received.Metadata["tool_calls"], "0") - } - - // duration_ms should be a non-negative number - - if received.Metadata["duration_ms"] == "" { - t.Error("duration_ms should be present") - } -} diff --git a/pkg/tools/web_test.go b/pkg/tools/web_test.go index 5758fed11..0737d2087 100644 --- a/pkg/tools/web_test.go +++ b/pkg/tools/web_test.go @@ -1,32 +1,39 @@ package tools import ( + "bytes" "context" "encoding/json" + "fmt" + "net" "net/http" "net/http/httptest" "strings" "testing" "time" + + "github.com/sipeed/picoclaw/pkg/logger" ) -// TestWebTool_WebFetch_Success verifies successful URL fetching +const testFetchLimit = int64(10 * 1024 * 1024) +// TestWebTool_WebFetch_Success verifies successful URL fetching func TestWebTool_WebFetch_Success(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html") - w.WriteHeader(http.StatusOK) - w.Write([]byte("

Test Page

Content here

")) })) - defer server.Close() - tool := NewWebFetchTool(50000) + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } ctx := context.Background() - args := map[string]any{ "url": server.URL, } @@ -34,45 +41,41 @@ func TestWebTool_WebFetch_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } - // ForUser should contain the fetched content - - if !strings.Contains(result.ForUser, "Test Page") { - t.Errorf("Expected ForUser to contain 'Test Page', got: %s", result.ForUser) + // ForLLM should contain the fetched content (full JSON result) + if !strings.Contains(result.ForLLM, "Test Page") { + t.Errorf("Expected ForLLM to contain 'Test Page', got: %s", result.ForLLM) } - // ForLLM should contain summary - - if !strings.Contains(result.ForLLM, "bytes") && !strings.Contains(result.ForLLM, "extractor") { - t.Errorf("Expected ForLLM to contain summary, got: %s", result.ForLLM) + // ForUser should contain summary + if !strings.Contains(result.ForUser, "bytes") && !strings.Contains(result.ForUser, "extractor") { + t.Errorf("Expected ForUser to contain summary, got: %s", result.ForUser) } } // TestWebTool_WebFetch_JSON verifies JSON content handling - func TestWebTool_WebFetch_JSON(t *testing.T) { - testData := map[string]string{"key": "value", "number": "123"} + withPrivateWebFetchHostsAllowed(t) + testData := map[string]string{"key": "value", "number": "123"} expectedJSON, _ := json.MarshalIndent(testData, "", " ") server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - w.Write(expectedJSON) })) - defer server.Close() - tool := NewWebFetchTool(50000) + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } ctx := context.Background() - args := map[string]any{ "url": server.URL, } @@ -80,25 +83,24 @@ func TestWebTool_WebFetch_JSON(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } - // ForUser should contain formatted JSON - - if !strings.Contains(result.ForUser, "key") && !strings.Contains(result.ForUser, "value") { - t.Errorf("Expected ForUser to contain JSON data, got: %s", result.ForUser) + // ForLLM should contain formatted JSON + if !strings.Contains(result.ForLLM, "key") && !strings.Contains(result.ForLLM, "value") { + t.Errorf("Expected ForLLM to contain JSON data, got: %s", result.ForLLM) } } // TestWebTool_WebFetch_InvalidURL verifies error handling for invalid URL - func TestWebTool_WebFetch_InvalidURL(t *testing.T) { - tool := NewWebFetchTool(50000) + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } ctx := context.Background() - args := map[string]any{ "url": "not-a-valid-url", } @@ -106,25 +108,24 @@ func TestWebTool_WebFetch_InvalidURL(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error for invalid URL") } // Should contain error message (either "invalid URL" or scheme error) - if !strings.Contains(result.ForLLM, "URL") && !strings.Contains(result.ForUser, "URL") { t.Errorf("Expected error message for invalid URL, got ForLLM: %s", result.ForLLM) } } // TestWebTool_WebFetch_UnsupportedScheme verifies error handling for non-http URLs - func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) { - tool := NewWebFetchTool(50000) + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } ctx := context.Background() - args := map[string]any{ "url": "ftp://example.com/file.txt", } @@ -132,61 +133,58 @@ func TestWebTool_WebFetch_UnsupportedScheme(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error for unsupported URL scheme") } // Should mention only http/https allowed - if !strings.Contains(result.ForLLM, "http/https") && !strings.Contains(result.ForUser, "http/https") { t.Errorf("Expected scheme error message, got ForLLM: %s", result.ForLLM) } } // TestWebTool_WebFetch_MissingURL verifies error handling for missing URL - func TestWebTool_WebFetch_MissingURL(t *testing.T) { - tool := NewWebFetchTool(50000) + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } ctx := context.Background() - args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when URL is missing") } // Should mention URL is required - if !strings.Contains(result.ForLLM, "url is required") && !strings.Contains(result.ForUser, "url is required") { t.Errorf("Expected 'url is required' message, got ForLLM: %s", result.ForLLM) } } // TestWebTool_WebFetch_Truncation verifies content truncation - func TestWebTool_WebFetch_Truncation(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + longContent := strings.Repeat("x", 20000) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/plain") - w.WriteHeader(http.StatusOK) - w.Write([]byte(longContent)) })) - defer server.Close() - tool := NewWebFetchTool(1000) // Limit to 1000 chars + tool, err := NewWebFetchTool(1000, testFetchLimit) // Limit to 1000 chars + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } ctx := context.Background() - args := map[string]any{ "url": server.URL, } @@ -194,17 +192,13 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } - // ForUser should contain truncated content (not the full 20000 chars) - + // ForLLM should contain truncated content (not the full 20000 chars) resultMap := make(map[string]any) - - json.Unmarshal([]byte(result.ForUser), &resultMap) - + json.Unmarshal([]byte(result.ForLLM), &resultMap) if text, ok := resultMap["text"].(string); ok { if len(text) > 1100 { // Allow some margin t.Errorf("Expected content to be truncated to ~1000 chars, got: %d", len(text)) @@ -212,79 +206,118 @@ func TestWebTool_WebFetch_Truncation(t *testing.T) { } // Should be marked as truncated - if truncated, ok := resultMap["truncated"].(bool); !ok || !truncated { t.Errorf("Expected 'truncated' to be true in result") } } -// TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing +func TestWebFetchTool_PayloadTooLarge(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + // Create a mock HTTP server + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/html") + w.WriteHeader(http.StatusOK) + + // Generate a payload intentionally larger than our limit. + // Limit: 10 * 1024 * 1024 (10MB). We generate 10MB + 100 bytes of the letter 'A'. + largeData := bytes.Repeat([]byte("A"), int(testFetchLimit)+100) + + w.Write(largeData) + })) + // Ensure the server is shut down at the end of the test + defer ts.Close() + + // Initialize the tool + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } + + // Prepare the arguments pointing to the URL of our local mock server + args := map[string]any{ + "url": ts.URL, + } + + // Execute the tool + ctx := context.Background() + result := tool.Execute(ctx, args) + + // Assuming ErrorResult sets the ForLLM field with the error text. + if result == nil { + t.Fatal("expected a ToolResult, got nil") + } + + // Search for the exact error string we set earlier in the Execute method + expectedErrorMsg := fmt.Sprintf("size exceeded %d bytes limit", testFetchLimit) + + if !strings.Contains(result.ForLLM, expectedErrorMsg) && !strings.Contains(result.ForUser, expectedErrorMsg) { + t.Errorf("test failed: expected error %q, but got: %+v", expectedErrorMsg, result) + } +} + +// TestWebTool_WebSearch_NoApiKey verifies that no tool is created when API key is missing func TestWebTool_WebSearch_NoApiKey(t *testing.T) { - tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: ""}) + tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKeys: nil}) if err != nil { t.Fatalf("Unexpected error: %v", err) } - if tool != nil { t.Errorf("Expected nil tool when Brave API key is empty") } // Also nil when nothing is enabled - tool, err = NewWebSearchTool(WebSearchToolOptions{}) if err != nil { t.Fatalf("Unexpected error: %v", err) } - if tool != nil { t.Errorf("Expected nil tool when no provider is enabled") } } // TestWebTool_WebSearch_MissingQuery verifies error handling for missing query - func TestWebTool_WebSearch_MissingQuery(t *testing.T) { - tool, err := NewWebSearchTool(WebSearchToolOptions{BraveEnabled: true, BraveAPIKey: "test-key", BraveMaxResults: 5}) + tool, err := NewWebSearchTool(WebSearchToolOptions{ + BraveEnabled: true, + BraveAPIKeys: []string{"test-key"}, + BraveMaxResults: 5, + }) if err != nil { t.Fatalf("Unexpected error: %v", err) } - ctx := context.Background() - args := map[string]any{} result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error when query is missing") } } // TestWebTool_WebFetch_HTMLExtraction verifies HTML text extraction - func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html") - w.WriteHeader(http.StatusOK) - w.Write( - []byte( `

Title

Content

`, ), ) })) - defer server.Close() - tool := NewWebFetchTool(50000) + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } ctx := context.Background() - args := map[string]any{ "url": server.URL, } @@ -292,105 +325,80 @@ func TestWebTool_WebFetch_HTMLExtraction(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } - // ForUser should contain extracted text (without script/style tags) - - if !strings.Contains(result.ForUser, "Title") && !strings.Contains(result.ForUser, "Content") { - t.Errorf("Expected ForUser to contain extracted text, got: %s", result.ForUser) + // ForLLM should contain extracted text (without script/style tags) + if !strings.Contains(result.ForLLM, "Title") && !strings.Contains(result.ForLLM, "Content") { + t.Errorf("Expected ForLLM to contain extracted text, got: %s", result.ForLLM) } - // Should NOT contain script or style tags - - if strings.Contains(result.ForUser, "

Keep this

", - wantFunc: func(t *testing.T, got string) { if strings.Contains(got, "alert") || strings.Contains(got, "body{}") { t.Errorf("Expected script/style content removed, got: %q", got) } - if !strings.Contains(got, "Keep this") { t.Errorf("Expected 'Keep this' to remain, got: %q", got) } }, }, - { - name: "collapses excessive blank lines", - + name: "collapses excessive blank lines", input: "

A

\n\n\n\n\n

B

", - wantFunc: func(t *testing.T, got string) { if strings.Contains(got, "\n\n\n") { t.Errorf("Expected excessive blank lines collapsed, got: %q", got) } }, }, - { - name: "collapses horizontal whitespace", - + name: "collapses horizontal whitespace", input: "

hello world

", - wantFunc: func(t *testing.T, got string) { if strings.Contains(got, " ") { t.Errorf("Expected spaces collapsed, got: %q", got) } - if !strings.Contains(got, "hello world") { t.Errorf("Expected 'hello world', got: %q", got) } }, }, - { - name: "empty input", - + name: "empty input", input: "", - wantFunc: func(t *testing.T, got string) { if got != "" { t.Errorf("Expected empty string, got: %q", got) @@ -402,19 +410,218 @@ func TestWebFetchTool_extractText(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := tool.extractText(tt.input) - tt.wantFunc(t, got) }) } } -// TestWebTool_WebFetch_MissingDomain verifies error handling for URL without domain +func withPrivateWebFetchHostsAllowed(t *testing.T) { + t.Helper() + previous := allowPrivateWebFetchHosts.Load() + allowPrivateWebFetchHosts.Store(true) + t.Cleanup(func() { + allowPrivateWebFetchHosts.Store(previous) + }) +} +func TestWebTool_WebFetch_PrivateHostBlocked(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://127.0.0.1:0", + }) + + if !result.IsError { + t.Errorf("expected error for private host URL, got success") + } + if !strings.Contains(result.ForLLM, "private or local network") && + !strings.Contains(result.ForUser, "private or local network") { + t.Errorf("expected private host block message, got %q", result.ForLLM) + } +} + +func TestWebTool_WebFetch_PrivateHostAllowedForTests(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/plain") + w.WriteHeader(http.StatusOK) + w.Write([]byte("ok")) + })) + defer server.Close() + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": server.URL, + }) + + if result.IsError { + t.Errorf("expected success when private host access is allowed in tests, got %q", result.ForLLM) + } +} + +// TestWebFetch_BlocksIPv4MappedIPv6Loopback verifies ::ffff:127.0.0.1 is blocked +func TestWebFetch_BlocksIPv4MappedIPv6Loopback(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[::ffff:127.0.0.1]:0", + }) + + if !result.IsError { + t.Error("expected error for IPv4-mapped IPv6 loopback URL, got success") + } +} + +// TestWebFetch_BlocksMetadataIP verifies 169.254.169.254 is blocked +func TestWebFetch_BlocksMetadataIP(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://169.254.169.254/latest/meta-data", + }) + + if !result.IsError { + t.Error("expected error for cloud metadata IP, got success") + } +} + +// TestWebFetch_BlocksIPv6UniqueLocal verifies fc00::/7 addresses are blocked +func TestWebFetch_BlocksIPv6UniqueLocal(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[fd00::1]:0", + }) + + if !result.IsError { + t.Error("expected error for IPv6 unique local address, got success") + } +} + +// TestWebFetch_Blocks6to4WithPrivateEmbed verifies 6to4 with private embedded IPv4 is blocked +func TestWebFetch_Blocks6to4WithPrivateEmbed(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + // 2002:7f00:0001::1 embeds 127.0.0.1 + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[2002:7f00:0001::1]:0", + }) + + if !result.IsError { + t.Error("expected error for 6to4 with private embedded IPv4, got success") + } +} + +// TestWebFetch_Allows6to4WithPublicEmbed verifies 6to4 with public embedded IPv4 is NOT blocked +func TestWebFetch_Allows6to4WithPublicEmbed(t *testing.T) { + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + // 2002:0801:0101::1 embeds 8.1.1.1 (public) — pre-flight should pass, + // connection will fail (no listener) but that's after the SSRF check. + result := tool.Execute(context.Background(), map[string]any{ + "url": "http://[2002:0801:0101::1]:0", + }) + + // Should NOT be blocked by SSRF check — error should be connection failure, not "private" + if result.IsError && strings.Contains(result.ForLLM, "private") { + t.Error("6to4 with public embedded IPv4 should not be blocked as private") + } +} + +// TestWebFetch_RedirectToPrivateBlocked verifies redirects to private IPs are blocked +func TestWebFetch_RedirectToPrivateBlocked(t *testing.T) { + withPrivateWebFetchHostsAllowed(t) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Redirect to a private IP + http.Redirect(w, r, "http://10.0.0.1/secret", http.StatusFound) + })) + defer server.Close() + + // Temporarily disable private host allowance for the redirect check + allowPrivateWebFetchHosts.Store(false) + defer allowPrivateWebFetchHosts.Store(true) + + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + t.Fatalf("Failed to create web fetch tool: %v", err) + } + result := tool.Execute(context.Background(), map[string]any{ + "url": server.URL, + }) + + if !result.IsError { + t.Error("expected error when redirecting to private IP, got success") + } +} + +// TestIsPrivateOrRestrictedIP_Table tests IP classification logic +func TestIsPrivateOrRestrictedIP_Table(t *testing.T) { + tests := []struct { + ip string + blocked bool + desc string + }{ + {"127.0.0.1", true, "IPv4 loopback"}, + {"10.0.0.1", true, "IPv4 private class A"}, + {"172.16.0.1", true, "IPv4 private class B"}, + {"192.168.1.1", true, "IPv4 private class C"}, + {"169.254.169.254", true, "link-local / cloud metadata"}, + {"100.64.0.1", true, "carrier-grade NAT"}, + {"0.0.0.0", true, "unspecified"}, + {"8.8.8.8", false, "public DNS"}, + {"1.1.1.1", false, "public DNS"}, + {"::1", true, "IPv6 loopback"}, + {"::ffff:127.0.0.1", true, "IPv4-mapped IPv6 loopback"}, + {"::ffff:10.0.0.1", true, "IPv4-mapped IPv6 private"}, + {"fc00::1", true, "IPv6 unique local"}, + {"fd00::1", true, "IPv6 unique local"}, + {"2002:7f00:0001::1", true, "6to4 with embedded 127.x (private)"}, + {"2002:0a00:0001::1", true, "6to4 with embedded 10.0.0.1 (private)"}, + {"2002:0801:0101::1", false, "6to4 with embedded 8.1.1.1 (public)"}, + {"2001:0000:4136:e378:8000:63bf:f5ff:fffe", true, "Teredo with client 10.0.0.1 (private)"}, + {"2001:0000:4136:e378:8000:63bf:f7f6:fefe", false, "Teredo with client 8.9.1.1 (public)"}, + {"2607:f8b0:4004:800::200e", false, "public IPv6 (Google)"}, + } + + for _, tt := range tests { + t.Run(tt.desc, func(t *testing.T) { + ip := net.ParseIP(tt.ip) + if ip == nil { + t.Fatalf("failed to parse IP: %s", tt.ip) + } + got := isPrivateOrRestrictedIP(ip) + if got != tt.blocked { + t.Errorf("isPrivateOrRestrictedIP(%s) = %v, want %v", tt.ip, got, tt.blocked) + } + }) + } +} + +// TestWebTool_WebFetch_MissingDomain verifies error handling for URL without domain func TestWebTool_WebFetch_MissingDomain(t *testing.T) { - tool := NewWebFetchTool(50000) + tool, err := NewWebFetchTool(50000, testFetchLimit) + if err != nil { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } ctx := context.Background() - args := map[string]any{ "url": "https://", } @@ -422,13 +629,11 @@ func TestWebTool_WebFetch_MissingDomain(t *testing.T) { result := tool.Execute(ctx, args) // Should return error result - if !result.IsError { t.Errorf("Expected error for URL without domain") } // Should mention missing domain - if !strings.Contains(result.ForLLM, "domain") && !strings.Contains(result.ForUser, "domain") { t.Errorf("Expected domain error message, got ForLLM: %s", result.ForLLM) } @@ -439,17 +644,14 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { if err != nil { t.Fatalf("createHTTPClient() error: %v", err) } - if client.Timeout != 12*time.Second { t.Fatalf("client.Timeout = %v, want %v", client.Timeout, 12*time.Second) } tr, ok := client.Transport.(*http.Transport) - if !ok { t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) } - if tr.Proxy == nil { t.Fatal("transport.Proxy is nil, want non-nil") } @@ -458,12 +660,10 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { if err != nil { t.Fatalf("http.NewRequest() error: %v", err) } - proxyURL, err := tr.Proxy(req) if err != nil { t.Fatalf("transport.Proxy(req) error: %v", err) } - if proxyURL == nil || proxyURL.String() != "http://127.0.0.1:7890" { t.Fatalf("proxy URL = %v, want %q", proxyURL, "http://127.0.0.1:7890") } @@ -471,7 +671,6 @@ func TestCreateHTTPClient_ProxyConfigured(t *testing.T) { func TestCreateHTTPClient_InvalidProxy(t *testing.T) { _, err := createHTTPClient("://bad-proxy", 10*time.Second) - if err == nil { t.Fatal("createHTTPClient() expected error for invalid proxy URL, got nil") } @@ -484,21 +683,17 @@ func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) { } tr, ok := client.Transport.(*http.Transport) - if !ok { t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) } - req, err := http.NewRequest("GET", "https://example.com", nil) if err != nil { t.Fatalf("http.NewRequest() error: %v", err) } - proxyURL, err := tr.Proxy(req) if err != nil { t.Fatalf("transport.Proxy(req) error: %v", err) } - if proxyURL == nil || proxyURL.String() != "socks5://127.0.0.1:1080" { t.Fatalf("proxy URL = %v, want %q", proxyURL, "socks5://127.0.0.1:1080") } @@ -506,11 +701,9 @@ func TestCreateHTTPClient_Socks5ProxyConfigured(t *testing.T) { func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) { _, err := createHTTPClient("ftp://127.0.0.1:21", 10*time.Second) - if err == nil { t.Fatal("createHTTPClient() expected error for unsupported scheme, got nil") } - if !strings.Contains(err.Error(), "unsupported proxy scheme") { t.Fatalf("error = %q, want to contain %q", err.Error(), "unsupported proxy scheme") } @@ -518,19 +711,12 @@ func TestCreateHTTPClient_UnsupportedProxyScheme(t *testing.T) { func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { t.Setenv("HTTP_PROXY", "http://127.0.0.1:8888") - t.Setenv("http_proxy", "http://127.0.0.1:8888") - t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8888") - t.Setenv("https_proxy", "http://127.0.0.1:8888") - t.Setenv("ALL_PROXY", "") - t.Setenv("all_proxy", "") - t.Setenv("NO_PROXY", "") - t.Setenv("no_proxy", "") client, err := createHTTPClient("", 10*time.Second) @@ -539,11 +725,9 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { } tr, ok := client.Transport.(*http.Transport) - if !ok { t.Fatalf("client.Transport type = %T, want *http.Transport", client.Transport) } - if tr.Proxy == nil { t.Fatal("transport.Proxy is nil, want proxy function from environment") } @@ -552,19 +736,16 @@ func TestCreateHTTPClient_ProxyFromEnvironmentWhenConfigEmpty(t *testing.T) { if err != nil { t.Fatalf("http.NewRequest() error: %v", err) } - if _, err := tr.Proxy(req); err != nil { t.Fatalf("transport.Proxy(req) error: %v", err) } } func TestNewWebFetchToolWithProxy(t *testing.T) { - tool, err := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890") + tool, err := NewWebFetchToolWithProxy(1024, "http://127.0.0.1:7890", testFetchLimit) if err != nil { - t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) - } - - if tool.maxChars != 1024 { + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) + } else if tool.maxChars != 1024 { t.Fatalf("maxChars = %d, want %d", tool.maxChars, 1024) } @@ -572,9 +753,9 @@ func TestNewWebFetchToolWithProxy(t *testing.T) { t.Fatalf("proxy = %q, want %q", tool.proxy, "http://127.0.0.1:7890") } - tool, err = NewWebFetchToolWithProxy(0, "http://127.0.0.1:7890") + tool, err = NewWebFetchToolWithProxy(0, "http://127.0.0.1:7890", testFetchLimit) if err != nil { - t.Fatalf("NewWebFetchToolWithProxy() error: %v", err) + logger.ErrorCF("agent", "Failed to create web fetch tool", map[string]any{"error": err.Error()}) } if tool.maxChars != 50000 { @@ -585,24 +766,18 @@ func TestNewWebFetchToolWithProxy(t *testing.T) { func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("perplexity", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ - PerplexityEnabled: true, - - PerplexityAPIKey: "k", - + PerplexityEnabled: true, + PerplexityAPIKeys: []string{"k"}, PerplexityMaxResults: 3, - - Proxy: "http://127.0.0.1:7890", + Proxy: "http://127.0.0.1:7890", }) if err != nil { t.Fatalf("NewWebSearchTool() error: %v", err) } - p, ok := tool.provider.(*PerplexitySearchProvider) - if !ok { t.Fatalf("provider type = %T, want *PerplexitySearchProvider", tool.provider) } - if p.proxy != "http://127.0.0.1:7890" { t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") } @@ -610,24 +785,18 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("brave", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ - BraveEnabled: true, - - BraveAPIKey: "k", - + BraveEnabled: true, + BraveAPIKeys: []string{"k"}, BraveMaxResults: 3, - - Proxy: "http://127.0.0.1:7890", + Proxy: "http://127.0.0.1:7890", }) if err != nil { t.Fatalf("NewWebSearchTool() error: %v", err) } - p, ok := tool.provider.(*BraveSearchProvider) - if !ok { t.Fatalf("provider type = %T, want *BraveSearchProvider", tool.provider) } - if p.proxy != "http://127.0.0.1:7890" { t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") } @@ -635,22 +804,17 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { t.Run("duckduckgo", func(t *testing.T) { tool, err := NewWebSearchTool(WebSearchToolOptions{ - DuckDuckGoEnabled: true, - + DuckDuckGoEnabled: true, DuckDuckGoMaxResults: 3, - - Proxy: "http://127.0.0.1:7890", + Proxy: "http://127.0.0.1:7890", }) if err != nil { t.Fatalf("NewWebSearchTool() error: %v", err) } - p, ok := tool.provider.(*DuckDuckGoSearchProvider) - if !ok { t.Fatalf("provider type = %T, want *DuckDuckGoSearchProvider", tool.provider) } - if p.proxy != "http://127.0.0.1:7890" { t.Fatalf("provider proxy = %q, want %q", p.proxy, "http://127.0.0.1:7890") } @@ -658,69 +822,50 @@ func TestNewWebSearchTool_PropagatesProxy(t *testing.T) { } // TestWebTool_TavilySearch_Success verifies successful Tavily search - func TestWebTool_TavilySearch_Success(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != "POST" { t.Errorf("Expected POST request, got %s", r.Method) } - if r.Header.Get("Content-Type") != "application/json" { t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) } // Verify payload - var payload map[string]any - json.NewDecoder(r.Body).Decode(&payload) - if payload["api_key"] != "test-key" { t.Errorf("Expected api_key test-key, got %v", payload["api_key"]) } - if payload["query"] != "test query" { t.Errorf("Expected query 'test query', got %v", payload["query"]) } // Return mock response - response := map[string]any{ "results": []map[string]any{ { - "title": "Test Result 1", - - "url": "https://example.com/1", - + "title": "Test Result 1", + "url": "https://example.com/1", "content": "Content for result 1", }, - { - "title": "Test Result 2", - - "url": "https://example.com/2", - + "title": "Test Result 2", + "url": "https://example.com/2", "content": "Content for result 2", }, }, } - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(http.StatusOK) - json.NewEncoder(w).Encode(response) })) - defer server.Close() tool, err := NewWebSearchTool(WebSearchToolOptions{ - TavilyEnabled: true, - - TavilyAPIKey: "test-key", - - TavilyBaseURL: server.URL, - + TavilyEnabled: true, + TavilyAPIKeys: []string{"test-key"}, + TavilyBaseURL: server.URL, TavilyMaxResults: 5, }) if err != nil { @@ -728,7 +873,6 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { } ctx := context.Background() - args := map[string]any{ "query": "test query", } @@ -736,22 +880,265 @@ func TestWebTool_TavilySearch_Success(t *testing.T) { result := tool.Execute(ctx, args) // Success should not be an error - if result.IsError { t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) } // ForUser should contain result titles and URLs - if !strings.Contains(result.ForUser, "Test Result 1") || - !strings.Contains(result.ForUser, "https://example.com/1") { t.Errorf("Expected results in output, got: %s", result.ForUser) } // Should mention via Tavily - if !strings.Contains(result.ForUser, "via Tavily") { t.Errorf("Expected 'via Tavily' in output, got: %s", result.ForUser) } } + +func TestAPIKeyPool(t *testing.T) { + pool := NewAPIKeyPool([]string{"key1", "key2", "key3"}) + if len(pool.keys) != 3 { + t.Fatalf("expected 3 keys, got %d", len(pool.keys)) + } + if pool.keys[0] != "key1" || pool.keys[1] != "key2" || pool.keys[2] != "key3" { + t.Fatalf("unexpected keys: %v", pool.keys) + } + + // Test Iterator: each iterator should cover all keys exactly once + iter := pool.NewIterator() + expected := []string{"key1", "key2", "key3"} + for i, want := range expected { + k, ok := iter.Next() + if !ok { + t.Fatalf("iter.Next() returned false at step %d", i) + } + if k != want { + t.Errorf("step %d: expected %s, got %s", i, want, k) + } + } + // Should be exhausted + if _, ok := iter.Next(); ok { + t.Errorf("expected iterator exhausted after all keys") + } + + // Second iterator starts at next position (load balancing) + iter2 := pool.NewIterator() + k, ok := iter2.Next() + if !ok { + t.Fatal("iter2.Next() returned false") + } + if k != "key2" { + t.Errorf("expected key2 (round-robin), got %s", k) + } + + // Empty pool + emptyPool := NewAPIKeyPool([]string{}) + emptyIter := emptyPool.NewIterator() + if _, ok := emptyIter.Next(); ok { + t.Errorf("expected false for empty pool") + } + + // Single key pool + singlePool := NewAPIKeyPool([]string{"single"}) + singleIter := singlePool.NewIterator() + if k, ok := singleIter.Next(); !ok || k != "single" { + t.Errorf("expected single, got %s (ok=%v)", k, ok) + } + if _, ok := singleIter.Next(); ok { + t.Errorf("expected exhausted after single key") + } +} + +func TestWebTool_TavilySearch_Failover(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode payload: %v", err) + } + + apiKey := payload["api_key"].(string) + + if apiKey == "key1" { + w.WriteHeader(http.StatusTooManyRequests) + w.Write([]byte("Rate limited")) + return + } + + if apiKey == "key2" { + // Success + response := map[string]any{ + "results": []map[string]any{ + { + "title": "Success Result", + "url": "https://example.com/success", + "content": "Success content", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(response) + return + } + + w.WriteHeader(http.StatusBadRequest) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + TavilyEnabled: true, + TavilyAPIKeys: []string{"key1", "key2"}, + TavilyBaseURL: server.URL, + TavilyMaxResults: 5, + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + ctx := context.Background() + args := map[string]any{ + "query": "test query", + } + + result := tool.Execute(ctx, args) + + if result.IsError { + t.Errorf("Expected success, got Error: %s", result.ForLLM) + } + if !strings.Contains(result.ForUser, "Success Result") { + t.Errorf("Expected failover to second key and success result, got: %s", result.ForUser) + } +} + +func TestWebTool_GLMSearch_Success(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + t.Errorf("Expected POST request, got %s", r.Method) + } + if r.Header.Get("Content-Type") != "application/json" { + t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type")) + } + if r.Header.Get("Authorization") != "Bearer test-glm-key" { + t.Errorf("Expected Authorization Bearer test-glm-key, got %s", r.Header.Get("Authorization")) + } + + var payload map[string]any + json.NewDecoder(r.Body).Decode(&payload) + if payload["search_query"] != "test query" { + t.Errorf("Expected search_query 'test query', got %v", payload["search_query"]) + } + if payload["search_engine"] != "search_std" { + t.Errorf("Expected search_engine 'search_std', got %v", payload["search_engine"]) + } + + response := map[string]any{ + "id": "web-search-test", + "created": 1709568000, + "search_result": []map[string]any{ + { + "title": "Test GLM Result", + "content": "GLM search snippet", + "link": "https://example.com/glm", + "media": "Example", + "publish_date": "2026-03-04", + }, + }, + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + json.NewEncoder(w).Encode(response) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + GLMSearchEnabled: true, + GLMSearchAPIKey: "test-glm-key", + GLMSearchBaseURL: server.URL, + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + }) + + if result.IsError { + t.Errorf("Expected success, got IsError=true: %s", result.ForLLM) + } + if !strings.Contains(result.ForUser, "Test GLM Result") { + t.Errorf("Expected 'Test GLM Result' in output, got: %s", result.ForUser) + } + if !strings.Contains(result.ForUser, "https://example.com/glm") { + t.Errorf("Expected URL in output, got: %s", result.ForUser) + } + if !strings.Contains(result.ForUser, "via GLM Search") { + t.Errorf("Expected 'via GLM Search' in output, got: %s", result.ForUser) + } +} + +func TestWebTool_GLMSearch_APIError(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + w.Write([]byte(`{"error":"invalid api key"}`)) + })) + defer server.Close() + + tool, err := NewWebSearchTool(WebSearchToolOptions{ + GLMSearchEnabled: true, + GLMSearchAPIKey: "bad-key", + GLMSearchBaseURL: server.URL, + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + result := tool.Execute(context.Background(), map[string]any{ + "query": "test query", + }) + + if !result.IsError { + t.Errorf("Expected IsError=true for 401 response") + } + if !strings.Contains(result.ForLLM, "status 401") { + t.Errorf("Expected status 401 in error, got: %s", result.ForLLM) + } +} + +func TestWebTool_GLMSearch_Priority(t *testing.T) { + // GLM Search should only be selected when all other providers are disabled + tool, err := NewWebSearchTool(WebSearchToolOptions{ + DuckDuckGoEnabled: true, + DuckDuckGoMaxResults: 5, + GLMSearchEnabled: true, + GLMSearchAPIKey: "test-key", + GLMSearchBaseURL: "https://example.com", + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + + // DuckDuckGo should win over GLM Search + if _, ok := tool.provider.(*DuckDuckGoSearchProvider); !ok { + t.Errorf("Expected DuckDuckGoSearchProvider when both enabled, got %T", tool.provider) + } + + // With DuckDuckGo disabled, GLM Search should be selected + tool2, err := NewWebSearchTool(WebSearchToolOptions{ + DuckDuckGoEnabled: false, + GLMSearchEnabled: true, + GLMSearchAPIKey: "test-key", + GLMSearchBaseURL: "https://example.com", + GLMSearchEngine: "search_std", + }) + if err != nil { + t.Fatalf("NewWebSearchTool() error: %v", err) + } + if _, ok := tool2.provider.(*GLMSearchProvider); !ok { + t.Errorf("Expected GLMSearchProvider when only GLM enabled, got %T", tool2.provider) + } +} diff --git a/pkg/utils/string_ext_test.go b/pkg/utils/string_ext_test.go new file mode 100644 index 000000000..a5b7fa9da --- /dev/null +++ b/pkg/utils/string_ext_test.go @@ -0,0 +1,179 @@ +package utils + +import ( + "strings" + "testing" +) + +func TestStripThinkBlocks_ClosedBlock(t *testing.T) { + in := "\nsecret reasoning\n\n\nVisible content" + got := StripThinkBlocks(in) + if got != "Visible content" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "Visible content") + } +} + +func TestStripThinkBlocks_UnclosedBlock(t *testing.T) { + in := "reasoning that never ends\nmore reasoning" + got := StripThinkBlocks(in) + if got != "" { + t.Fatalf("StripThinkBlocks() = %q, want empty", got) + } +} + +func TestStripThinkBlocks_MultipleBlocks(t *testing.T) { + in := "firstmiddlesecondend" + got := StripThinkBlocks(in) + if got != "middleend" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "middleend") + } +} + +func TestStripThinkBlocks_NoBlocks(t *testing.T) { + in := "plain text without think blocks" + got := StripThinkBlocks(in) + if got != in { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, in) + } +} + +func TestStripThinkBlocks_CaseInsensitive(t *testing.T) { + in := "upper casevisible" + got := StripThinkBlocks(in) + if got != "visible" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "visible") + } +} + +func TestStripThinkBlocks_ClosedThenUnclosed(t *testing.T) { + in := "closedmiddleunclosed tail" + got := StripThinkBlocks(in) + if got != "middle" { + t.Fatalf("StripThinkBlocks() = %q, want %q", got, "middle") + } +} + +func TestDetectRepetitionLoop_HighRepetition(t *testing.T) { + + phrase := "結構本格的なコード" + repeated := strings.Repeat(phrase, 300) + if !DetectRepetitionLoop(repeated) { + t.Fatal("DetectRepetitionLoop should return true for highly repetitive text") + } +} + +func TestDetectRepetitionLoop_NormalText(t *testing.T) { + + normal := "The quick brown fox jumps over the lazy dog. " + + "Pack my box with five dozen liquor jugs. " + + "How vexingly quick daft zebras jump. " + + "Sphinx of black quartz, judge my vow. " + + "Two driven jocks help fax my big quiz. " + + "The five boxing wizards jump quickly. " + + "Jackdaws love my big sphinx of quartz. " + + "Grumpy wizards make a toxic brew for the jovial queen." + + long := strings.Repeat(normal+" ", 10) + if DetectRepetitionLoop(long) { + t.Fatal("DetectRepetitionLoop should return false for normal text") + } +} + +func TestDetectRepetitionLoop_ShortText(t *testing.T) { + + if DetectRepetitionLoop("short") { + t.Fatal("DetectRepetitionLoop should return false for short text") + } +} + +func TestDetectRepetitionLoop_EmptyString(t *testing.T) { + if DetectRepetitionLoop("") { + t.Fatal("DetectRepetitionLoop should return false for empty string") + } +} + +func TestDetectRepetitionLoop_SingleCharRepeat(t *testing.T) { + + repeated := strings.Repeat("あ", 2500) + if !DetectRepetitionLoop(repeated) { + t.Fatal("DetectRepetitionLoop should return true for single-char repetition") + } +} + +func TestDetectRepetitionLoop_BelowSampleSize(t *testing.T) { + + phrase := "abcdefghij" + repeated := strings.Repeat(phrase, 50) + if !DetectRepetitionLoop(repeated) { + t.Fatal("DetectRepetitionLoop should return true for repetitive text below sample size") + } +} + +func TestTailPad_FewerThanN(t *testing.T) { + got := TailPad("a\nb", 5, 80) + lines := strings.Split(got, "\n") + if len(lines) != 5 { + t.Fatalf("TailPad line count = %d, want 5", len(lines)) + } + for i := 0; i < 3; i++ { + if lines[i] != "\u2800" { + t.Errorf("TailPad line %d = %q, want padding", i, lines[i]) + } + } + if lines[3] != "a" || lines[4] != "b" { + t.Errorf("TailPad content = %q %q, want a b", lines[3], lines[4]) + } +} + +func TestTailPad_ExactlyN(t *testing.T) { + in := "a\nb\nc" + got := TailPad(in, 3, 80) + if got != in { + t.Fatalf("TailPad exact = %q, want %q", got, in) + } +} + +func TestTailPad_MoreThanN(t *testing.T) { + got := TailPad("a\nb\nc\nd\ne", 3, 80) + if got != "c\nd\ne" { + t.Fatalf("TailPad tail = %q, want %q", got, "c\nd\ne") + } +} + +func TestTailPad_Empty(t *testing.T) { + got := TailPad("", 4, 80) + lines := strings.Split(got, "\n") + if len(lines) != 4 { + t.Fatalf("TailPad empty line count = %d, want 4", len(lines)) + } + for i, l := range lines { + if i == len(lines)-1 { + if l != "" { + t.Errorf("TailPad empty last line = %q, want empty", l) + } + } else if l != "\u2800" { + t.Errorf("TailPad empty line %d = %q, want padding", i, l) + } + } +} + +func TestTailPad_LongLineWraps(t *testing.T) { + + got := TailPad("abcdefghij", 4, 5) + lines := strings.Split(got, "\n") + if len(lines) != 4 { + t.Fatalf("TailPad wrap line count = %d, want 4", len(lines)) + } + + if lines[2] != "abcde" || lines[3] != "fghij" { + t.Errorf("TailPad wrap content = %v", lines) + } +} + +func TestTailPad_WrapPushesOldLines(t *testing.T) { + + got := TailPad("short\nabcdefghij", 2, 5) + if got != "abcde\nfghij" { + t.Fatalf("TailPad wrap push = %q, want %q", got, "abcde\nfghij") + } +} diff --git a/pkg/utils/string_test.go b/pkg/utils/string_test.go index 7b4b54098..e3b5af052 100644 --- a/pkg/utils/string_test.go +++ b/pkg/utils/string_test.go @@ -1,191 +1,6 @@ package utils -import ( - "strings" - "testing" -) - -// --- StripThinkBlocks --- - -func TestStripThinkBlocks_ClosedBlock(t *testing.T) { - in := "\nsecret reasoning\n\n\nVisible content" - got := StripThinkBlocks(in) - if got != "Visible content" { - t.Fatalf("StripThinkBlocks() = %q, want %q", got, "Visible content") - } -} - -func TestStripThinkBlocks_UnclosedBlock(t *testing.T) { - in := "reasoning that never ends\nmore reasoning" - got := StripThinkBlocks(in) - if got != "" { - t.Fatalf("StripThinkBlocks() = %q, want empty", got) - } -} - -func TestStripThinkBlocks_MultipleBlocks(t *testing.T) { - in := "firstmiddlesecondend" - got := StripThinkBlocks(in) - if got != "middleend" { - t.Fatalf("StripThinkBlocks() = %q, want %q", got, "middleend") - } -} - -func TestStripThinkBlocks_NoBlocks(t *testing.T) { - in := "plain text without think blocks" - got := StripThinkBlocks(in) - if got != in { - t.Fatalf("StripThinkBlocks() = %q, want %q", got, in) - } -} - -func TestStripThinkBlocks_CaseInsensitive(t *testing.T) { - in := "upper casevisible" - got := StripThinkBlocks(in) - if got != "visible" { - t.Fatalf("StripThinkBlocks() = %q, want %q", got, "visible") - } -} - -func TestStripThinkBlocks_ClosedThenUnclosed(t *testing.T) { - in := "closedmiddleunclosed tail" - got := StripThinkBlocks(in) - if got != "middle" { - t.Fatalf("StripThinkBlocks() = %q, want %q", got, "middle") - } -} - -// --- DetectRepetitionLoop --- - -func TestDetectRepetitionLoop_HighRepetition(t *testing.T) { - // Repeat a short phrase many times → should be detected - phrase := "結構本格的なコード" //nolint:gosmopolitan // CJK test data - repeated := strings.Repeat(phrase, 300) - if !DetectRepetitionLoop(repeated) { - t.Fatal("DetectRepetitionLoop should return true for highly repetitive text") - } -} - -func TestDetectRepetitionLoop_NormalText(t *testing.T) { - // Normal varied text should not trigger - normal := "The quick brown fox jumps over the lazy dog. " + - "Pack my box with five dozen liquor jugs. " + - "How vexingly quick daft zebras jump. " + - "Sphinx of black quartz, judge my vow. " + - "Two driven jocks help fax my big quiz. " + - "The five boxing wizards jump quickly. " + - "Jackdaws love my big sphinx of quartz. " + - "Grumpy wizards make a toxic brew for the jovial queen." - // Extend to be long enough - long := strings.Repeat(normal+" ", 10) - if DetectRepetitionLoop(long) { - t.Fatal("DetectRepetitionLoop should return false for normal text") - } -} - -func TestDetectRepetitionLoop_ShortText(t *testing.T) { - // Text shorter than N-gram size should never trigger - if DetectRepetitionLoop("short") { - t.Fatal("DetectRepetitionLoop should return false for short text") - } -} - -func TestDetectRepetitionLoop_EmptyString(t *testing.T) { - if DetectRepetitionLoop("") { - t.Fatal("DetectRepetitionLoop should return false for empty string") - } -} - -func TestDetectRepetitionLoop_SingleCharRepeat(t *testing.T) { - // "aaaa..." repeated → only 1 unique N-gram → detected - repeated := strings.Repeat("あ", 2500) - if !DetectRepetitionLoop(repeated) { - t.Fatal("DetectRepetitionLoop should return true for single-char repetition") - } -} - -func TestDetectRepetitionLoop_BelowSampleSize(t *testing.T) { - // Repetitive but under sample size still detected - phrase := "abcdefghij" - repeated := strings.Repeat(phrase, 50) // 500 chars - if !DetectRepetitionLoop(repeated) { - t.Fatal("DetectRepetitionLoop should return true for repetitive text below sample size") - } -} - -// --- TailPad --- - -func TestTailPad_FewerThanN(t *testing.T) { - got := TailPad("a\nb", 5, 80) - lines := strings.Split(got, "\n") - if len(lines) != 5 { - t.Fatalf("TailPad line count = %d, want 5", len(lines)) - } - for i := 0; i < 3; i++ { - if lines[i] != "\u2800" { - t.Errorf("TailPad line %d = %q, want padding", i, lines[i]) - } - } - if lines[3] != "a" || lines[4] != "b" { - t.Errorf("TailPad content = %q %q, want a b", lines[3], lines[4]) - } -} - -func TestTailPad_ExactlyN(t *testing.T) { - in := "a\nb\nc" - got := TailPad(in, 3, 80) - if got != in { - t.Fatalf("TailPad exact = %q, want %q", got, in) - } -} - -func TestTailPad_MoreThanN(t *testing.T) { - got := TailPad("a\nb\nc\nd\ne", 3, 80) - if got != "c\nd\ne" { - t.Fatalf("TailPad tail = %q, want %q", got, "c\nd\ne") - } -} - -func TestTailPad_Empty(t *testing.T) { - got := TailPad("", 4, 80) - lines := strings.Split(got, "\n") - if len(lines) != 4 { - t.Fatalf("TailPad empty line count = %d, want 4", len(lines)) - } - for i, l := range lines { - if i == len(lines)-1 { - if l != "" { - t.Errorf("TailPad empty last line = %q, want empty", l) - } - } else if l != "\u2800" { - t.Errorf("TailPad empty line %d = %q, want padding", i, l) - } - } -} - -func TestTailPad_LongLineWraps(t *testing.T) { - // One 10-char line wraps into 2 visual lines at width 5. - got := TailPad("abcdefghij", 4, 5) - lines := strings.Split(got, "\n") - if len(lines) != 4 { - t.Fatalf("TailPad wrap line count = %d, want 4", len(lines)) - } - // 2 padding + "abcde" + "fghij" - if lines[2] != "abcde" || lines[3] != "fghij" { - t.Errorf("TailPad wrap content = %v", lines) - } -} - -func TestTailPad_WrapPushesOldLines(t *testing.T) { - // "short" (1 visual) + "abcdefghij" (2 visual at width 5) = 3 visual. - // With n=2, only tail 2 visual lines remain. - got := TailPad("short\nabcdefghij", 2, 5) - if got != "abcde\nfghij" { - t.Fatalf("TailPad wrap push = %q, want %q", got, "abcde\nfghij") - } -} - -// --- Truncate --- +import "testing" func TestTruncate(t *testing.T) { tests := []struct {