diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index 00b1e7db3..61e7ac0e8 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -134,10 +134,11 @@ func (g *outputBlocker) Check(_ context.Context, message llm.Message) (*agent.Gu type recordingHook struct { agent.NoOpHooks - runStarted bool - runEnded bool - toolNames []string - handoffs []string + runStarted bool + runEnded bool + toolStartNames []string + toolNames []string + handoffs []string } func (h *recordingHook) OnRunStart(_ context.Context, _ *agent.Agent, _ []llm.Message) { @@ -148,6 +149,10 @@ func (h *recordingHook) OnRunEnd(_ context.Context, _ *agent.Agent, _ *agent.Res h.runEnded = true } +func (h *recordingHook) OnToolStart(_ context.Context, _ *agent.Agent, tool agent.Tool, _ string) { + h.toolStartNames = append(h.toolStartNames, tool.Name()) +} + func (h *recordingHook) OnToolEnd(_ context.Context, _ *agent.Agent, tool agent.Tool, _ agent.ToolResult, _ error) { h.toolNames = append(h.toolNames, tool.Name()) } diff --git a/pkg/agent/agent_tool_test.go b/pkg/agent/agent_tool_test.go index bdcf2447b..09440df1e 100644 --- a/pkg/agent/agent_tool_test.go +++ b/pkg/agent/agent_tool_test.go @@ -791,6 +791,70 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { assert.Equal(t, "agent_a", result.LastAgent.Name()) }, ) + + t.Run( + "nested interruption emits paired OnToolStart and OnToolEnd on outer agent", + func(t *testing.T) { + t.Parallel() + + deleteTool, err := agent.FunctionTool[struct{}]( + "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{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "inner_tc1", + Function: llm.FunctionCall{Name: "delete_file", Arguments: `{}`}, + }), + }, + } + + 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( + context.Background(), + []llm.Message{userMessage("Delete the file")}, + ) + + var interrupted *agent.InterruptedError + require.ErrorAs(t, err, &interrupted) + require.Len(t, hook.toolStartNames, 1) + assert.Equal(t, hook.toolStartNames, hook.toolNames, "every OnToolStart must have a matching OnToolEnd") + }, + ) } func TestAgentTool_Execute_DepthLimit(t *testing.T) { diff --git a/pkg/agent/run.go b/pkg/agent/run.go index 41c874415..64bd2f00d 100644 --- a/pkg/agent/run.go +++ b/pkg/agent/run.go @@ -818,6 +818,11 @@ func executeSingleTool( if _, ok := errors.AsType[*InterruptedError](err); ok { toolSpan.SetAttributes(attribute.Bool("tool.interrupted", true)) toolSpan.End() + + onEvent(ctx, StreamEvent{Type: StreamEventToolEnd, Agent: agent, Tool: tool}) + emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}, err) }) + emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}) }) + return ToolResult{}, err } @@ -825,7 +830,9 @@ func executeSingleTool( toolSpan.SetStatus(codes.Error, err.Error()) toolSpan.End() + onEvent(ctx, StreamEvent{Type: StreamEventToolEnd, Agent: agent, Tool: tool, Err: err}) emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, result, err) }) + emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, result) }) logger.ErrorCtx( ctx,