Style

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/agent/agent_tool_test.go

@@ -394,31 +394,33 @@ func TestAgentTool_Execute(t *testing.T) {
 	)
 }
 
-func TestAgentTool_Execute_NestedApproval(t *testing.T) {
+type nestedApprovalFixture struct {
-	t.Parallel()
+	innerProvider *mockProvider
+	outerProvider *mockProvider
+	outerAgent    *agent.Agent
+}
 
-	t.Run(
+func noopDeleteFile(_ context.Context, _ struct{}) (agent.ToolResult, error) {
-		"nested agent approval surfaces as InterruptedError",
+	return agent.ToolResult{Content: "file deleted"}, nil
-		func(t *testing.T) {
+}
-			t.Parallel()
+
+func newNestedApprovalFixture(
+	t *testing.T,
+	deleteFunc func(context.Context, struct{}) (agent.ToolResult, error),
+	innerResponses []*llm.ChatCompletionResponse,
+	outerResponses []*llm.ChatCompletionResponse,
+	outerOpts ...agent.Option,
+) nestedApprovalFixture {
+	t.Helper()
 
 	deleteTool, err := agent.FunctionTool[struct{}](
 		"delete_file",
 		"Delete a file",
-				func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
+		deleteFunc,
-					return agent.ToolResult{Content: "file deleted"}, nil
-				},
 	)
 	require.NoError(t, err)
 
-			innerProvider := &mockProvider{
+	innerProvider := &mockProvider{responses: innerResponses}
-				responses: []*llm.ChatCompletionResponse{
-					toolCallResponse(llm.ToolCall{
-						ID:       "inner_tc1",
-						Function: llm.FunctionCall{Name: "delete_file", Arguments: `{}`},
-					}),
-				},
-			}
 
 	innerAgent := agent.New(
 		"file_manager",
@@ -430,23 +432,55 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
 		}),
 	)
 
-			outerProvider := &mockProvider{
+	outerProvider := &mockProvider{responses: outerResponses}
-				responses: []*llm.ChatCompletionResponse{
+
+	opts := append(
+		[]agent.Option{
+			agent.WithModel("test-model"),
+			agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
+		},
+		outerOpts...,
+	)
+
+	outerAgent := agent.New(
+		"assistant",
+		newTestClient(outerProvider),
+		opts...,
+	)
+
+	return nestedApprovalFixture{
+		innerProvider: innerProvider,
+		outerProvider: outerProvider,
+		outerAgent:    outerAgent,
+	}
+}
+
+func TestAgentTool_Execute_NestedApproval(t *testing.T) {
+	t.Parallel()
+
+	t.Run(
+		"nested agent approval surfaces as InterruptedError",
+		func(t *testing.T) {
+			t.Parallel()
+
+			f := newNestedApprovalFixture(
+				t,
+				noopDeleteFile,
+				[]*llm.ChatCompletionResponse{
+					toolCallResponse(llm.ToolCall{
+						ID:       "inner_tc1",
+						Function: llm.FunctionCall{Name: "delete_file", Arguments: `{}`},
+					}),
+				},
+				[]*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "outer_tc1",
 						Function: llm.FunctionCall{Name: "file_expert", Arguments: `{"input":"delete the file"}`},
 					}),
 				},
-			}
-
-			outerAgent := agent.New(
-				"assistant",
-				newTestClient(outerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
 			)
 
-			_, err = outerAgent.Run(
+			_, err := f.outerAgent.Run(
 				context.Background(),
 				[]llm.Message{userMessage("Delete the file")},
 			)
@@ -467,54 +501,29 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
 
 			var toolExecuted bool
 
-			deleteTool, err := agent.FunctionTool[struct{}](
+			f := newNestedApprovalFixture(
-				"delete_file",
+				t,
-				"Delete a file",
 				func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
 					toolExecuted = true
 					return agent.ToolResult{Content: "file deleted"}, nil
 				},
-			)
+				[]*llm.ChatCompletionResponse{
-			require.NoError(t, err)
-
-			innerProvider := &mockProvider{
-				responses: []*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "inner_tc1",
 						Function: llm.FunctionCall{Name: "delete_file", Arguments: `{}`},
 					}),
 					stopResponse("File has been deleted."),
 				},
-			}
+				[]*llm.ChatCompletionResponse{
-
-			innerAgent := agent.New(
-				"file_manager",
-				newTestClient(innerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(deleteTool),
-				agent.WithApproval(agent.ApprovalConfig{
-					ToolNames: []string{"delete_file"},
-				}),
-			)
-
-			outerProvider := &mockProvider{
-				responses: []*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "outer_tc1",
 						Function: llm.FunctionCall{Name: "file_expert", Arguments: `{"input":"delete the file"}`},
 					}),
 					stopResponse("Done, the file has been deleted."),
 				},
-			}
-
-			outerAgent := agent.New(
-				"assistant",
-				newTestClient(outerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
 			)
 
-			_, err = outerAgent.Run(
+			_, err := f.outerAgent.Run(
 				context.Background(),
 				[]llm.Message{userMessage("Delete the file")},
 			)
@@ -545,54 +554,29 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
 		func(t *testing.T) {
 			t.Parallel()
 
-			deleteTool, err := agent.FunctionTool[struct{}](
+			f := newNestedApprovalFixture(
-				"delete_file",
+				t,
-				"Delete a file",
 				func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
 					t.Fatal("tool should not be called")
 					return agent.ToolResult{}, nil
 				},
-			)
+				[]*llm.ChatCompletionResponse{
-			require.NoError(t, err)
-
-			innerProvider := &mockProvider{
-				responses: []*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "inner_tc1",
 						Function: llm.FunctionCall{Name: "delete_file", Arguments: `{}`},
 					}),
 					stopResponse("OK, I won't delete the file."),
 				},
-			}
+				[]*llm.ChatCompletionResponse{
-
-			innerAgent := agent.New(
-				"file_manager",
-				newTestClient(innerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(deleteTool),
-				agent.WithApproval(agent.ApprovalConfig{
-					ToolNames: []string{"delete_file"},
-				}),
-			)
-
-			outerProvider := &mockProvider{
-				responses: []*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "outer_tc1",
 						Function: llm.FunctionCall{Name: "file_expert", Arguments: `{"input":"delete the file"}`},
 					}),
 					stopResponse("The file manager declined."),
 				},
-			}
-
-			outerAgent := agent.New(
-				"assistant",
-				newTestClient(outerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
 			)
 
-			_, err = outerAgent.Run(
+			_, err := f.outerAgent.Run(
 				context.Background(),
 				[]llm.Message{userMessage("Delete the file")},
 			)
@@ -634,37 +618,17 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
 			)
 			require.NoError(t, err)
 
-			deleteTool, err := agent.FunctionTool[struct{}](
+			f := newNestedApprovalFixture(
-				"delete_file",
+				t,
-				"Delete a file",
+				noopDeleteFile,
-				func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
+				[]*llm.ChatCompletionResponse{
-					return agent.ToolResult{Content: "file deleted"}, nil
-				},
-			)
-			require.NoError(t, err)
-
-			innerProvider := &mockProvider{
-				responses: []*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "inner_tc1",
 						Function: llm.FunctionCall{Name: "delete_file", Arguments: `{}`},
 					}),
 					stopResponse("File has been deleted."),
 				},
-			}
+				[]*llm.ChatCompletionResponse{
-
-			innerAgent := agent.New(
-				"file_manager",
-				newTestClient(innerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(deleteTool),
-				agent.WithApproval(agent.ApprovalConfig{
-					ToolNames: []string{"delete_file"},
-				}),
-			)
-
-			outerProvider := &mockProvider{
-				responses: []*llm.ChatCompletionResponse{
 					toolCallResponse(
 						llm.ToolCall{
 							ID:       "outer_tc1",
@@ -677,19 +641,10 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
 					),
 					stopResponse("Files listed and deleted."),
 				},
-			}
+				agent.WithTools(siblingTool),
-
-			outerAgent := agent.New(
-				"assistant",
-				newTestClient(outerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(
-					siblingTool,
-					innerAgent.AsTool("file_expert", "Manage files"),
-				),
 			)
 
-			_, err = outerAgent.Run(
+			_, err = f.outerAgent.Run(
 				context.Background(),
 				[]llm.Message{userMessage("List and delete files")},
 			)
@@ -818,54 +773,27 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
 		func(t *testing.T) {
 			t.Parallel()
 
-			deleteTool, err := agent.FunctionTool[struct{}](
+			hook := &recordingHook{}
-				"delete_file",
-				"Delete a file",
-				func(_ context.Context, _ struct{}) (agent.ToolResult, error) {
-					return agent.ToolResult{Content: "file deleted"}, nil
-				},
-			)
-			require.NoError(t, err)
 
-			innerProvider := &mockProvider{
+			f := newNestedApprovalFixture(
-				responses: []*llm.ChatCompletionResponse{
+				t,
+				noopDeleteFile,
+				[]*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "inner_tc1",
 						Function: llm.FunctionCall{Name: "delete_file", Arguments: `{}`},
 					}),
 				},
-			}
+				[]*llm.ChatCompletionResponse{
-
-			innerAgent := agent.New(
-				"file_manager",
-				newTestClient(innerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(deleteTool),
-				agent.WithApproval(agent.ApprovalConfig{
-					ToolNames: []string{"delete_file"},
-				}),
-			)
-
-			outerProvider := &mockProvider{
-				responses: []*llm.ChatCompletionResponse{
 					toolCallResponse(llm.ToolCall{
 						ID:       "outer_tc1",
 						Function: llm.FunctionCall{Name: "file_expert", Arguments: `{"input":"delete the file"}`},
 					}),
 				},
-			}
-
-			hook := &recordingHook{}
-
-			outerAgent := agent.New(
-				"assistant",
-				newTestClient(outerProvider),
-				agent.WithModel("test-model"),
-				agent.WithTools(innerAgent.AsTool("file_expert", "Manage files")),
 				agent.WithHooks(hook),
 			)
 
-			_, err = outerAgent.Run(
+			_, err := f.outerAgent.Run(
 				context.Background(),
 				[]llm.Message{userMessage("Delete the file")},
 			)