opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: replace test files with upstream versions, extract fork-only tests

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
dj-oyu 2026-03-13 03:16:04 +09:00
parent c14e3769e2
commit 1c12de7630
68 changed files with 10498 additions and 9252 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
cmd/picoclaw/internal/helpers_ext_test.go Normal file
View file

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

73
cmd/picoclaw/internal/helpers_test.go
View file

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

3
cmd/picoclaw/main_test.go
View file

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

206
pkg/agent/context_cache_test.go
View file

@ -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
history []providers.Message
summary string
message string
}{
{
name: "no summary, no history",
summary: "",
message: "hello",
},
{
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
checkField string // substring to verify in rebuilt prompt
}{
{
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.",
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",
})
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
checkField string // substring to verify in rebuilt prompt
}{
{
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.",
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.",
"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")
}

62
pkg/agent/context_test.go
View file

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

52
pkg/agent/instance_ext_test.go Normal file
View file

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

117
pkg/agent/instance_test.go
View file

@ -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,
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,
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,
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) {
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",
},
{
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",
},
}
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)
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: tmpDir,
Model: "step-3.5-flash",
Model: tt.aliasName,
},
},
ModelList: []config.ModelConfig{
{
ModelName: "step-3.5-flash",
Model: "openrouter/stepfun/step-3.5-flash:free",
APIBase: "https://openrouter.ai/api/v1",
ModelName: tt.aliasName,
Model: tt.modelName,
APIBase: tt.apiBase,
},
},
}
provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider)
if len(agent.Candidates) != 1 {
t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates))
}
if agent.Candidates[0].Provider != "openrouter" {
t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, "openrouter")
if agent.Candidates[0].Provider != tt.wantProvider {
t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, tt.wantProvider)
}
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 agent.Candidates[0].Model != tt.wantModel {
t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, tt.wantModel)
}
})
}
}

2918
pkg/agent/loop_ext_test.go Normal file
View file

File diff suppressed because it is too large Load diff

4080
pkg/agent/loop_test.go
View file

File diff suppressed because it is too large Load diff

5
pkg/agent/mock_provider_test.go
View file

@ -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",
ToolCalls: []providers.ToolCall{},
}, nil
}

51
pkg/agent/registry_test.go
View file

