1package agent
2
3import (
4 "encoding/json"
5 "errors"
6 "fmt"
7 "net/http"
8 "net/http/httptest"
9 "strings"
10 "sync/atomic"
11 "testing"
12
13 "github.com/bnema/zut/packages/agent/subagents"
14 "github.com/bnema/zut/packages/core"
15 "github.com/bnema/zut/packages/provider"
16)
17
18func TestResidentChildRegistryUsesExactToolListAndForbidsDelegation(t *testing.T) {
19 catalogue := core.Registry{
20 "read": nil,
21 "bash": nil,
22 "subagent_spawn": nil,
23 "update_goal": nil,
24 }
25 registry, err := residentChildRegistry(catalogue, []string{"read"})
26 if err != nil {
27 t.Fatal(err)
28 }
29 if len(registry) != 1 {
30 t.Fatalf("registry = %#v", registry)
31 }
32 for _, name := range []string{"subagent_spawn", "update_goal", "missing"} {
33 _, err := residentChildRegistry(catalogue, []string{name})
34 if err == nil || !strings.Contains(err.Error(), name) {
35 t.Fatalf("residentChildRegistry(%q) error = %v", name, err)
36 }
37 }
38}
39
40func TestConfigureResidentBudgetOnlyStopsAtExhaustion(t *testing.T) {
41 agent := core.NewAgent(nil, "model", "system", core.Registry{"read": nil})
42 budget := subagents.NewRolloutBudget(1_000, provider.Usage{})
43 configureResidentBudget(agent, budget)
44
45 budget.Observe(provider.Usage{InputTokens: 900})
46 allowed, reason, contextText := agent.BeforeTurnContext(t.Context(), 1)
47 if !allowed || reason != "" || contextText != "" {
48 t.Fatalf("near-limit context = allowed=%t reason=%q context=%q", allowed, reason, contextText)
49 }
50 if len(agent.ToolsSnapshot()) != 1 || agent.MaxTokens != 0 {
51 t.Fatalf("near-limit agent = tools=%#v max_tokens=%d", agent.ToolsSnapshot(), agent.MaxTokens)
52 }
53
54 budget.Observe(provider.Usage{InputTokens: 1_000})
55 allowed, reason, _ = agent.BeforeTurnContext(t.Context(), 2)
56 if allowed || reason != subagents.ErrBudgetExceeded.Error() {
57 t.Fatalf("exhausted context = allowed=%t reason=%q", allowed, reason)
58 }
59}
60
61func TestResidentChildRunnerKeepsToolsAtNinetyPercent(t *testing.T) {
62 var requests atomic.Int32
63 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
64 requestNumber := requests.Add(1)
65 var body map[string]any
66 if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
67 t.Errorf("decode request: %v", err)
68 }
69 w.Header().Set("content-type", "text/event-stream")
70 if requestNumber == 1 {
71 if got := requestMaxTokens(t, body); got != 16_384 {
72 t.Errorf("initial max_tokens = %d, want 16384", got)
73 }
74 writeOpenAIChunk(t, w, map[string]any{
75 "choices": []any{map[string]any{
76 "index": 0,
77 "delta": map[string]any{"tool_calls": []any{map[string]any{
78 "index": 0, "id": "call-1", "type": "function",
79 "function": map[string]string{"name": "read", "arguments": `{"path":"x"}`},
80 }}},
81 "finish_reason": "tool_calls",
82 }},
83 "usage": map[string]int{"prompt_tokens": 900, "completion_tokens": 0},
84 })
85 return
86 }
87 if tools, ok := body["tools"].([]any); !ok || len(tools) == 0 {
88 t.Errorf("near-limit request tools = %#v, want registered tools", body["tools"])
89 }
90 if got := requestMaxTokens(t, body); got <= 100 {
91 t.Errorf("near-limit max_tokens = %d, want model output limit", got)
92 }
93 writeOpenAIChunk(t, w, map[string]any{
94 "choices": []any{map[string]any{"index": 0, "delta": map[string]string{"content": "final result"}, "finish_reason": "stop"}},
95 "usage": map[string]int{"prompt_tokens": 0, "completion_tokens": 50},
96 })
97 }))
98 defer server.Close()
99
100 runner, journal := newBudgetedResidentTestRunner(t, server.URL, 1_000)
101 if err := runner(t.Context(), "review"); err != nil {
102 t.Fatal(err)
103 }
104 if got := requests.Load(); got != 2 {
105 t.Fatalf("provider requests = %d, want 2", got)
106 }
107 if summary := latestResidentAssistantText(t, journal); summary != "final result" {
108 t.Fatalf("latest summary = %q", summary)
109 }
110}
111
112func TestResidentChildRunnerTerminatesAtBudgetLimit(t *testing.T) {
113 var requests atomic.Int32
114 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
115 requests.Add(1)
116 var body map[string]any
117 if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
118 t.Errorf("decode request: %v", err)
119 }
120 if got := requestMaxTokens(t, body); got != 16_384 {
121 t.Errorf("initial max_tokens = %d, want 16384", got)
122 }
123 w.Header().Set("content-type", "text/event-stream")
124 writeOpenAIChunk(t, w, map[string]any{
125 "choices": []any{map[string]any{"index": 0, "delta": map[string]string{"content": "best available result"}, "finish_reason": "stop"}},
126 "usage": map[string]int{"prompt_tokens": 1_200, "completion_tokens": 0},
127 })
128 }))
129 defer server.Close()
130
131 runner, journal := newBudgetedResidentTestRunner(t, server.URL, 1_000)
132 if err := runner(t.Context(), "review"); err != nil {
133 t.Fatalf("runner error = %v, want terminal result", err)
134 }
135 if got := requests.Load(); got != 1 {
136 t.Fatalf("provider requests = %d, want 1", got)
137 }
138 if summary := latestResidentAssistantText(t, journal); summary != "best available result" {
139 t.Fatalf("latest summary = %q", summary)
140 }
141 if err := runner(t.Context(), "continue"); !errors.Is(err, subagents.ErrBudgetExceeded) {
142 t.Fatalf("exhausted follow-up error = %v", err)
143 }
144 if got := requests.Load(); got != 1 {
145 t.Fatalf("exhausted follow-up made %d provider requests, want 1 total", got)
146 }
147}
148
149func newBudgetedResidentTestRunner(t *testing.T, baseURL string, limit int64) (subagents.ResidentTurnRunner, *subagents.ResidentJournal) {
150 t.Helper()
151 spec := subagents.ResidentChildSpec{
152 ID: "budget-child", SessionID: "budget-session", Provider: "openai", Model: "gpt-4o",
153 Tools: []string{"read"}, BudgetLimit: limit, BudgetSource: "model_context",
154 }
155 journal, err := subagents.OpenResidentJournal(t.TempDir(), spec.ID)
156 if err != nil {
157 t.Fatal(err)
158 }
159 t.Cleanup(func() { _ = journal.Close() })
160 if err := journal.Accept(spec, "review"); err != nil {
161 t.Fatal(err)
162 }
163 runner, err := newResidentChildRunner(Args{
164 Provider: "openai", Model: "gpt-4o", APIKey: "synthetic", BaseURL: baseURL,
165 CWD: t.TempDir(), NoContextFiles: true, NoSkill: true,
166 }, spec, journal)
167 if err != nil {
168 t.Fatal(err)
169 }
170 return runner, journal
171}
172
173func requestMaxTokens(t *testing.T, body map[string]any) int {
174 t.Helper()
175 value, ok := body["max_tokens"].(float64)
176 if !ok {
177 t.Errorf("request max_tokens = %#v", body["max_tokens"])
178 return 0
179 }
180 return int(value)
181}
182
183func latestResidentAssistantText(t *testing.T, journal *subagents.ResidentJournal) string {
184 t.Helper()
185 messages, err := subagents.ReadResidentTranscriptMessages(journal.Dir())
186 if err != nil {
187 t.Fatal(err)
188 }
189 for index := len(messages) - 1; index >= 0; index-- {
190 if messages[index].Role != provider.RoleAssistant {
191 continue
192 }
193 var text strings.Builder
194 for _, content := range messages[index].Content {
195 if block, ok := content.(provider.TextBlock); ok {
196 text.WriteString(block.Text)
197 }
198 }
199 return text.String()
200 }
201 return ""
202}
203
204func writeOpenAIChunk(t *testing.T, w http.ResponseWriter, chunk map[string]any) {
205 t.Helper()
206 data, err := json.Marshal(chunk)
207 if err != nil {
208 t.Fatal(err)
209 }
210 _, _ = fmt.Fprintf(w, "data: %s\n\ndata: [DONE]\n\n", data)
211}
212
213func TestResidentChildArgsPreserveDurableProfileInheritance(t *testing.T) {
214 no := false
215 next := residentChildArgs(Args{Orchestrate: true, NoSkill: false, NoContextFiles: false, BaseURL: "https://stale.example/v1", InsecureTLS: false}, "openai", subagents.ResidentChildSpec{
216 Provider: "openai", BaseURL: "https://current.example/v1", InsecureTLS: true, Model: "gpt-5", Workspace: "/repo/child",
217 InheritSkills: &no, InheritProjectContext: &no,
218 })
219 if next.Provider != "openai" || next.BaseURL != "https://current.example/v1" || !next.InsecureTLS || next.Model != "gpt-5" || next.CWD != "/repo/child" || !next.NoSkill || !next.NoContextFiles || next.Orchestrate || !next.ResidentChild {
220 t.Fatalf("resident child args = %#v", next)
221 }
222}
223
224func TestResidentChildArgsDoNotForwardCLIKeyAcrossProviders(t *testing.T) {
225 parent := Args{Provider: "openai", APIKey: "parent-key"}
226 child := subagents.ResidentChildSpec{Provider: "anthropic", Model: "claude"}
227 if next := residentChildArgs(parent, "openai", child); next.APIKey != "" {
228 t.Fatalf("cross-provider child inherited API key %q", next.APIKey)
229 }
230 child.Provider = "openai"
231 if next := residentChildArgs(parent, "openai", child); next.APIKey != "parent-key" {
232 t.Fatalf("same-provider child API key = %q, want parent key", next.APIKey)
233 }
234}