Merge pull request #515 from dgageot/fix-514

Temporary fix to exclude todo tools from code Mode

Author: dgageot

pkg/codemode/codemode.go

@@ -46,8 +46,17 @@ func (c *codeModeTool) Instructions() string {
 	return ""
 }
 
+func isExcludedTool(tool tools.Tool) bool {
+	// TODO(dga): make this more robust. It really a temporary hack to exclude the todo tools.
+	// See #514
+	return strings.Contains(tool.Name, "_todo")
+}
+
 func (c *codeModeTool) Tools(ctx context.Context) ([]tools.Tool, error) {
-	var functionsDoc []string
+	var (
+		functionsDoc  []string
+		excludedTools []tools.Tool
+	)
 
 	for _, toolset := range c.toolsets {
 		allTools, err := toolset.Tools(ctx)
@@ -56,11 +65,15 @@ func (c *codeModeTool) Tools(ctx context.Context) ([]tools.Tool, error) {
 		}
 
 		for _, tool := range allTools {
-			functionsDoc = append(functionsDoc, toolToJsDoc(tool))
+			if isExcludedTool(tool) {
+				excludedTools = append(excludedTools, tool)
+			} else {
+				functionsDoc = append(functionsDoc, toolToJsDoc(tool))
+			}
 		}
 	}
 
-	return []tools.Tool{{
+	allTools := []tools.Tool{{
 		Name:        "run_tools_with_javascript",
 		Description: prompt + strings.Join(functionsDoc, "\n"),
 		Parameters:  tools.MustSchemaFor[RunToolsWithJavascriptArgs](),
@@ -88,7 +101,11 @@ func (c *codeModeTool) Tools(ctx context.Context) ([]tools.Tool, error) {
 		Annotations: tools.ToolAnnotations{
 			Title: "Run tools with Javascript",
 		},
-	}}, nil
+	}}
+
+	allTools = append(allTools, excludedTools...)
+
+	return allTools, nil
 }
 
 func (c *codeModeTool) Start(ctx context.Context) error {

pkg/teamloader/teamloader.go

@@ -173,13 +173,12 @@ func LoadWithOverrides(ctx context.Context, path string, runtimeConfig config.Ru
 			return nil, fmt.Errorf("failed to get tools: %w", err)
 		}
 
+		if len(a.SubAgents) > 0 {
+			agentTools = append(agentTools, builtin.NewTransferTaskTool())
+		}
+
 		if len(agentTools) > 0 {
-			if agentConfig.CodeModeTools || runtimeConfig.GlobalCodeMode {
-				codemodeTool := codemode.Wrap(agentTools)
-				opts = append(opts, agent.WithToolSets(codemodeTool))
-			} else {
-				opts = append(opts, agent.WithToolSets(agentTools...))
-			}
+			opts = append(opts, agent.WithToolSets(agentTools...))
 		}
 
 		ag := agent.New(name, agentConfig.Instruction, opts...)
@@ -237,10 +236,6 @@ func getModelsForAgent(ctx context.Context, cfg *latest.Config, a *latest.AgentC
 func getToolsForAgent(ctx context.Context, a *latest.AgentConfig, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) ([]tools.ToolSet, error) {
 	var t []tools.ToolSet
 
-	if len(a.SubAgents) > 0 {
-		t = append(t, builtin.NewTransferTaskTool())
-	}
-
 	for i := range a.Toolsets {
 		toolset := a.Toolsets[i]
 
@@ -252,7 +247,16 @@ func getToolsForAgent(ctx context.Context, a *latest.AgentConfig, parentDir stri
 		t = append(t, WithInstructions(tool, a.Instruction))
 	}
 
-	return t, nil
+	if !a.CodeModeTools && !runtimeConfig.GlobalCodeMode {
+		return t, nil
+	}
+
+	// Wrap all tools in a single Code Mode toolset.
+	// This allows the agent to call multiple tools in a single response.
+	// It also allows to combine the results of multiple tools in a single response.
+	return []tools.ToolSet{
+		codemode.Wrap(t),
+	}, nil
 }
 
 func createTool(ctx context.Context, toolset latest.Toolset, a *latest.AgentConfig, parentDir string, envProvider environment.Provider, runtimeConfig config.RuntimeConfig) (tools.ToolSet, error) {