@ -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
@ -33,14 +29,10 @@ func testCfg(agents []config.AgentConfig) *config.Config {
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
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")
}
@ -140,35 +115,26 @@ func TestAgentRegistry_CanSpawnSubagent(t *testing.T) {
cfg := testCfg([]config.AgentConfig{
{
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)")
}
@ -178,23 +144,18 @@ func TestAgentRegistry_CanSpawnSubagent_Wildcard(t *testing.T) {
cfg := testCfg([]config.AgentConfig{
{
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))
}
@ -235,21 +190,15 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) {
func TestAgentInstance_FallbackExplicitEmpty(t *testing.T) {
model := &config.AgentModelConfig{
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)
}

881
pkg/channels/manager_ext_test.go Normal file
View file

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

1075
pkg/channels/manager_test.go
View file

File diff suppressed because it is too large Load diff

96
pkg/channels/matrix/matrix_test.go
View file

@ -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**", "<strong>hello</strong>"},
{"italic", "_world_", "<em>world</em>"},
{"header", "### Title", "<h3"},
{"code block", "```\nfoo()\n```", "<code>"},
{"inline code", "`x`", "<code>x</code>"},
{"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, "<strong>hi</strong>") {
t.Errorf("format %q: FormattedBody %q missing <strong>", 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)
}
}

52
pkg/channels/telegram/telegram_ext_test.go Normal file
View file

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

482
pkg/channels/telegram/telegram_test.go
View file

@ -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"
"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)
}
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)
type stubCall struct {
URL string
Data *ta.RequestData
}
return
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)
}
if err != nil {
t.Fatalf("parseChatID(%q) unexpected error: %v", tc.input, err)
// stubConstructor implements ta.RequestConstructor for testing.
type stubConstructor struct{}
func (s *stubConstructor) JSONRequest(parameters any) (*ta.RequestData, error) {
return &ta.RequestData{}, nil
}
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 (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 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 TestFormatChatID(t *testing.T) {
if got := formatChatID(-100, 42); got != "-100/42" {
t.Fatalf("formatChatID(-100, 42) = %q, want %q", got, "-100/42")
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
},
}
if got := formatChatID(12345, 0); got != "12345" {
t.Fatalf("formatChatID(12345, 0) = %q, want %q", got, "12345")
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 "<b>a</b> " (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"])
}

6
pkg/channels/wecom/app_test.go
View file

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

7
pkg/channels/wecom/bot_test.go
View file

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

2
pkg/channels/wecom/common.go
View file

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

123
pkg/config/config_ext_test.go Normal file
View file

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

264
pkg/config/config_test.go
View file

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

1
pkg/config/defaults.go
View file

@ -413,6 +413,7 @@ func DefaultConfig() *Config {
Enabled: true,
},
EnableDenyPatterns: true,
AllowRemote: true,
TimeoutSeconds: 60,
},
Skills: SkillsToolsConfig{

6
pkg/config/migration.go
View file

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

21
pkg/config/migration_test.go
View file

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

103
pkg/heartbeat/service_ext_test.go Normal file
View file

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

95
pkg/heartbeat/service_test.go
View file

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

289
pkg/logger/logger_ext_test.go Normal file
View file

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

287
pkg/logger/logger_test.go
View file

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

52
pkg/memory/migration_test.go
View file

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

1
pkg/migrate/sources/openclaw/openclaw_config.go
View file

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

14
pkg/migrate/sources/openclaw/openclaw_config_test.go
View file

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

6
pkg/providers/anthropic/provider_test.go
View file

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

2
pkg/providers/antigravity_provider_test.go
View file

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

283
pkg/providers/claude_cli_provider_ext_test.go Normal file
View file

@ -0,0 +1,283 @@
package providers
import (
"strings"
"testing"
)
func TestExtractXMLToolCalls_Single(t *testing.T) {
text := `<vendor:toolcall>
<invoke name="exec">
<parameter name="command">echo hello</parameter>
</invoke>
</vendor:toolcall>`
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 := `<vendor:toolcall>
<invoke name="web_search">
<parameter name="query">golang testing</parameter>
</invoke>
<invoke name="exec">
<parameter name="command">go test ./...</parameter>
<parameter name="timeout">30</parameter>
</invoke>
</vendor:toolcall>`
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.
<vendor:toolcall>
<invoke name="exec">
<parameter name="command">echo hello</parameter>
</invoke>
</vendor:toolcall>
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 := `<minimax:toolcall>
<invoke name="readfile">
<parameter name="path">/home/user/project/pyproject.toml</parameter>
</invoke>
</minimax:tool_call>`
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 := `今テスト走らせるね。` +
`
<minimax:toolcall>
<invoke name="exec">
<parameter name="command">cd /home/user && pytest</parameter>
</invoke>
</minimax:tool_call>`
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 := `<minimax:tool_call>
<invoke name="exec">
<parameter name="command">ls -la</parameter>
</invoke>
</minimax:tool_call>`
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 := `<vendor:Tool-Call>
<invoke name="read_file">
<parameter name="path">/etc/hosts</parameter>
</invoke>
</vendor:tool-call>`
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.
<minimax:tool_call>
<invoke name="exec">
<parameter name="command">ls</parameter>
</invoke>
</minimax:toolcall>
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<invoke name=\"listdir\">\n<parameter name=\"path\">/home/user/workspace</parameter>\n</invoke>\n</minimax:tool_call>"
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<invoke name=\"listdir\">\n<parameter name=\"path\">/home/user</parameter>\n</invoke>\n</minimax:tool_call>"
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)
}
}
}

289
pkg/providers/claude_cli_provider_test.go
View file

