git.iko.soy/crush
on main help sign in
refactor: use params struct for runSubAgent and add unit tests
42aee6dwanghuaiyu@qiniu.com
internal/agent/agent_tool.go+8−1
compare with main
unifiedsplit
@@ -55,6 +55,13 @@ func (c *coordinator) agentTool(ctx context.Context) (fantasy.AgentTool, error)
55 return fantasy.ToolResponse{}, errors.New("agent message id missing from context")
55 return fantasy.ToolResponse{}, errors.New("agent message id missing from context")
56 }
56 }
57
57
58− return c.runSubAgent(ctx, agent, sessionID, agentMessageID, call.ID, params.Prompt, "New Agent Session")
58+ return c.runSubAgent(ctx, subAgentParams{
59+ Agent: agent,
60+ SessionID: sessionID,
61+ AgentMessageID: agentMessageID,
62+ ToolCallID: call.ID,
63+ Prompt: params.Prompt,
64+ SessionTitle: "New Agent Session",
65+ })
59 }), nil
66 }), nil
60}
67}
internal/agent/agentic_fetch_tool.go+9−10
@@ -184,17 +184,16 @@ func (c *coordinator) agenticFetchTool(_ context.Context, client *http.Client) (
184 Tools: fetchTools,
184 Tools: fetchTools,
185 })
185 })
186
186
187− return c.runSubAgentWithOptions(
187+ return c.runSubAgent(ctx, subAgentParams{
188− ctx,
189− agent,
188+ Agent: agent,
190− validationResult.SessionID,
189+ SessionID: validationResult.SessionID,
191− validationResult.AgentMessageID,
190+ AgentMessageID: validationResult.AgentMessageID,
192− call.ID,
191+ ToolCallID: call.ID,
193− fullPrompt,
192+ Prompt: fullPrompt,
194− "Fetch Analysis",
193+ SessionTitle: "Fetch Analysis",
195− func(sessionID string) {
194+ SessionSetup: func(sessionID string) {
196 c.permissions.AutoApproveSession(sessionID)
195 c.permissions.AutoApproveSession(sessionID)
197 },
196 },
198− )
197+ })
199 }), nil
198 }), nil
200}
199}
internal/agent/coordinator.go+22−28
@@ -941,43 +941,37 @@ func (c *coordinator) refreshApiKeyTemplate(ctx context.Context, providerCfg con
941 return nil
941 return nil
942}
942}
943
943
944+// subAgentParams holds the parameters for running a sub-agent.
945+type subAgentParams struct {
946+ Agent SessionAgent
947+ SessionID string
948+ AgentMessageID string
949+ ToolCallID string
950+ Prompt string
951+ SessionTitle string
952+ // SessionSetup is an optional callback invoked after session creation
953+ // but before agent execution, for custom session configuration.
954+ SessionSetup func(sessionID string)
955+}
956+
944// runSubAgent runs a sub-agent and handles session management and cost accumulation.
957// runSubAgent runs a sub-agent and handles session management and cost accumulation.
945// It creates a sub-session, runs the agent with the given prompt, and propagates
958// It creates a sub-session, runs the agent with the given prompt, and propagates
946// the cost to the parent session.
959// the cost to the parent session.
947−func (c *coordinator) runSubAgent(
960+func (c *coordinator) runSubAgent(ctx context.Context, params subAgentParams) (fantasy.ToolResponse, error) {
948− ctx context.Context,
949− agent SessionAgent,
950− sessionID, agentMessageID, toolCallID string,
951− prompt string,
952− sessionTitle string,
953−) (fantasy.ToolResponse, error) {
954− return c.runSubAgentWithOptions(ctx, agent, sessionID, agentMessageID, toolCallID, prompt, sessionTitle, nil)
955−}
956−
957−// runSubAgentWithOptions runs a sub-agent with additional session configuration options.
958−// The sessionSetup function is called after session creation but before agent execution.
959−func (c *coordinator) runSubAgentWithOptions(
960− ctx context.Context,
961− agent SessionAgent,
962− sessionID, agentMessageID, toolCallID string,
963− prompt string,
964− sessionTitle string,
965− sessionSetup func(sessionID string),
966−) (fantasy.ToolResponse, error) {
967 // Create sub-session
961 // Create sub-session
968− agentToolSessionID := c.sessions.CreateAgentToolSessionID(agentMessageID, toolCallID)
962+ agentToolSessionID := c.sessions.CreateAgentToolSessionID(params.AgentMessageID, params.ToolCallID)
969− session, err := c.sessions.CreateTaskSession(ctx, agentToolSessionID, sessionID, sessionTitle)
963+ session, err := c.sessions.CreateTaskSession(ctx, agentToolSessionID, params.SessionID, params.SessionTitle)
970 if err != nil {
964 if err != nil {
971 return fantasy.ToolResponse{}, fmt.Errorf("create session: %w", err)
965 return fantasy.ToolResponse{}, fmt.Errorf("create session: %w", err)
972 }
966 }
973
967
974 // Call session setup function if provided
968 // Call session setup function if provided
975− if sessionSetup != nil {
969+ if params.SessionSetup != nil {
976− sessionSetup(session.ID)
970+ params.SessionSetup(session.ID)
977 }
971 }
978
972
979 // Get model configuration
973 // Get model configuration
980− model := agent.Model()
974+ model := params.Agent.Model()
981 maxTokens := model.CatwalkCfg.DefaultMaxTokens
975 maxTokens := model.CatwalkCfg.DefaultMaxTokens
982 if model.ModelCfg.MaxTokens != 0 {
976 if model.ModelCfg.MaxTokens != 0 {
983 maxTokens = model.ModelCfg.MaxTokens
977 maxTokens = model.ModelCfg.MaxTokens
@@ -989,9 +983,9 @@ func (c *coordinator) runSubAgentWithOptions(
989 }
983 }
990
984
991 // Run the agent
985 // Run the agent
992− result, err := agent.Run(ctx, SessionAgentCall{
986+ result, err := params.Agent.Run(ctx, SessionAgentCall{
993 SessionID: session.ID,
987 SessionID: session.ID,
994− Prompt: prompt,
988+ Prompt: params.Prompt,
995 MaxOutputTokens: maxTokens,
989 MaxOutputTokens: maxTokens,
996 ProviderOptions: getProviderOptions(model, providerCfg),
990 ProviderOptions: getProviderOptions(model, providerCfg),
997 Temperature: model.ModelCfg.Temperature,
991 Temperature: model.ModelCfg.Temperature,
@@ -1005,7 +999,7 @@ func (c *coordinator) runSubAgentWithOptions(
1005 }
999 }
1006
1000
1007 // Update parent session cost
1001 // Update parent session cost
1008− if err := c.updateParentSessionCost(ctx, session.ID, sessionID); err != nil {
1002+ if err := c.updateParentSessionCost(ctx, session.ID, params.SessionID); err != nil {
1009 return fantasy.ToolResponse{}, err
1003 return fantasy.ToolResponse{}, err
1010 }
1004 }
1011
1005
internal/agent/coordinator_test.gonew file+385
@@ -0,0 +1,385 @@
1+package agent
2+
3+import (
4+ "context"
5+ "errors"
6+ "testing"
7+
8+ "charm.land/catwalk/pkg/catwalk"
9+ "charm.land/fantasy"
10+ "github.com/charmbracelet/crush/internal/config"
11+ "github.com/stretchr/testify/assert"
12+ "github.com/stretchr/testify/require"
13+)
14+
15+// mockSessionAgent is a minimal mock for the SessionAgent interface.
16+type mockSessionAgent struct {
17+ model Model
18+ runFunc func(ctx context.Context, call SessionAgentCall) (*fantasy.AgentResult, error)
19+ cancelled []string
20+}
21+
22+func (m *mockSessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) {
23+ return m.runFunc(ctx, call)
24+}
25+
26+func (m *mockSessionAgent) Model() Model { return m.model }
27+func (m *mockSessionAgent) SetModels(large, small Model) {}
28+func (m *mockSessionAgent) SetTools(tools []fantasy.AgentTool) {}
29+func (m *mockSessionAgent) SetSystemPrompt(systemPrompt string) {}
30+func (m *mockSessionAgent) Cancel(sessionID string) {
31+ m.cancelled = append(m.cancelled, sessionID)
32+}
33+func (m *mockSessionAgent) CancelAll() {}
34+func (m *mockSessionAgent) IsSessionBusy(sessionID string) bool { return false }
35+func (m *mockSessionAgent) IsBusy() bool { return false }
36+func (m *mockSessionAgent) QueuedPrompts(sessionID string) int { return 0 }
37+func (m *mockSessionAgent) QueuedPromptsList(sessionID string) []string { return nil }
38+func (m *mockSessionAgent) ClearQueue(sessionID string) {}
39+func (m *mockSessionAgent) Summarize(context.Context, string, fantasy.ProviderOptions) error {
40+ return nil
41+}
42+
43+// newTestCoordinator creates a minimal coordinator for unit testing runSubAgent.
44+func newTestCoordinator(t *testing.T, env fakeEnv, providerID string, providerCfg config.ProviderConfig) *coordinator {
45+ cfg, err := config.Init(env.workingDir, "", false)
46+ require.NoError(t, err)
47+ cfg.Providers.Set(providerID, providerCfg)
48+ return &coordinator{
49+ cfg: cfg,
50+ sessions: env.sessions,
51+ }
52+}
53+
54+// newMockAgent creates a mockSessionAgent with the given provider and run function.
55+func newMockAgent(providerID string, maxTokens int64, runFunc func(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)) *mockSessionAgent {
56+ return &mockSessionAgent{
57+ model: Model{
58+ CatwalkCfg: catwalk.Model{
59+ DefaultMaxTokens: maxTokens,
60+ },
61+ ModelCfg: config.SelectedModel{
62+ Provider: providerID,
63+ },
64+ },
65+ runFunc: runFunc,
66+ }
67+}
68+
69+// agentResultWithText creates a minimal AgentResult with the given text response.
70+func agentResultWithText(text string) *fantasy.AgentResult {
71+ return &fantasy.AgentResult{
72+ Response: fantasy.Response{
73+ Content: fantasy.ResponseContent{
74+ fantasy.TextContent{Text: text},
75+ },
76+ },
77+ }
78+}
79+
80+func TestRunSubAgent(t *testing.T) {
81+ const providerID = "test-provider"
82+ providerCfg := config.ProviderConfig{ID: providerID}
83+
84+ t.Run("happy path", func(t *testing.T) {
85+ env := testEnv(t)
86+ coord := newTestCoordinator(t, env, providerID, providerCfg)
87+
88+ parentSession, err := env.sessions.Create(t.Context(), "Parent")
89+ require.NoError(t, err)
90+
91+ agent := newMockAgent(providerID, 4096, func(_ context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) {
92+ assert.Equal(t, "do something", call.Prompt)
93+ assert.Equal(t, int64(4096), call.MaxOutputTokens)
94+ return agentResultWithText("done"), nil
95+ })
96+
97+ resp, err := coord.runSubAgent(t.Context(), subAgentParams{
98+ Agent: agent,
99+ SessionID: parentSession.ID,
100+ AgentMessageID: "msg-1",
101+ ToolCallID: "call-1",
102+ Prompt: "do something",
103+ SessionTitle: "Test Session",
104+ })
105+ require.NoError(t, err)
106+ assert.Equal(t, "done", resp.Content)
107+ assert.False(t, resp.IsError)
108+ })
109+
110+ t.Run("ModelCfg.MaxTokens overrides default", func(t *testing.T) {
111+ env := testEnv(t)
112+ coord := newTestCoordinator(t, env, providerID, providerCfg)
113+
114+ parentSession, err := env.sessions.Create(t.Context(), "Parent")
115+ require.NoError(t, err)
116+
117+ agent := &mockSessionAgent{
118+ model: Model{
119+ CatwalkCfg: catwalk.Model{
120+ DefaultMaxTokens: 4096,
121+ },
122+ ModelCfg: config.SelectedModel{
123+ Provider: providerID,
124+ MaxTokens: 8192,
125+ },
126+ },
127+ runFunc: func(_ context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) {
128+ assert.Equal(t, int64(8192), call.MaxOutputTokens)
129+ return agentResultWithText("ok"), nil
130+ },
131+ }
132+
133+ resp, err := coord.runSubAgent(t.Context(), subAgentParams{
134+ Agent: agent,
135+ SessionID: parentSession.ID,
136+ AgentMessageID: "msg-1",
137+ ToolCallID: "call-1",
138+ Prompt: "test",
139+ SessionTitle: "Test",
140+ })
141+ require.NoError(t, err)
142+ assert.Equal(t, "ok", resp.Content)
143+ })
144+
145+ t.Run("session creation failure with canceled context", func(t *testing.T) {
146+ env := testEnv(t)
147+ coord := newTestCoordinator(t, env, providerID, providerCfg)
148+
149+ parentSession, err := env.sessions.Create(t.Context(), "Parent")
150+ require.NoError(t, err)
151+
152+ agent := newMockAgent(providerID, 4096, nil)
153+
154+ // Use a canceled context to trigger CreateTaskSession failure.
155+ ctx, cancel := context.WithCancel(t.Context())
156+ cancel()
157+
158+ _, err = coord.runSubAgent(ctx, subAgentParams{
159+ Agent: agent,
160+ SessionID: parentSession.ID,
161+ AgentMessageID: "msg-1",
162+ ToolCallID: "call-1",
163+ Prompt: "test",
164+ SessionTitle: "Test",
165+ })
166+ require.Error(t, err)
167+ })
168+
169+ t.Run("provider not configured", func(t *testing.T) {
170+ env := testEnv(t)
171+ coord := newTestCoordinator(t, env, providerID, providerCfg)
172+
173+ parentSession, err := env.sessions.Create(t.Context(), "Parent")
174+ require.NoError(t, err)
175+
176+ // Agent references a provider that doesn't exist in config.
177+ agent := newMockAgent("unknown-provider", 4096, nil)
178+
179+ _, err = coord.runSubAgent(t.Context(), subAgentParams{
180+ Agent: agent,
181+ SessionID: parentSession.ID,
182+ AgentMessageID: "msg-1",
183+ ToolCallID: "call-1",
184+ Prompt: "test",
185+ SessionTitle: "Test",
186+ })
187+ require.Error(t, err)
188+ assert.Contains(t, err.Error(), "model provider not configured")
189+ })
190+
191+ t.Run("agent run error returns error response", func(t *testing.T) {
192+ env := testEnv(t)
193+ coord := newTestCoordinator(t, env, providerID, providerCfg)
194+
195+ parentSession, err := env.sessions.Create(t.Context(), "Parent")
196+ require.NoError(t, err)
197+
198+ agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) {
199+ return nil, errors.New("agent exploded")
200+ })
201+
202+ resp, err := coord.runSubAgent(t.Context(), subAgentParams{
203+ Agent: agent,
204+ SessionID: parentSession.ID,
205+ AgentMessageID: "msg-1",
206+ ToolCallID: "call-1",
207+ Prompt: "test",
208+ SessionTitle: "Test",
209+ })
210+ // runSubAgent returns (errorResponse, nil) when agent.Run fails — not a Go error.
211+ require.NoError(t, err)
212+ assert.True(t, resp.IsError)
213+ assert.Equal(t, "error generating response", resp.Content)
214+ })
215+
216+ t.Run("session setup callback is invoked", func(t *testing.T) {
217+ env := testEnv(t)
218+ coord := newTestCoordinator(t, env, providerID, providerCfg)
219+
220+ parentSession, err := env.sessions.Create(t.Context(), "Parent")
221+ require.NoError(t, err)
222+
223+ var setupCalledWith string
224+ agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) {
225+ return agentResultWithText("ok"), nil
226+ })
227+
228+ _, err = coord.runSubAgent(t.Context(), subAgentParams{
229+ Agent: agent,
230+ SessionID: parentSession.ID,
231+ AgentMessageID: "msg-1",
232+ ToolCallID: "call-1",
233+ Prompt: "test",
234+ SessionTitle: "Test",
235+ SessionSetup: func(sessionID string) {
236+ setupCalledWith = sessionID
237+ },
238+ })
239+ require.NoError(t, err)
240+ assert.NotEmpty(t, setupCalledWith, "SessionSetup should have been called")
241+ })
242+
243+ t.Run("cost propagation to parent session", func(t *testing.T) {
244+ env := testEnv(t)
245+ coord := newTestCoordinator(t, env, providerID, providerCfg)
246+
247+ parentSession, err := env.sessions.Create(t.Context(), "Parent")
248+ require.NoError(t, err)
249+
250+ agent := newMockAgent(providerID, 4096, func(ctx context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) {
251+ // Simulate the agent incurring cost by updating the child session.
252+ childSession, err := env.sessions.Get(ctx, call.SessionID)
253+ if err != nil {
254+ return nil, err
255+ }
256+ childSession.Cost = 0.05
257+ _, err = env.sessions.Save(ctx, childSession)
258+ if err != nil {
259+ return nil, err
260+ }
261+ return agentResultWithText("ok"), nil
262+ })
263+
264+ _, err = coord.runSubAgent(t.Context(), subAgentParams{
265+ Agent: agent,
266+ SessionID: parentSession.ID,
267+ AgentMessageID: "msg-1",
268+ ToolCallID: "call-1",
269+ Prompt: "test",
270+ SessionTitle: "Test",
271+ })
272+ require.NoError(t, err)
273+
274+ updated, err := env.sessions.Get(t.Context(), parentSession.ID)
275+ require.NoError(t, err)
276+ assert.InDelta(t, 0.05, updated.Cost, 1e-9)
277+ })
278+}
279+
280+func TestUpdateParentSessionCost(t *testing.T) {
281+ t.Run("accumulates cost correctly", func(t *testing.T) {
282+ env := testEnv(t)
283+ cfg, err := config.Init(env.workingDir, "", false)
284+ require.NoError(t, err)
285+ coord := &coordinator{cfg: cfg, sessions: env.sessions}
286+
287+ parent, err := env.sessions.Create(t.Context(), "Parent")
288+ require.NoError(t, err)
289+
290+ child, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child")
291+ require.NoError(t, err)
292+
293+ // Set child cost.
294+ child.Cost = 0.10
295+ _, err = env.sessions.Save(t.Context(), child)
296+ require.NoError(t, err)
297+
298+ err = coord.updateParentSessionCost(t.Context(), child.ID, parent.ID)
299+ require.NoError(t, err)
300+
301+ updated, err := env.sessions.Get(t.Context(), parent.ID)
302+ require.NoError(t, err)
303+ assert.InDelta(t, 0.10, updated.Cost, 1e-9)
304+ })
305+
306+ t.Run("accumulates multiple child costs", func(t *testing.T) {
307+ env := testEnv(t)
308+ cfg, err := config.Init(env.workingDir, "", false)
309+ require.NoError(t, err)
310+ coord := &coordinator{cfg: cfg, sessions: env.sessions}
311+
312+ parent, err := env.sessions.Create(t.Context(), "Parent")
313+ require.NoError(t, err)
314+
315+ child1, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child1")
316+ require.NoError(t, err)
317+ child1.Cost = 0.05
318+ _, err = env.sessions.Save(t.Context(), child1)
319+ require.NoError(t, err)
320+
321+ child2, err := env.sessions.CreateTaskSession(t.Context(), "tool-2", parent.ID, "Child2")
322+ require.NoError(t, err)
323+ child2.Cost = 0.03
324+ _, err = env.sessions.Save(t.Context(), child2)
325+ require.NoError(t, err)
326+
327+ err = coord.updateParentSessionCost(t.Context(), child1.ID, parent.ID)
328+ require.NoError(t, err)
329+ err = coord.updateParentSessionCost(t.Context(), child2.ID, parent.ID)
330+ require.NoError(t, err)
331+
332+ updated, err := env.sessions.Get(t.Context(), parent.ID)
333+ require.NoError(t, err)
334+ assert.InDelta(t, 0.08, updated.Cost, 1e-9)
335+ })
336+
337+ t.Run("child session not found", func(t *testing.T) {
338+ env := testEnv(t)
339+ cfg, err := config.Init(env.workingDir, "", false)
340+ require.NoError(t, err)
341+ coord := &coordinator{cfg: cfg, sessions: env.sessions}
342+
343+ parent, err := env.sessions.Create(t.Context(), "Parent")
344+ require.NoError(t, err)
345+
346+ err = coord.updateParentSessionCost(t.Context(), "non-existent", parent.ID)
347+ require.Error(t, err)
348+ assert.Contains(t, err.Error(), "get child session")
349+ })
350+
351+ t.Run("parent session not found", func(t *testing.T) {
352+ env := testEnv(t)
353+ cfg, err := config.Init(env.workingDir, "", false)
354+ require.NoError(t, err)
355+ coord := &coordinator{cfg: cfg, sessions: env.sessions}
356+
357+ parent, err := env.sessions.Create(t.Context(), "Parent")
358+ require.NoError(t, err)
359+ child, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child")
360+ require.NoError(t, err)
361+
362+ err = coord.updateParentSessionCost(t.Context(), child.ID, "non-existent")
363+ require.Error(t, err)
364+ assert.Contains(t, err.Error(), "get parent session")
365+ })
366+
367+ t.Run("zero cost handled correctly", func(t *testing.T) {
368+ env := testEnv(t)
369+ cfg, err := config.Init(env.workingDir, "", false)
370+ require.NoError(t, err)
371+ coord := &coordinator{cfg: cfg, sessions: env.sessions}
372+
373+ parent, err := env.sessions.Create(t.Context(), "Parent")
374+ require.NoError(t, err)
375+ child, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child")
376+ require.NoError(t, err)
377+
378+ err = coord.updateParentSessionCost(t.Context(), child.ID, parent.ID)
379+ require.NoError(t, err)
380+
381+ updated, err := env.sessions.Get(t.Context(), parent.ID)
382+ require.NoError(t, err)
383+ assert.InDelta(t, 0.0, updated.Cost, 1e-9)
384+ })
385+}