on main help sign in
fix: agent tool not working when switching models
commit8e0d9b5b9aaed5841d44457858716de08405917e authorkujtimiihoxha <kujtimii.h@gmail.com> dateThu 11 Sep 2025, 12:53 (1 year ago) parent49a1e21
3 files changed+36−21
internal/app/app.go−1 internal/llm/agent/agent.go+34−18 internal/llm/agent/mcp-tools.go+2−2
internal/app/app.go−1
@@ -278,7 +278,6 @@ func (app *App) InitCoderAgent() error {
278278 }
279279 var err error
280280 app.CoderAgent, err = agent.NewAgent(
281− app.globalCtx,
282281 coderAgentCfg,
283282 app.Permissions,
284283 app.Sessions,
internal/llm/agent/agent.go+34−18
@@ -72,6 +72,8 @@ type agent struct {
7272 mcpTools []McpTool
7373
7474 tools *csync.LazySlice[tools.BaseTool]
75+ // We need this to be able to update it when model changes
76+ agentToolFn func() (tools.BaseTool, error)
7577
7678 provider provider.Provider
7779 providerID string
@@ -91,7 +93,6 @@ var agentPromptMap = map[string]prompt.PromptID{
9193}
9294
9395func NewAgent(
94− ctx context.Context,
9596 agentCfg config.Agent,
9697 // These services are needed in the tools
9798 permissions permission.Service,
@@ -102,18 +103,19 @@ func NewAgent(
102103) (Service, error) {
103104 cfg := config.Get()
104105
105− var agentTool tools.BaseTool
106+ var agentToolFn func() (tools.BaseTool, error)
106107 if agentCfg.ID == "coder" {
107− taskAgentCfg := config.Get().Agents["task"]
108− if taskAgentCfg.ID == "" {
109− return nil, fmt.Errorf("task agent not found in config")
110− }
111− taskAgent, err := NewAgent(ctx, taskAgentCfg, permissions, sessions, messages, history, lspClients)
112− if err != nil {
113− return nil, fmt.Errorf("failed to create task agent: %w", err)
108+ agentToolFn = func() (tools.BaseTool, error) {
109+ taskAgentCfg := config.Get().Agents["task"]
110+ if taskAgentCfg.ID == "" {
111+ return nil, fmt.Errorf("task agent not found in config")
112+ }
113+ taskAgent, err := NewAgent(taskAgentCfg, permissions, sessions, messages, history, lspClients)
114+ if err != nil {
115+ return nil, fmt.Errorf("failed to create task agent: %w", err)
116+ }
117+ return NewAgentTool(taskAgent, sessions, messages), nil
114118 }
115−
116− agentTool = NewAgentTool(taskAgent, sessions, messages)
117119 }
118120
119121 providerCfg := config.Get().GetProviderForModel(agentCfg.Model)
@@ -195,7 +197,7 @@ func NewAgent(
195197 }
196198
197199 mcpToolsOnce.Do(func() {
198− mcpTools = doGetMCPTools(ctx, permissions, cfg)
200+ mcpTools = doGetMCPTools(permissions, cfg)
199201 })
200202 allTools = append(allTools, mcpTools...)
201203
@@ -203,10 +205,6 @@ func NewAgent(
203205 allTools = append(allTools, tools.NewDiagnosticsTool(lspClients))
204206 }
205207
206− if agentTool != nil {
207− allTools = append(allTools, agentTool)
208− }
209−
210208 if agentCfg.AllowedTools == nil {
211209 return allTools
212210 }
@@ -230,6 +228,7 @@ func NewAgent(
230228 titleProvider: titleProvider,
231229 summarizeProvider: summarizeProvider,
232230 summarizeProviderID: string(providerCfg.ID),
231+ agentToolFn: agentToolFn,
233232 activeRequests: csync.NewMap[string, context.CancelFunc](),
234233 tools: csync.NewLazySlice(toolFn),
235234 promptQueue: csync.NewMap[string, []string](),
@@ -500,6 +499,18 @@ func (a *agent) createUserMessage(ctx context.Context, sessionID, content string
500499 })
501500}
502501
502+func (a *agent) getAllTools() ([]tools.BaseTool, error) {
503+ allTools := slices.Collect(a.tools.Seq())
504+ if a.agentToolFn != nil {
505+ agentTool, agentToolErr := a.agentToolFn()
506+ if agentToolErr != nil {
507+ return nil, agentToolErr
508+ }
509+ allTools = append(allTools, agentTool)
510+ }
511+ return allTools, nil
512+}
513+
503514func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {
504515 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
505516
@@ -514,8 +525,12 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg
514525 return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)
515526 }
516527
528+ allTools, toolsErr := a.getAllTools()
529+ if toolsErr != nil {
530+ return assistantMsg, nil, toolsErr
531+ }
517532 // Now collect tools (which may block on MCP initialization)
518− eventChan := a.provider.StreamResponse(ctx, msgHistory, slices.Collect(a.tools.Seq()))
533+ eventChan := a.provider.StreamResponse(ctx, msgHistory, allTools)
519534
520535 // Add the session and message ID into the context if needed by tools.
521536 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
@@ -554,7 +569,8 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg
554569 default:
555570 // Continue processing
556571 var tool tools.BaseTool
557− for availableTool := range a.tools.Seq() {
572+ allTools, _ := a.getAllTools()
573+ for _, availableTool := range allTools {
558574 if availableTool.Info().Name == toolCall.Name {
559575 tool = availableTool
560576 break
internal/llm/agent/mcp-tools.go+2−2
@@ -275,7 +275,7 @@ var mcpInitRequest = mcp.InitializeRequest{
275275 },
276276}
277277
278−func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {
278+func doGetMCPTools(permissions permission.Service, cfg *config.Config) []tools.BaseTool {
279279 var wg sync.WaitGroup
280280 result := csync.NewSlice[tools.BaseTool]()
281281
@@ -309,7 +309,7 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con
309309 }
310310 }()
311311
312− ctx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
312+ ctx, cancel := context.WithTimeout(context.Background(), mcpTimeout(m))
313313 defer cancel()
314314 c, err := createAndInitializeClient(ctx, name, m)
315315 if err != nil {