From 9d999559f1eb3d4ff1d34fecee0ed921ae592f55 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@getprobo.com>
Date: Fri, 13 Mar 2026 19:53:08 +0100
Subject: [PATCH] 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>
---
 pkg/agent/agent_test.go      | 13 +++++---
 pkg/agent/agent_tool_test.go | 64 ++++++++++++++++++++++++++++++++++++
 pkg/agent/run.go             |  7 ++++
 3 files changed, 80 insertions(+), 4 deletions(-)

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,