git.iko.soy/crush
on main help sign in
refactor: extract common sub-agent execution logic
e0d5454wanghuaiyu@qiniu.com
internal/agent/agent_tool.go+1−46
compare with main
unifiedsplit
@@ -4,7 +4,6 @@ import (
4 "context"
4 "context"
5 _ "embed"
5 _ "embed"
6 "errors"
6 "errors"
7− "fmt"
8
7
9 "charm.land/fantasy"
8 "charm.land/fantasy"
10
9
@@ -56,50 +55,6 @@ func (c *coordinator) agentTool(ctx context.Context) (fantasy.AgentTool, error)
56 return fantasy.ToolResponse{}, errors.New("agent message id missing from context")
55 return fantasy.ToolResponse{}, errors.New("agent message id missing from context")
57 }
56 }
58
57
59− agentToolSessionID := c.sessions.CreateAgentToolSessionID(agentMessageID, call.ID)
58+ return c.runSubAgent(ctx, agent, sessionID, agentMessageID, call.ID, params.Prompt, "New Agent Session")
60− session, err := c.sessions.CreateTaskSession(ctx, agentToolSessionID, sessionID, "New Agent Session")
61− if err != nil {
62− return fantasy.ToolResponse{}, fmt.Errorf("error creating session: %s", err)
63− }
64− model := agent.Model()
65− maxTokens := model.CatwalkCfg.DefaultMaxTokens
66− if model.ModelCfg.MaxTokens != 0 {
67− maxTokens = model.ModelCfg.MaxTokens
68− }
69−
70− providerCfg, ok := c.cfg.Providers.Get(model.ModelCfg.Provider)
71− if !ok {
72− return fantasy.ToolResponse{}, errors.New("model provider not configured")
73− }
74− result, err := agent.Run(ctx, SessionAgentCall{
75− SessionID: session.ID,
76− Prompt: params.Prompt,
77− MaxOutputTokens: maxTokens,
78− ProviderOptions: getProviderOptions(model, providerCfg),
79− Temperature: model.ModelCfg.Temperature,
80− TopP: model.ModelCfg.TopP,
81− TopK: model.ModelCfg.TopK,
82− FrequencyPenalty: model.ModelCfg.FrequencyPenalty,
83− PresencePenalty: model.ModelCfg.PresencePenalty,
84− })
85− if err != nil {
86− return fantasy.NewTextErrorResponse("error generating response"), nil
87− }
88− updatedSession, err := c.sessions.Get(ctx, session.ID)
89− if err != nil {
90− return fantasy.ToolResponse{}, fmt.Errorf("error getting session: %s", err)
91− }
92− parentSession, err := c.sessions.Get(ctx, sessionID)
93− if err != nil {
94− return fantasy.ToolResponse{}, fmt.Errorf("error getting parent session: %s", err)
95− }
96−
97− parentSession.Cost += updatedSession.Cost
98−
99− _, err = c.sessions.Save(ctx, parentSession)
100− if err != nil {
101− return fantasy.ToolResponse{}, fmt.Errorf("error saving parent session: %s", err)
102− }
103− return fantasy.NewTextResponse(result.Response.Content.Text()), nil
104 }), nil
59 }), nil
105}
60}
internal/agent/agentic_fetch_tool.go+12−46
@@ -184,51 +184,17 @@ func (c *coordinator) agenticFetchTool(_ context.Context, client *http.Client) (
184 Tools: fetchTools,
184 Tools: fetchTools,
185 })
185 })
186
186
187+ return c.runSubAgentWithOptions(
188+ ctx,
189+ agent,
190+ validationResult.SessionID,
187− agentToolSessionID := c.sessions.CreateAgentToolSessionID(validationResult.AgentMessageID, call.ID)
191+ validationResult.AgentMessageID,
188− session, err := c.sessions.CreateTaskSession(ctx, agentToolSessionID, validationResult.SessionID, "Fetch Analysis")
192+ call.ID,
189− if err != nil {
193+ fullPrompt,
190− return fantasy.ToolResponse{}, fmt.Errorf("error creating session: %s", err)
194+ "Fetch Analysis",
191− }
195+ func(sessionID string) {
192−
193− c.permissions.AutoApproveSession(session.ID)
196+ c.permissions.AutoApproveSession(sessionID)
194−
197+ },
195− // Use small model for web content analysis (faster and cheaper)
198+ )
196− maxTokens := small.CatwalkCfg.DefaultMaxTokens
197− if small.ModelCfg.MaxTokens != 0 {
198− maxTokens = small.ModelCfg.MaxTokens
199− }
200−
201− result, err := agent.Run(ctx, SessionAgentCall{
202− SessionID: session.ID,
203− Prompt: fullPrompt,
204− MaxOutputTokens: maxTokens,
205− ProviderOptions: getProviderOptions(small, smallProviderCfg),
206− Temperature: small.ModelCfg.Temperature,
207− TopP: small.ModelCfg.TopP,
208− TopK: small.ModelCfg.TopK,
209− FrequencyPenalty: small.ModelCfg.FrequencyPenalty,
210− PresencePenalty: small.ModelCfg.PresencePenalty,
211− })
212− if err != nil {
213− return fantasy.NewTextErrorResponse("error generating response"), nil
214− }
215−
216− updatedSession, err := c.sessions.Get(ctx, session.ID)
217− if err != nil {
218− return fantasy.ToolResponse{}, fmt.Errorf("error getting session: %s", err)
219− }
220− parentSession, err := c.sessions.Get(ctx, validationResult.SessionID)
221− if err != nil {
222− return fantasy.ToolResponse{}, fmt.Errorf("error getting parent session: %s", err)
223− }
224−
225− parentSession.Cost += updatedSession.Cost
226−
227− _, err = c.sessions.Save(ctx, parentSession)
228− if err != nil {
229− return fantasy.ToolResponse{}, fmt.Errorf("error saving parent session: %s", err)
230− }
231−
232− return fantasy.NewTextResponse(result.Response.Content.Text()), nil
233 }), nil
199 }), nil
234}
200}
internal/agent/coordinator.go+92
@@ -940,3 +940,95 @@ func (c *coordinator) refreshApiKeyTemplate(ctx context.Context, providerCfg con
940 }
940 }
941 return nil
941 return nil
942}
942}
943+
944+// 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
946+// the cost to the parent session.
947+func (c *coordinator) runSubAgent(
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
968+ agentToolSessionID := c.sessions.CreateAgentToolSessionID(agentMessageID, toolCallID)
969+ session, err := c.sessions.CreateTaskSession(ctx, agentToolSessionID, sessionID, sessionTitle)
970+ if err != nil {
971+ return fantasy.ToolResponse{}, fmt.Errorf("create session: %w", err)
972+ }
973+
974+ // Call session setup function if provided
975+ if sessionSetup != nil {
976+ sessionSetup(session.ID)
977+ }
978+
979+ // Get model configuration
980+ model := agent.Model()
981+ maxTokens := model.CatwalkCfg.DefaultMaxTokens
982+ if model.ModelCfg.MaxTokens != 0 {
983+ maxTokens = model.ModelCfg.MaxTokens
984+ }
985+
986+ providerCfg, ok := c.cfg.Providers.Get(model.ModelCfg.Provider)
987+ if !ok {
988+ return fantasy.ToolResponse{}, errors.New("model provider not configured")
989+ }
990+
991+ // Run the agent
992+ result, err := agent.Run(ctx, SessionAgentCall{
993+ SessionID: session.ID,
994+ Prompt: prompt,
995+ MaxOutputTokens: maxTokens,
996+ ProviderOptions: getProviderOptions(model, providerCfg),
997+ Temperature: model.ModelCfg.Temperature,
998+ TopP: model.ModelCfg.TopP,
999+ TopK: model.ModelCfg.TopK,
1000+ FrequencyPenalty: model.ModelCfg.FrequencyPenalty,
1001+ PresencePenalty: model.ModelCfg.PresencePenalty,
1002+ })
1003+ if err != nil {
1004+ return fantasy.NewTextErrorResponse("error generating response"), nil
1005+ }
1006+
1007+ // Update parent session cost
1008+ if err := c.updateParentSessionCost(ctx, session.ID, sessionID); err != nil {
1009+ return fantasy.ToolResponse{}, err
1010+ }
1011+
1012+ return fantasy.NewTextResponse(result.Response.Content.Text()), nil
1013+}
1014+
1015+// updateParentSessionCost accumulates the cost from a child session to its parent session.
1016+func (c *coordinator) updateParentSessionCost(ctx context.Context, childSessionID, parentSessionID string) error {
1017+ childSession, err := c.sessions.Get(ctx, childSessionID)
1018+ if err != nil {
1019+ return fmt.Errorf("get child session: %w", err)
1020+ }
1021+
1022+ parentSession, err := c.sessions.Get(ctx, parentSessionID)
1023+ if err != nil {
1024+ return fmt.Errorf("get parent session: %w", err)
1025+ }
1026+
1027+ parentSession.Cost += childSession.Cost
1028+
1029+ if _, err := c.sessions.Save(ctx, parentSession); err != nil {
1030+ return fmt.Errorf("save parent session: %w", err)
1031+ }
1032+
1033+ return nil
1034+}