tool set refactor
This commit is contained in:
Luxferre
+113
-116
@@ -6,10 +6,12 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
@@ -183,7 +185,7 @@ func TestParseStreamToolCallSplitAcrossChunks(t *testing.T) {
|
||||
|
||||
func TestParseStreamMultipleToolCallsKeepFirstAppearanceOrder(t *testing.T) {
|
||||
sseData := strings.Join([]string{
|
||||
`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"c2","function":{"name":"run_subagent","arguments":"{\"prompt\":\"p\"}"}}]}}]}`,
|
||||
`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"c2","function":{"name":"write_file","arguments":"{\"path\":\"a.txt\",\"content\":\"hello\"}"}}]}}]}`,
|
||||
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"shell_exec","arguments":"{\"command\":\"ls\"}"}}]}}]}`,
|
||||
`data: [DONE]`,
|
||||
}, "\n")
|
||||
@@ -240,24 +242,21 @@ func TestParseStreamReasoningAfterContent(t *testing.T) {
|
||||
|
||||
// ---------- prompt / getCfg / atoiD ----------
|
||||
|
||||
func TestDefaultSystemPromptFallback(t *testing.T) {
|
||||
p := prompt("non_existent_file.txt")
|
||||
if !strings.Contains(p, "You are Bantam, a tiny, powerful AI agent.") {
|
||||
t.Errorf("expected prompt to contain base description, got: %q", p)
|
||||
func TestDefaultSystemPrompt(t *testing.T) {
|
||||
if !strings.Contains(defaultSystemPrompt, "You are Bantam, a tiny, powerful AI agent.") {
|
||||
t.Errorf("expected prompt to contain base description, got: %q", defaultSystemPrompt)
|
||||
}
|
||||
if !strings.Contains(p, "shell_exec") || !strings.Contains(p, "run_subagent") {
|
||||
t.Errorf("expected prompt to list shell_exec and run_subagent tools, got: %q", p)
|
||||
if !strings.Contains(defaultSystemPrompt, "shell_exec") || !strings.Contains(defaultSystemPrompt, "write_file") {
|
||||
t.Errorf("expected prompt to list shell_exec and write_file tools, got: %q", defaultSystemPrompt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptReadsAndTrimsFile(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "sys.txt")
|
||||
if err := os.WriteFile(p, []byte(" hello\nworld \n"), 0644); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
if strings.Contains(defaultSystemPrompt, "run_subagent") {
|
||||
t.Errorf("prompt should not mention run_subagent: %q", defaultSystemPrompt)
|
||||
}
|
||||
got := prompt(p)
|
||||
if got != "hello\nworld" {
|
||||
t.Errorf("expected trimmed content, got %q", got)
|
||||
if !strings.Contains(defaultSystemPrompt, "AGENTS.md") {
|
||||
t.Errorf("expected prompt to mention AGENTS.md, got: %q", defaultSystemPrompt)
|
||||
}
|
||||
if strings.Contains(defaultSystemPrompt, "GEMINI.md") || strings.Contains(defaultSystemPrompt, "CLAUDE.md") {
|
||||
t.Errorf("prompt should not mention GEMINI.md or CLAUDE.md: %q", defaultSystemPrompt)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1100,7 +1099,7 @@ func TestALToolLoop(t *testing.T) {
|
||||
cfg.Endpoint = srv.URL
|
||||
cfg.Stream = false
|
||||
cfg.APIKey = "-"
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("run")}}, "sys", 0)
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("run")}})
|
||||
if err != nil {
|
||||
t.Fatalf("AL: %v", err)
|
||||
}
|
||||
@@ -1149,7 +1148,7 @@ func TestALReasoningOnlyAutoContinue(t *testing.T) {
|
||||
cfg.Endpoint = srv.URL
|
||||
cfg.Stream = false
|
||||
cfg.APIKey = "-"
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("start")}}, "sys", 0)
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("start")}})
|
||||
if err != nil {
|
||||
t.Fatalf("AL: %v", err)
|
||||
}
|
||||
@@ -1167,37 +1166,92 @@ func TestALReasoningOnlyAutoContinue(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestALRunSubagent(t *testing.T) {
|
||||
var n int
|
||||
var mu sync.Mutex
|
||||
func TestWriteFile(t *testing.T) {
|
||||
tmp := t.TempDir()
|
||||
target := filepath.Join(tmp, "sub", "dir", "test.txt")
|
||||
|
||||
// 1. Create file and write initial content
|
||||
res, err := writeFile(target, 0, 0, "Hello World")
|
||||
if err != nil {
|
||||
t.Fatalf("writeFile create: %v", err)
|
||||
}
|
||||
if !strings.Contains(res, "Successfully wrote 11 bytes") {
|
||||
t.Errorf("unexpected res: %q", res)
|
||||
}
|
||||
data, err := os.ReadFile(target)
|
||||
if err != nil || string(data) != "Hello World" {
|
||||
t.Fatalf("read = %q, want Hello World", string(data))
|
||||
}
|
||||
|
||||
// 2. Overwrite / replace "World" with "Bantam" (offset 6, del_bytes 5)
|
||||
res, err = writeFile(target, 6, 5, "Bantam")
|
||||
if err != nil {
|
||||
t.Fatalf("writeFile replace: %v", err)
|
||||
}
|
||||
data, _ = os.ReadFile(target)
|
||||
if string(data) != "Hello Bantam" {
|
||||
t.Fatalf("read = %q, want Hello Bantam", string(data))
|
||||
}
|
||||
|
||||
// 3. Insert without deletion (offset 5, del_bytes 0, content " dear")
|
||||
res, err = writeFile(target, 5, 0, " dear")
|
||||
if err != nil {
|
||||
t.Fatalf("writeFile insert: %v", err)
|
||||
}
|
||||
data, _ = os.ReadFile(target)
|
||||
if string(data) != "Hello dear Bantam" {
|
||||
t.Fatalf("read = %q, want Hello dear Bantam", string(data))
|
||||
}
|
||||
|
||||
// 4. Pure deletion (offset 5, del_bytes 5, content "")
|
||||
res, err = writeFile(target, 5, 5, "")
|
||||
if err != nil {
|
||||
t.Fatalf("writeFile delete: %v", err)
|
||||
}
|
||||
data, _ = os.ReadFile(target)
|
||||
if string(data) != "Hello Bantam" {
|
||||
t.Fatalf("read = %q, want Hello Bantam", string(data))
|
||||
}
|
||||
|
||||
// 5. Offset beyond file length -> padded with null bytes
|
||||
res, err = writeFile(target, 15, 0, "end")
|
||||
if err != nil {
|
||||
t.Fatalf("writeFile beyond len: %v", err)
|
||||
}
|
||||
data, _ = os.ReadFile(target)
|
||||
if len(data) != 18 || !strings.HasSuffix(string(data), "end") {
|
||||
t.Fatalf("read length = %d, want 18", len(data))
|
||||
}
|
||||
|
||||
// 6. Error on empty path
|
||||
_, err = writeFile("", 0, 0, "abc")
|
||||
if err == nil {
|
||||
t.Fatalf("expected error for empty path")
|
||||
}
|
||||
}
|
||||
|
||||
func TestALWriteFile(t *testing.T) {
|
||||
tmp := t.TempDir()
|
||||
target := filepath.Join(tmp, "out.txt")
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
n++
|
||||
cur := n
|
||||
mu.Unlock()
|
||||
var req struct {
|
||||
Messages []Message `json:"messages"`
|
||||
}
|
||||
json.NewDecoder(r.Body).Decode(&req)
|
||||
switch cur {
|
||||
case 1: // parent's first call -> delegate to subagent
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"run_subagent","arguments":"{\"prompt\":\"inner task\"}"}}]}}]}`))
|
||||
case 2: // child's call -> verify it got the child system prompt, then finish
|
||||
isChild := false
|
||||
for _, m := range req.Messages {
|
||||
if m.Role == "system" && m.Content != nil && strings.Contains(*m.Content, "Important: this is a child agent") {
|
||||
isChild = true
|
||||
}
|
||||
for _, m := range req.Messages {
|
||||
if m.Role == "tool" {
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"file is ready"}}]}`))
|
||||
return
|
||||
}
|
||||
if !isChild {
|
||||
t.Errorf("request #2 did not carry the child system prompt")
|
||||
}
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"child done"}}]}`))
|
||||
case 3: // parent's second call -> finish
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"parent done"}}]}`))
|
||||
default:
|
||||
t.Errorf("unexpected request #%d", cur)
|
||||
}
|
||||
args, _ := json.Marshal(map[string]any{
|
||||
"path": target,
|
||||
"offset": 0,
|
||||
"del_bytes": 0,
|
||||
"content": "sample file content",
|
||||
})
|
||||
w.Write([]byte(fmt.Sprintf(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"write_file","arguments":%s}}]}}]}`, strconv.Quote(string(args)))))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
@@ -1205,67 +1259,16 @@ func TestALRunSubagent(t *testing.T) {
|
||||
cfg.Endpoint = srv.URL
|
||||
cfg.Stream = false
|
||||
cfg.APIKey = "-"
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("parent task")}}, "sys", 0)
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("write a file")}})
|
||||
if err != nil {
|
||||
t.Fatalf("AL: %v", err)
|
||||
}
|
||||
if got := last(msgs); got != "parent done" {
|
||||
t.Errorf("last = %q, want parent done", got)
|
||||
if got := last(msgs); got != "file is ready" {
|
||||
t.Errorf("last = %q, want 'file is ready'", got)
|
||||
}
|
||||
found := false
|
||||
for _, m := range msgs {
|
||||
if m.Role == "tool" && m.Content != nil && strings.Contains(*m.Content, "child done") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("expected subagent result 'child done' in tool messages")
|
||||
}
|
||||
}
|
||||
|
||||
func TestALSubagentDepthLimit(t *testing.T) {
|
||||
// depth >= MAX_DEPTH must not spawn a child request; the LLM gets the
|
||||
// depth-limit tool result and is expected to acknowledge and finish.
|
||||
var n int
|
||||
var mu sync.Mutex
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
mu.Lock()
|
||||
n++
|
||||
cur := n
|
||||
mu.Unlock()
|
||||
switch cur {
|
||||
case 1:
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"run_subagent","arguments":"{\"prompt\":\"deep\"}"}}]}}]}`))
|
||||
case 2:
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"done"}}]}`))
|
||||
default:
|
||||
t.Errorf("unexpected request #%d (child must not be spawned)", cur)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := defCfg
|
||||
cfg.Endpoint = srv.URL
|
||||
cfg.Stream = false
|
||||
cfg.APIKey = "-"
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "user", Content: strp("go")}}, "sys", MAX_DEPTH)
|
||||
if err != nil {
|
||||
t.Fatalf("AL: %v", err)
|
||||
}
|
||||
if n != 2 {
|
||||
t.Errorf("expected exactly 2 LLM requests, got %d", n)
|
||||
}
|
||||
if got := last(msgs); got != "done" {
|
||||
t.Errorf("last = %q, want done", got)
|
||||
}
|
||||
found := false
|
||||
for _, m := range msgs {
|
||||
if m.Role == "tool" && m.Content != nil && strings.Contains(*m.Content, "depth limit") {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("expected depth-limit tool result")
|
||||
content, err := os.ReadFile(target)
|
||||
if err != nil || string(content) != "sample file content" {
|
||||
t.Errorf("file content = %q, want 'sample file content'", string(content))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1295,7 +1298,7 @@ func TestALStripsInvalidAssistantAndRetries(t *testing.T) {
|
||||
cfg.Endpoint = srv.URL
|
||||
cfg.Stream = false
|
||||
cfg.APIKey = "-"
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("go")}}, "sys", 0)
|
||||
msgs, _, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("go")}})
|
||||
if err != nil {
|
||||
t.Fatalf("AL: %v", err)
|
||||
}
|
||||
@@ -1714,7 +1717,7 @@ func TestALContextCancellation(t *testing.T) {
|
||||
|
||||
origMsgs := []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("hello")}}
|
||||
inputMsgs := append([]Message{}, origMsgs...)
|
||||
msgs, _, err := AL(ctx, &cfg, inputMsgs, "sys", 0)
|
||||
msgs, _, err := AL(ctx, &cfg, inputMsgs)
|
||||
if err == nil {
|
||||
t.Fatalf("expected context cancellation error, got nil")
|
||||
}
|
||||
@@ -2048,8 +2051,8 @@ func TestSanitizeMessagesWithInvisibles(t *testing.T) {
|
||||
|
||||
// ---------- new coverage from review ----------
|
||||
|
||||
// #14: subagent token usage must be accumulated into the parent turn usage.
|
||||
func TestALRunSubagentAccumulatesUsage(t *testing.T) {
|
||||
// #14: tool call loop token usage must be accumulated into turn usage.
|
||||
func TestALToolLoopAccumulatesUsage(t *testing.T) {
|
||||
var n int
|
||||
var mu sync.Mutex
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
@@ -2059,11 +2062,9 @@ func TestALRunSubagentAccumulatesUsage(t *testing.T) {
|
||||
mu.Unlock()
|
||||
switch cur {
|
||||
case 1:
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"run_subagent","arguments":"{\"prompt\":\"inner\"}"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,"prompt_tokens_details":{"cached_tokens":7}}}`))
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"shell_exec","arguments":"{\"command\":\"echo 1\"}"}}]}}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,"prompt_tokens_details":{"cached_tokens":7}}}`))
|
||||
case 2:
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"child done"}}],"usage":{"prompt_tokens":20,"completion_tokens":3,"total_tokens":23,"prompt_tokens_details":{"cached_tokens":12}}}`))
|
||||
case 3:
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"parent done"}}],"usage":{"prompt_tokens":30,"completion_tokens":4,"total_tokens":34,"prompt_tokens_details":{"cached_tokens":15}}}`))
|
||||
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"done"}}],"usage":{"prompt_tokens":30,"completion_tokens":4,"total_tokens":34,"prompt_tokens_details":{"cached_tokens":15}}}`))
|
||||
default:
|
||||
t.Errorf("unexpected request #%d", cur)
|
||||
}
|
||||
@@ -2074,22 +2075,18 @@ func TestALRunSubagentAccumulatesUsage(t *testing.T) {
|
||||
cfg.Endpoint = srv.URL
|
||||
cfg.Stream = false
|
||||
cfg.APIKey = "-"
|
||||
_, usg, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("parent task")}}, "sys", 0)
|
||||
_, usg, err := AL(context.Background(), &cfg, []Message{{Role: "system", Content: strp("sys")}, {Role: "user", Content: strp("parent task")}})
|
||||
if err != nil {
|
||||
t.Fatalf("AL: %v", err)
|
||||
}
|
||||
// PromptTokens/Cached reflect the parent's own final context (30 / 15); they are
|
||||
// overwritten per iteration, not accumulated. Completion/Total accumulate every
|
||||
// assistant call: parent's tool-call call (5/15) + child (3/23) + parent final (4/34)
|
||||
// => completion 12, total 72.
|
||||
if usg.PromptTokens != 30 {
|
||||
t.Errorf("PromptTokens = %d, want 30", usg.PromptTokens)
|
||||
}
|
||||
if usg.CompletionTokens != 12 {
|
||||
t.Errorf("CompletionTokens = %d, want 12", usg.CompletionTokens)
|
||||
if usg.CompletionTokens != 9 {
|
||||
t.Errorf("CompletionTokens = %d, want 9", usg.CompletionTokens)
|
||||
}
|
||||
if usg.TotalTokens != 72 {
|
||||
t.Errorf("TotalTokens = %d, want 72", usg.TotalTokens)
|
||||
if usg.TotalTokens != 49 {
|
||||
t.Errorf("TotalTokens = %d, want 49", usg.TotalTokens)
|
||||
}
|
||||
if usg.Cached() != 15 {
|
||||
t.Errorf("CachedTokens = %d, want 15", usg.Cached())
|
||||
|
||||
Reference in New Issue
Block a user