@ -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 := `<vendor:toolcall>
<invoke name="exec">
<parameter name="command">echo hello</parameter>
</invoke>
</vendor:toolcall>`
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 := `<vendor:toolcall>
<invoke name="web_search">
<parameter name="query">golang testing</parameter>
</invoke>
<invoke name="exec">
<parameter name="command">go test ./...</parameter>
<parameter name="timeout">30</parameter>
</invoke>
</vendor:toolcall>`
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.
<vendor:toolcall>
<invoke name="exec">
<parameter name="command">echo hello</parameter>
</invoke>
</vendor:toolcall>
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 <minimax:toolcall> but closes with </minimax:tool_call> (underscore)
text := `<minimax:toolcall>
<invoke name="readfile">
<parameter name="path">/home/user/project/pyproject.toml</parameter>
</invoke>
</minimax:tool_call>`
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
`
<minimax:toolcall>
<invoke name="exec">
<parameter name="command">cd /home/user && pytest</parameter>
</invoke>
</minimax:tool_call>`
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: <minimax:tool_call>
text := `<minimax:tool_call>
<invoke name="exec">
<parameter name="command">ls -la</parameter>
</invoke>
</minimax:tool_call>`
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: <vendor:Tool-Call>
text := `<vendor:Tool-Call>
<invoke name="read_file">
<parameter name="path">/etc/hosts</parameter>
</invoke>
</vendor:tool-call>`
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.
<minimax:tool_call>
<invoke name="exec">
<parameter name="command">ls</parameter>
</invoke>
</minimax:toolcall>
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 + <invoke> with orphaned closing tag (no opening tag)
text := "了解！確認するね。\n[TOOLCALL]\n<invoke name=\"listdir\">\n<parameter name=\"path\">/home/user/workspace</parameter>\n</invoke>\n</minimax:tool_call>" //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<invoke name=\"listdir\">\n<parameter name=\"path\">/home/user</parameter>\n</invoke>\n</minimax:tool_call>" //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)
}
}
}

12
pkg/providers/codex_cli_provider_test.go
View file

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

18
pkg/providers/codex_provider_test.go
View file

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

73
pkg/providers/factory_ext_test.go Normal file
View file

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

24
pkg/providers/factory_provider_test.go
View file

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

86
pkg/providers/factory_test.go
View file

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

456
pkg/providers/openai_compat/provider_ext_test.go Normal file
View file

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

864
pkg/providers/openai_compat/provider_test.go
View file

@ -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: "<!DOCTYPE html><html><body>gateway login</body></html>",
},
{
name: "html error response",
contentType: "text/html; charset=utf-8",
statusCode: http.StatusBadGateway,
body: "<!DOCTYPE html><html><body>bad gateway</body></html>",
},
{
name: "mislabeled html success response",
contentType: "application/json",
statusCode: http.StatusOK,
body: " \r\n\t<!DOCTYPE html><html><body>gateway login</body></html>",
},
}
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("<!DOCTYPE html><html><body>"), bytes.Repeat([]byte("A"), 2048)...)
body = append(body, []byte("</body></html>")...)
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: <!DOCTYPE html><html><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")
}
if got := normalizeModel("vivgrid/managed", "https://api.vivgrid.com/v1"); got != "managed" {
t.Fatalf("normalizeModel(vivgrid) = %q, want %q", got, "managed")
}
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) // 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")
}
}

5
pkg/session/manager.go
View file

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

121
pkg/session/manager_ext_test.go Normal file
View file

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

163
pkg/session/manager_test.go
View file

@ -4,33 +4,25 @@ import (
"os"
"path/filepath"
"testing"
"github.com/sipeed/picoclaw/pkg/providers"
)
func TestSanitizeFilename(t *testing.T) {
tests := []struct {
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)")
}
}

75
pkg/skills/loader_test.go
View file

@ -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# COPYRIGHT NOTICE\n# This file is part of the \"Universal Biomedical Skills\" project.\n# Copyright (c) 2026 MD BABU MIA, PhD <md.babu.mia@mssm.edu>\n# All Rights Reserved.\n#\n# This code is proprietary and confidential.\n# Unauthorized copying of this file, via any medium is strictly prohibited.\n#\n# Provenance: Authenticated by MD BABU MIA\n\n-->\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)
}

4
pkg/state/state.go
View file

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

38
pkg/state/state_ext_test.go Normal file
View file

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

44
pkg/state/state_test.go
View file

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

172
pkg/tools/edit_test.go
View file

