Emit OnToolEnd hooks for interrupted and failed tools

executeSingleTool had three exit paths but only the success path
emitted all end signals. The interrupted path (nested agent
approval) skipped OnToolEnd and StreamEventToolEnd entirely,
leaving hook consumers with an unpaired OnToolStart. The error
path also missed StreamEventToolEnd and AgentHooks.OnToolEnd.

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2026-03-13 19:53:08 +01:00
parent 5c889a1d47
commit 9d999559f1
3 changed files with 80 additions and 4 deletions

View File

@@ -134,10 +134,11 @@ func (g *outputBlocker) Check(_ context.Context, message llm.Message) (*agent.Gu
type recordingHook struct { type recordingHook struct {
agent.NoOpHooks agent.NoOpHooks
runStarted bool runStarted bool
runEnded bool runEnded bool
toolNames []string toolStartNames []string
handoffs []string toolNames []string
handoffs []string
} }
func (h *recordingHook) OnRunStart(_ context.Context, _ *agent.Agent, _ []llm.Message) { 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 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) { func (h *recordingHook) OnToolEnd(_ context.Context, _ *agent.Agent, tool agent.Tool, _ agent.ToolResult, _ error) {
h.toolNames = append(h.toolNames, tool.Name()) h.toolNames = append(h.toolNames, tool.Name())
} }

View File

@@ -791,6 +791,70 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) {
assert.Equal(t, "agent_a", result.LastAgent.Name()) 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) { func TestAgentTool_Execute_DepthLimit(t *testing.T) {

View File

@@ -818,6 +818,11 @@ func executeSingleTool(
if _, ok := errors.AsType[*InterruptedError](err); ok { if _, ok := errors.AsType[*InterruptedError](err); ok {
toolSpan.SetAttributes(attribute.Bool("tool.interrupted", true)) toolSpan.SetAttributes(attribute.Bool("tool.interrupted", true))
toolSpan.End() 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 return ToolResult{}, err
} }
@@ -825,7 +830,9 @@ func executeSingleTool(
toolSpan.SetStatus(codes.Error, err.Error()) toolSpan.SetStatus(codes.Error, err.Error())
toolSpan.End() 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) }) 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( logger.ErrorCtx(
ctx, ctx,