diff --git a/internal/llminternal/base_flow.go b/internal/llminternal/base_flow.go index 388c25bba..c774aa034 100644 --- a/internal/llminternal/base_flow.go +++ b/internal/llminternal/base_flow.go @@ -760,6 +760,9 @@ func mergeEventActions(base, other *session.EventActions) *session.EventActions if base == nil { return other } + if other.ExitLoop { + base.ExitLoop = true + } if other.SkipSummarization { base.SkipSummarization = true } diff --git a/internal/toolinternal/context.go b/internal/toolinternal/context.go index 36d2a19ee..753e8ea78 100644 --- a/internal/toolinternal/context.go +++ b/internal/toolinternal/context.go @@ -123,11 +123,13 @@ func (c *toolContext) RequestConfirmation(hint string, payload any) error { Confirmed: false, Payload: payload, } - // SkipSummarization stops the agent loop after this tool call. Without it, - // the function response event becomes lastEvent and IsFinalResponse() returns - // false (hasFunctionResponses == true), causing the loop to continue. - // This matches the behavior of the built-in RequireConfirmation path in - // functiontool (function.go). - c.eventActions.SkipSummarization = true + c.ExitLoop() return nil } + +func (c *toolContext) ExitLoop() { + c.eventActions.ExitLoop = true + // Also set SkipSummarization for backward compatibility with older code + // that may not know about ExitLoop yet. + c.eventActions.SkipSummarization = true +} diff --git a/internal/toolinternal/context_test.go b/internal/toolinternal/context_test.go index a846e78e3..6b7ca46b0 100644 --- a/internal/toolinternal/context_test.go +++ b/internal/toolinternal/context_test.go @@ -34,7 +34,7 @@ func TestToolContext(t *testing.T) { } } -func TestRequestConfirmation_SetsSkipSummarization(t *testing.T) { +func TestRequestConfirmation_SetsExitLoop(t *testing.T) { inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) actions := &session.EventActions{} toolCtx := NewToolContext(inv, "fn1", actions, nil) @@ -44,8 +44,8 @@ func TestRequestConfirmation_SetsSkipSummarization(t *testing.T) { t.Fatalf("RequestConfirmation returned unexpected error: %v", err) } - if !actions.SkipSummarization { - t.Error("RequestConfirmation did not set SkipSummarization to true") + if !actions.ExitLoop { + t.Error("RequestConfirmation did not set ExitLoop to true") } if actions.RequestedToolConfirmations == nil { @@ -74,8 +74,8 @@ func TestRequestConfirmation_AutoGeneratesIDWhenEmpty(t *testing.T) { t.Fatalf("RequestConfirmation returned unexpected error: %v", err) } - if !actions.SkipSummarization { - t.Error("SkipSummarization should be set even with auto-generated function call ID") + if !actions.ExitLoop { + t.Error("ExitLoop should be set even with auto-generated function call ID") } if len(actions.RequestedToolConfirmations) != 1 { t.Fatalf("expected 1 confirmation entry, got %d", len(actions.RequestedToolConfirmations)) @@ -89,3 +89,23 @@ func TestRequestConfirmation_AutoGeneratesIDWhenEmpty(t *testing.T) { } } } + +func TestExitLoop(t *testing.T) { + inv := contextinternal.NewInvocationContext(t.Context(), contextinternal.InvocationContextParams{}) + actions := &session.EventActions{} + toolCtx := NewToolContext(inv, "fn1", actions, nil) + + if actions.ExitLoop { + t.Error("ExitLoop should be false initially") + } + + toolCtx.ExitLoop() + + if !actions.ExitLoop { + t.Error("ExitLoop should be true after calling ExitLoop()") + } + // SkipSummarization should also be set for backward compatibility + if !actions.SkipSummarization { + t.Error("SkipSummarization should be true for backward compatibility") + } +} diff --git a/session/session.go b/session/session.go index 471015212..cc25586cf 100644 --- a/session/session.go +++ b/session/session.go @@ -122,7 +122,7 @@ type Event struct { // Note: when multiple agents participate in one invocation, there could be // multiple events with IsFinalResponse() as True, for each participating agent. func (e *Event) IsFinalResponse() bool { - if (e.Actions.SkipSummarization) || len(e.LongRunningToolIDs) > 0 { + if e.Actions.ExitLoop || e.Actions.SkipSummarization || len(e.LongRunningToolIDs) > 0 { return true } @@ -150,8 +150,12 @@ type EventActions struct { RequestedToolConfirmations map[string]toolconfirmation.ToolConfirmation - // If true, it won't call model to summarize function response. - // Only valid for function response event. + // ExitLoop signals the agent to stop its event loop after processing this event. + // Use this when a tool needs to halt the agent's execution, such as when + // requiring user confirmation or when the tool's result is the final output. + ExitLoop bool + + // Deprecated: Use ExitLoop instead. SkipSummarization is kept for backward compatibility. SkipSummarization bool // If set, the event transfers to the specified agent. TransferToAgent string diff --git a/tool/agenttool/agent_tool.go b/tool/agenttool/agent_tool.go index 2cfe1a890..1200469d3 100644 --- a/tool/agenttool/agent_tool.go +++ b/tool/agenttool/agent_tool.go @@ -124,9 +124,7 @@ func (t *agentTool) Run(toolCtx tool.Context, args any) (map[string]any, error) } if t.skipSummarization { - if actions := toolCtx.Actions(); actions != nil { - actions.SkipSummarization = true - } + toolCtx.ExitLoop() } var agentInputSchema *genai.Schema diff --git a/tool/agenttool/agent_tool_test.go b/tool/agenttool/agent_tool_test.go index 2f1b3da0b..18521865f 100644 --- a/tool/agenttool/agent_tool_test.go +++ b/tool/agenttool/agent_tool_test.go @@ -272,7 +272,7 @@ func TestAgentTool_Run_SkipSummarization(t *testing.T) { agent := createAgentWithModel(t, nil, nil, testLLM) toolCtx := createToolContext(t, agent) - // Test with skipSummarization = true + // Test with skipSummarization = true (sets ExitLoop) agentToolSkip := agenttool.New(agent, &agenttool.Config{SkipSummarization: true}) actions := toolCtx.Actions() toolImpl, ok := agentToolSkip.(toolinternal.FunctionTool) @@ -283,8 +283,8 @@ func TestAgentTool_Run_SkipSummarization(t *testing.T) { if err != nil { t.Fatalf("Run() with skipSummarization=true failed unexpectedly: %v", err) } - if !actions.SkipSummarization { - t.Errorf("SkipSummarization flag not set when AgentTool was created with skipSummarization=true") + if !actions.ExitLoop { + t.Errorf("ExitLoop flag not set when AgentTool was created with skipSummarization=true") } // Test with skipSummarization = false @@ -293,7 +293,7 @@ func TestAgentTool_Run_SkipSummarization(t *testing.T) { if !ok { t.Fatal("agentToolNoSkip does not implement FunctionTool") } - actions.SkipSummarization = false // Reset + actions.ExitLoop = false // Reset // Reset mock for the second call testLLM.Responses = []*genai.Content{ genai.NewContentFromText("test response", genai.RoleModel), @@ -303,8 +303,8 @@ func TestAgentTool_Run_SkipSummarization(t *testing.T) { if err != nil { t.Fatalf("Run() with skipSummarization=false failed unexpectedly: %v", err) } - if actions.SkipSummarization { - t.Errorf("SkipSummarization flag was set when AgentTool was created with skipSummarization=false") + if actions.ExitLoop { + t.Errorf("ExitLoop flag was set when AgentTool was created with skipSummarization=false") } } diff --git a/tool/exitlooptool/tool.go b/tool/exitlooptool/tool.go index 30703ca28..d9008c810 100644 --- a/tool/exitlooptool/tool.go +++ b/tool/exitlooptool/tool.go @@ -27,7 +27,7 @@ type EmptyArgs struct{} func exitLoop(ctx tool.Context, myArgs EmptyArgs) (map[string]string, error) { ctx.Actions().Escalate = true - ctx.Actions().SkipSummarization = true + ctx.ExitLoop() return map[string]string{}, nil } diff --git a/tool/functiontool/function.go b/tool/functiontool/function.go index 42f417509..d352db7c2 100644 --- a/tool/functiontool/function.go +++ b/tool/functiontool/function.go @@ -218,7 +218,7 @@ func (f *functionTool[TArgs, TResults]) Run(ctx tool.Context, args any) (result if err != nil { return nil, err } - ctx.Actions().SkipSummarization = true + // RequestConfirmation calls ExitLoop() internally return nil, fmt.Errorf("error tool %q requires confirmation, please approve or reject", f.Name()) } } diff --git a/tool/mcptoolset/tool.go b/tool/mcptoolset/tool.go index 640c0cfc3..179d37235 100644 --- a/tool/mcptoolset/tool.go +++ b/tool/mcptoolset/tool.go @@ -111,7 +111,7 @@ func (t *mcpTool) Run(ctx tool.Context, args any) (map[string]any, error) { if err != nil { return nil, err } - ctx.Actions().SkipSummarization = true + // RequestConfirmation calls ExitLoop() internally return nil, fmt.Errorf("error tool %q requires confirmation, please approve or reject", t.Name()) } } diff --git a/tool/tool.go b/tool/tool.go index 7b9ceb567..e82a5eebb 100644 --- a/tool/tool.go +++ b/tool/tool.go @@ -87,6 +87,11 @@ type Context interface { // - error: If there was a failure in initiating the confirmation process itself (e.g., invalid // arguments, issue with the event system). The request to ask the user has not been sent. RequestConfirmation(hint string, payload any) error + + // ExitLoop signals the agent to stop its event loop after processing this tool's response. + // Call this when the tool's result should be the final output, or when the agent should + // halt execution (e.g., waiting for user input, confirmation, or external action). + ExitLoop() } // Toolset is an interface for a collection of tools. It allows grouping