8e0d9b5 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix: agent tool not working when switching models

8e0d9b5kujtimiihoxha1 year ago

internal/app/app.go−1

compare with main

unifiedsplit

@@ -278,7 +278,6 @@ func (app *App) InitCoderAgent() error {

278 }

279 var err error

280 app.CoderAgent, err = agent.NewAgent(

281− app.globalCtx,

282 coderAgentCfg,

281 coderAgentCfg,

283 app.Permissions,

282 app.Permissions,

284 app.Sessions,

283 app.Sessions,

internal/llm/agent/agent.go+34−18

@@ -72,6 +72,8 @@ type agent struct {

72 mcpTools []McpTool

74 tools *csync.LazySlice[tools.BaseTool]

75+ // We need this to be able to update it when model changes

76+ agentToolFn func() (tools.BaseTool, error)

76 provider provider.Provider

78 provider provider.Provider

77 providerID string

79 providerID string

@@ -91,7 +93,6 @@ var agentPromptMap = map[string]prompt.PromptID{

91}

93}

93func NewAgent(

95func NewAgent(

94− ctx context.Context,

95 agentCfg config.Agent,

96 agentCfg config.Agent,

96 // These services are needed in the tools

97 // These services are needed in the tools

97 permissions permission.Service,

98 permissions permission.Service,

@@ -102,18 +103,19 @@ func NewAgent(

102) (Service, error) {

103) (Service, error) {

103 cfg := config.Get()

104 cfg := config.Get()

104

105

105− var agentTool tools.BaseTool

106+ var agentToolFn func() (tools.BaseTool, error)

106 if agentCfg.ID == "coder" {

107 if agentCfg.ID == "coder" {

108+ agentToolFn = func() (tools.BaseTool, error) {

107− taskAgentCfg := config.Get().Agents["task"]

109+ taskAgentCfg := config.Get().Agents["task"]

108− if taskAgentCfg.ID == "" {

110+ if taskAgentCfg.ID == "" {

109− return nil, fmt.Errorf("task agent not found in config")

111+ return nil, fmt.Errorf("task agent not found in config")

110− }

112+ }

111− taskAgent, err := NewAgent(ctx, taskAgentCfg, permissions, sessions, messages, history, lspClients)

113+ taskAgent, err := NewAgent(taskAgentCfg, permissions, sessions, messages, history, lspClients)

112− if err != nil {

114+ if err != nil {

113− return nil, fmt.Errorf("failed to create task agent: %w", err)

115+ return nil, fmt.Errorf("failed to create task agent: %w", err)

116+ }

117+ return NewAgentTool(taskAgent, sessions, messages), nil

114 }

118 }

115−

116− agentTool = NewAgentTool(taskAgent, sessions, messages)

117 }

119 }

118

120

119 providerCfg := config.Get().GetProviderForModel(agentCfg.Model)

121 providerCfg := config.Get().GetProviderForModel(agentCfg.Model)

@@ -195,7 +197,7 @@ func NewAgent(

195 }

197 }

196

198

197 mcpToolsOnce.Do(func() {

199 mcpToolsOnce.Do(func() {

198− mcpTools = doGetMCPTools(ctx, permissions, cfg)

200+ mcpTools = doGetMCPTools(permissions, cfg)

199 })

201 })

200 allTools = append(allTools, mcpTools...)

202 allTools = append(allTools, mcpTools...)

201

203

@@ -203,10 +205,6 @@ func NewAgent(

203 allTools = append(allTools, tools.NewDiagnosticsTool(lspClients))

205 allTools = append(allTools, tools.NewDiagnosticsTool(lspClients))

204 }

206 }

205

207

206− if agentTool != nil {

207− allTools = append(allTools, agentTool)

208− }

209−

210 if agentCfg.AllowedTools == nil {

208 if agentCfg.AllowedTools == nil {

211 return allTools

209 return allTools

212 }

210 }

@@ -230,6 +228,7 @@ func NewAgent(

230 titleProvider: titleProvider,

228 titleProvider: titleProvider,

231 summarizeProvider: summarizeProvider,

229 summarizeProvider: summarizeProvider,

232 summarizeProviderID: string(providerCfg.ID),

230 summarizeProviderID: string(providerCfg.ID),

231+ agentToolFn: agentToolFn,

233 activeRequests: csync.NewMap[string, context.CancelFunc](),

232 activeRequests: csync.NewMap[string, context.CancelFunc](),

234 tools: csync.NewLazySlice(toolFn),

233 tools: csync.NewLazySlice(toolFn),

235 promptQueue: csync.NewMap[string, []string](),

234 promptQueue: csync.NewMap[string, []string](),

@@ -500,6 +499,18 @@ func (a *agent) createUserMessage(ctx context.Context, sessionID, content string

500 })

499 })

501}

500}

502

501

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+

503func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {

514func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {

504 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)

515 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)

505

516

@@ -514,8 +525,12 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg

514 return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)

525 return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)

515 }

526 }

516

527

528+ allTools, toolsErr := a.getAllTools()

529+ if toolsErr != nil {

530+ return assistantMsg, nil, toolsErr

531+ }

517 // Now collect tools (which may block on MCP initialization)

532 // 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)

519

534

520 // Add the session and message ID into the context if needed by tools.

535 // Add the session and message ID into the context if needed by tools.

521 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)

536 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)

@@ -554,7 +569,8 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg

554 default:

569 default:

555 // Continue processing

570 // Continue processing

556 var tool tools.BaseTool

571 var tool tools.BaseTool

557− for availableTool := range a.tools.Seq() {

572+ allTools, _ := a.getAllTools()

573+ for _, availableTool := range allTools {

558 if availableTool.Info().Name == toolCall.Name {

574 if availableTool.Info().Name == toolCall.Name {

559 tool = availableTool

575 tool = availableTool

560 break

576 break

internal/llm/agent/mcp-tools.go+2−2

@@ -275,7 +275,7 @@ var mcpInitRequest = mcp.InitializeRequest{

275 },

276}

277

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 {

279 var wg sync.WaitGroup

280 result := csync.NewSlice[tools.BaseTool]()

281

@@ -309,7 +309,7 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con

309 }

310 }()

311

312− ctx, cancel := context.WithTimeout(ctx, mcpTimeout(m))

312+ ctx, cancel := context.WithTimeout(context.Background(), mcpTimeout(m))

313 defer cancel()

314 c, err := createAndInitializeClient(ctx, name, m)

315 if err != nil {