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