1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "net/http"
9 "net/http/httptest"
10 "strings"
11 "sync/atomic"
12 "testing"
13 "time"
14
15 "github.com/bnema/zut/packages/agent/subagents"
16 "github.com/bnema/zut/packages/core"
17 "github.com/bnema/zut/packages/provider"
18)
19
20func TestResidentChildRegistryUsesExactToolListAndForbidsDelegation(t *testing.T) {
21 catalogue := core.Registry{
22 "read": nil,
23 "bash": nil,
24 "subagent_spawn": nil,
25 "update_goal": nil,
26 }
27 registry, err := residentChildRegistry(catalogue, []string{"read"})
28 if err != nil {
29 t.Fatal(err)
30 }
31 if len(registry) != 1 {
32 t.Fatalf("registry = %#v", registry)
33 }
34 for _, name := range []string{"subagent_spawn", "update_goal", "missing"} {
35 _, err := residentChildRegistry(catalogue, []string{name})
36 if err == nil || !strings.Contains(err.Error(), name) {
37 t.Fatalf("residentChildRegistry(%q) error = %v", name, err)
38 }
39 }
40}
41
42func TestConfigureResidentBudgetOnlyStopsAtExhaustion(t *testing.T) {
43 agent := core.NewAgent(nil, "model", "system", core.Registry{"read": nil})
44 budget := subagents.NewRolloutBudget(1_000, provider.Usage{})
45 configureResidentBudget(agent, budget)
46
47 budget.Observe(provider.Usage{InputTokens: 900})
48 allowed, reason, contextText := agent.BeforeTurnContext(t.Context(), 1)
49 if !allowed || reason != "" || contextText != "" {
50 t.Fatalf("near-limit context = allowed=%t reason=%q context=%q", allowed, reason, contextText)
51 }
52 if len(agent.ToolsSnapshot()) != 1 || agent.MaxTokens != 0 {
53 t.Fatalf("near-limit agent = tools=%#v max_tokens=%d", agent.ToolsSnapshot(), agent.MaxTokens)
54 }
55
56 budget.Observe(provider.Usage{InputTokens: 1_000})
57 allowed, reason, _ = agent.BeforeTurnContext(t.Context(), 2)
58 if allowed || reason != subagents.ErrBudgetExceeded.Error() {
59 t.Fatalf("exhausted context = allowed=%t reason=%q", allowed, reason)
60 }
61}
62
63func TestResidentChildRunnerKeepsToolsAtNinetyPercent(t *testing.T) {
64 var requests atomic.Int32
65 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
66 requestNumber := requests.Add(1)
67 var body map[string]any
68 if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
69 t.Errorf("decode request: %v", err)
70 }
71 w.Header().Set("content-type", "text/event-stream")
72 if requestNumber == 1 {
73 if got := requestMaxTokens(t, body); got != 16_384 {
74 t.Errorf("initial max_tokens = %d, want 16384", got)
75 }
76 writeOpenAIChunk(t, w, map[string]any{
77 "choices": []any{map[string]any{
78 "index": 0,
79 "delta": map[string]any{"tool_calls": []any{map[string]any{
80 "index": 0, "id": "call-1", "type": "function",
81 "function": map[string]string{"name": "read", "arguments": `{"path":"x"}`},
82 }}},
83 "finish_reason": "tool_calls",
84 }},
85 "usage": map[string]int{"prompt_tokens": 900, "completion_tokens": 0},
86 })
87 return
88 }
89 if tools, ok := body["tools"].([]any); !ok || len(tools) == 0 {
90 t.Errorf("near-limit request tools = %#v, want registered tools", body["tools"])
91 }
92 if got := requestMaxTokens(t, body); got <= 100 {
93 t.Errorf("near-limit max_tokens = %d, want model output limit", got)
94 }
95 writeOpenAIChunk(t, w, map[string]any{
96 "choices": []any{map[string]any{"index": 0, "delta": map[string]string{"content": "final result"}, "finish_reason": "stop"}},
97 "usage": map[string]int{"prompt_tokens": 0, "completion_tokens": 50},
98 })
99 }))
100 defer server.Close()
101
102 runner, journal := newBudgetedResidentTestRunner(t, server.URL, 1_000)
103 if err := runner(t.Context(), "review"); err != nil {
104 t.Fatal(err)
105 }
106 if got := requests.Load(); got != 2 {
107 t.Fatalf("provider requests = %d, want 2", got)
108 }
109 if summary := latestResidentAssistantText(t, journal); summary != "final result" {
110 t.Fatalf("latest summary = %q", summary)
111 }
112}
113
114func TestResidentChildRunnerTerminatesAtBudgetLimit(t *testing.T) {
115 var requests atomic.Int32
116 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
117 requests.Add(1)
118 var body map[string]any
119 if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
120 t.Errorf("decode request: %v", err)
121 }
122 if got := requestMaxTokens(t, body); got != 16_384 {
123 t.Errorf("initial max_tokens = %d, want 16384", got)
124 }
125 w.Header().Set("content-type", "text/event-stream")
126 writeOpenAIChunk(t, w, map[string]any{
127 "choices": []any{map[string]any{"index": 0, "delta": map[string]string{"content": "best available result"}, "finish_reason": "stop"}},
128 "usage": map[string]int{"prompt_tokens": 1_200, "completion_tokens": 0},
129 })
130 }))
131 defer server.Close()
132
133 runner, journal := newBudgetedResidentTestRunner(t, server.URL, 1_000)
134 if err := runner(t.Context(), "review"); err != nil {
135 t.Fatalf("runner error = %v, want terminal result", err)
136 }
137 if got := requests.Load(); got != 1 {
138 t.Fatalf("provider requests = %d, want 1", got)
139 }
140 if summary := latestResidentAssistantText(t, journal); summary != "best available result" {
141 t.Fatalf("latest summary = %q", summary)
142 }
143 if err := runner(t.Context(), "continue"); !errors.Is(err, subagents.ErrBudgetExceeded) {
144 t.Fatalf("exhausted follow-up error = %v", err)
145 }
146 if got := requests.Load(); got != 1 {
147 t.Fatalf("exhausted follow-up made %d provider requests, want 1 total", got)
148 }
149}
150
151func TestResidentChildRunnerReportsBudgetStopAfterToolWork(t *testing.T) {
152 var requests atomic.Int32
153 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
154 requests.Add(1)
155 w.Header().Set("content-type", "text/event-stream")
156 writeOpenAIChunk(t, w, map[string]any{
157 "choices": []any{map[string]any{
158 "index": 0,
159 "delta": map[string]any{"tool_calls": []any{map[string]any{
160 "index": 0, "id": "call-budget", "type": "function",
161 "function": map[string]string{"name": "read", "arguments": `{"path":"x"}`},
162 }}},
163 "finish_reason": "tool_calls",
164 }},
165 "usage": map[string]int{"prompt_tokens": 1200, "completion_tokens": 0},
166 })
167 }))
168 defer server.Close()
169 runner, _ := newBudgetedResidentTestRunner(t, server.URL, 1000)
170 if err := runner(t.Context(), "review"); !errors.Is(err, subagents.ErrBudgetExceeded) {
171 t.Fatalf("budget-stopped tool turn = %v, want ErrBudgetExceeded", err)
172 }
173 if got := requests.Load(); got != 1 {
174 t.Fatalf("requests = %d, want 1", got)
175 }
176}
177
178func TestResidentChildRunnerResumesWithSameSizeAllowance(t *testing.T) {
179 var requests atomic.Int32
180 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
181 request := requests.Add(1)
182 w.Header().Set("content-type", "text/event-stream")
183 if request == 1 {
184 writeOpenAIChunk(t, w, map[string]any{
185 "choices": []any{map[string]any{"index": 0, "delta": map[string]string{"content": "first result"}, "finish_reason": "stop"}},
186 "usage": map[string]int{"prompt_tokens": 1200, "completion_tokens": 0},
187 })
188 return
189 }
190 var body map[string]any
191 if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
192 t.Error(err)
193 }
194 encoded, _ := json.Marshal(body["messages"])
195 if !strings.Contains(string(encoded), "first result") {
196 t.Errorf("resume lost prior messages: %s", encoded)
197 }
198 writeOpenAIChunk(t, w, map[string]any{
199 "choices": []any{map[string]any{"index": 0, "delta": map[string]string{"content": "recovered result"}, "finish_reason": "stop"}},
200 "usage": map[string]int{"prompt_tokens": 20, "completion_tokens": 10},
201 })
202 }))
203 defer server.Close()
204 cwd := t.TempDir()
205 manager := subagents.NewResidentManager(t.TempDir(), func(spec subagents.ResidentChildSpec, journal *subagents.ResidentJournal) (subagents.ResidentTurnRunner, error) {
206 return newResidentChildRunner(Args{Provider: "openai", Model: "gpt-4o", APIKey: "synthetic", BaseURL: server.URL, CWD: cwd, NoContextFiles: true, NoSkill: true}, spec, journal)
207 })
208 t.Cleanup(func() { _ = manager.Close(context.Background()) })
209 completed := make(chan subagents.ResidentCompletion, 2)
210 manager.SetCompletionObserver(func(c subagents.ResidentCompletion) { completed <- c })
211 spec := subagents.ResidentChildSpec{ID: "resume-budget", SessionID: "session", Provider: "openai", Model: "gpt-4o", BudgetLimit: 1000, BudgetSource: "model_context"}
212 if _, err := manager.Spawn(t.Context(), spec, "review"); err != nil {
213 t.Fatal(err)
214 }
215 for turn := 0; turn < 2; turn++ {
216 select {
217 case completion := <-completed:
218 if completion.Err != nil {
219 t.Fatal(completion.Err)
220 }
221 case <-time.After(5 * time.Second):
222 t.Fatal("no completion")
223 }
224 if turn == 0 {
225 if err := manager.Resume(t.Context(), spec.ID, "follow up using previous results"); err != nil {
226 t.Fatal(err)
227 }
228 }
229 }
230 snapshot, _ := manager.SnapshotFor(spec.ID)
231 if requests.Load() != 2 || snapshot.Budget.Used != 30 || snapshot.Budget.Limit != 1000 || snapshot.Usage.InputTokens != 1220 {
232 t.Fatalf("requests=%d, snapshot=%#v", requests.Load(), snapshot)
233 }
234}
235
236func newBudgetedResidentTestRunner(t *testing.T, baseURL string, limit int64) (subagents.ResidentTurnRunner, *subagents.ResidentJournal) {
237 t.Helper()
238 spec := subagents.ResidentChildSpec{
239 ID: "budget-child", SessionID: "budget-session", Provider: "openai", Model: "gpt-4o",
240 Tools: []string{"read"}, BudgetLimit: limit, BudgetSource: "model_context",
241 }
242 journal, err := subagents.OpenResidentJournal(t.TempDir(), spec.ID)
243 if err != nil {
244 t.Fatal(err)
245 }
246 t.Cleanup(func() { _ = journal.Close() })
247 if err := journal.Accept(spec, "review"); err != nil {
248 t.Fatal(err)
249 }
250 runner, err := newResidentChildRunner(Args{
251 Provider: "openai", Model: "gpt-4o", APIKey: "synthetic", BaseURL: baseURL,
252 CWD: t.TempDir(), NoContextFiles: true, NoSkill: true,
253 }, spec, journal)
254 if err != nil {
255 t.Fatal(err)
256 }
257 return runner, journal
258}
259
260func requestMaxTokens(t *testing.T, body map[string]any) int {
261 t.Helper()
262 value, ok := body["max_tokens"].(float64)
263 if !ok {
264 t.Errorf("request max_tokens = %#v", body["max_tokens"])
265 return 0
266 }
267 return int(value)
268}
269
270func latestResidentAssistantText(t *testing.T, journal *subagents.ResidentJournal) string {
271 t.Helper()
272 messages, err := subagents.ReadResidentTranscriptMessages(journal.Dir())
273 if err != nil {
274 t.Fatal(err)
275 }
276 for index := len(messages) - 1; index >= 0; index-- {
277 if messages[index].Role != provider.RoleAssistant {
278 continue
279 }
280 var text strings.Builder
281 for _, content := range messages[index].Content {
282 if block, ok := content.(provider.TextBlock); ok {
283 text.WriteString(block.Text)
284 }
285 }
286 return text.String()
287 }
288 return ""
289}
290
291func writeOpenAIChunk(t *testing.T, w http.ResponseWriter, chunk map[string]any) {
292 t.Helper()
293 data, err := json.Marshal(chunk)
294 if err != nil {
295 t.Fatal(err)
296 }
297 _, _ = fmt.Fprintf(w, "data: %s\n\ndata: [DONE]\n\n", data)
298}
299
300func TestResidentChildArgsPreserveDurableProfileInheritance(t *testing.T) {
301 no := false
302 next := residentChildArgs(Args{Orchestrate: true, NoSkill: false, NoContextFiles: false, BaseURL: "https://stale.example/v1", InsecureTLS: false}, "openai", subagents.ResidentChildSpec{
303 Provider: "openai", BaseURL: "https://current.example/v1", InsecureTLS: true, Model: "gpt-5", Workspace: "/repo/child",
304 InheritSkills: &no, InheritProjectContext: &no,
305 })
306 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 {
307 t.Fatalf("resident child args = %#v", next)
308 }
309}
310
311func TestResidentChildArgsDoNotForwardCLIKeyAcrossProviders(t *testing.T) {
312 parent := Args{Provider: "openai", APIKey: "parent-key"}
313 child := subagents.ResidentChildSpec{Provider: "anthropic", Model: "claude"}
314 if next := residentChildArgs(parent, "openai", child); next.APIKey != "" {
315 t.Fatalf("cross-provider child inherited API key %q", next.APIKey)
316 }
317 child.Provider = "openai"
318 if next := residentChildArgs(parent, "openai", child); next.APIKey != "parent-key" {
319 t.Fatalf("same-provider child API key = %q, want parent key", next.APIKey)
320 }
321}