chore: revert ctx change

kujtimiihoxha created 3 months ago

Change summary

internal/llm/agent/agent.go     | 5 +++--
internal/llm/agent/mcp-tools.go | 4 ++--
2 files changed, 5 insertions(+), 4 deletions(-)

Detailed changes

internal/llm/agent/agent.go

@@ -93,6 +93,7 @@ var agentPromptMap = map[string]prompt.PromptID{
 }
 
 func NewAgent(
+	ctx context.Context,
 	agentCfg config.Agent,
 	// These services are needed in the tools
 	permissions permission.Service,
@@ -110,7 +111,7 @@ func NewAgent(
 			if taskAgentCfg.ID == "" {
 				return nil, fmt.Errorf("task agent not found in config")
 			}
-			taskAgent, err := NewAgent(taskAgentCfg, permissions, sessions, messages, history, lspClients)
+			taskAgent, err := NewAgent(ctx, taskAgentCfg, permissions, sessions, messages, history, lspClients)
 			if err != nil {
 				return nil, fmt.Errorf("failed to create task agent: %w", err)
 			}
@@ -197,7 +198,7 @@ func NewAgent(
 		}
 
 		mcpToolsOnce.Do(func() {
-			mcpTools = doGetMCPTools(permissions, cfg)
+			mcpTools = doGetMCPTools(ctx, permissions, cfg)
 		})
 		allTools = append(allTools, mcpTools...)
 

internal/llm/agent/mcp-tools.go

@@ -275,7 +275,7 @@ var mcpInitRequest = mcp.InitializeRequest{
 	},
 }
 
-func doGetMCPTools(permissions permission.Service, cfg *config.Config) []tools.BaseTool {
+func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {
 	var wg sync.WaitGroup
 	result := csync.NewSlice[tools.BaseTool]()
 
@@ -309,7 +309,7 @@ func doGetMCPTools(permissions permission.Service, cfg *config.Config) []tools.B
 				}
 			}()
 
-			ctx, cancel := context.WithTimeout(context.Background(), mcpTimeout(m))
+			ctx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
 			defer cancel()
 			c, err := createAndInitializeClient(ctx, name, m)
 			if err != nil {