@ -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,
"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,
"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,
"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,
"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,
"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",
"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",
"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,
"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
expectError bool
}{
{
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,
expectError: true,
},
{
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",
"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,
"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,
"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",
"old_text": "old",
"new_text": "new",
}
result := tool.Execute(ctx, args)
assert.True(t, result.IsError)
assert.Contains(t, result.ForLLM, "not found")
}

183
pkg/tools/filesystem_ext_test.go Normal file
View file

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

369
pkg/tools/filesystem_test.go
View file

@ -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,
"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,
"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,
"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)
assert.Contains(t, err.Error(), "access denied")
assert.Contains(t, err.Error(), workspace)
// 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)
}
// Read from non-whitelisted path outside workspace should fail.
otherDir := t.TempDir()
otherFile := filepath.Join(otherDir, "blocked.txt")
os.WriteFile(otherFile, []byte("blocked"), 0o644)
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)
}
}

81
pkg/tools/message_test.go
View file

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

269
pkg/tools/registry_ext_test.go Normal file
View file

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

360
pkg/tools/registry_test.go
View file

@ -13,46 +13,36 @@ import (
type mockRegistryTool struct {
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) 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 ---
@ -60,52 +50,19 @@ func (m *mockAsyncRegistryTool) SetCallback(cb AsyncCallback) {
func newMockTool(name, desc string) *mockRegistryTool {
return &mockRegistryTool{
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",
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",
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",
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)
}

50
pkg/tools/result_test.go
View file

@ -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")
}
@ -115,36 +98,26 @@ func TestUserResult(t *testing.T) {
func TestToolResultJSONSerialization(t *testing.T) {
tests := []struct {
name string
result *ToolResult
}{
{
name: "basic result",
result: NewToolResult("basic content"),
},
{
name: "silent result",
result: SilentResult("silent content"),
},
{
name: "async result",
result: AsyncResult("async content"),
},
{
name: "error result",
result: ErrorResult("error content"),
},
{
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"])
}

896
pkg/tools/shell_ext_test.go Normal file
View file

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

1242
pkg/tools/shell_test.go
View file

File diff suppressed because it is too large Load diff

35
pkg/tools/shell_timeout_unix_ext_test.go Normal file
View file

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

31
pkg/tools/shell_timeout_unix_test.go
View file

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

35
pkg/tools/skills_install_test.go
View file

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

30
pkg/tools/skills_search_test.go
View file

@ -11,81 +11,62 @@ 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")
}
@ -93,28 +74,17 @@ func TestFormatSearchResultsWithData(t *testing.T) {
results := []skills.SearchResult{
{
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")
}

32
pkg/tools/spawn_test.go
View file

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

102
pkg/tools/subagent_tool_ext_test.go Normal file
View file

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

259
pkg/tools/subagent_tool_test.go
View file

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

751
pkg/tools/web_test.go
View file

File diff suppressed because it is too large Load diff

179
pkg/utils/string_ext_test.go Normal file
View file

@ -0,0 +1,179 @@
package utils
import (
"strings"
"testing"
)
func TestStripThinkBlocks_ClosedBlock(t *testing.T) {
in := "<think>\nsecret reasoning\n</think>\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 := "<think>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 := "<think>first</think>middle<think>second</think>end"
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 := "<THINK>upper case</THINK>visible"
got := StripThinkBlocks(in)
if got != "visible" {
t.Fatalf("StripThinkBlocks() = %q, want %q", got, "visible")
}
}
func TestStripThinkBlocks_ClosedThenUnclosed(t *testing.T) {
in := "<think>closed</think>middle<think>unclosed 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")
}
}

187
pkg/utils/string_test.go
View file

@ -1,191 +1,6 @@
package utils
import (
"strings"
"testing"
)
// --- StripThinkBlocks ---
func TestStripThinkBlocks_ClosedBlock(t *testing.T) {
in := "<think>\nsecret reasoning\n</think>\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 := "<think>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 := "<think>first</think>middle<think>second</think>end"
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 := "<THINK>upper case</THINK>visible"
got := StripThinkBlocks(in)
if got != "visible" {
t.Fatalf("StripThinkBlocks() = %q, want %q", got, "visible")
}
}
func TestStripThinkBlocks_ClosedThenUnclosed(t *testing.T) {
in := "<think>closed</think>middle<think>unclosed 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 {