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:
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user