From 2fdceaedb809fd01bb74165776e50b5d796a6059 Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Tue, 18 Aug 2026 14:33:20 +0200 Subject: [PATCH 1/2] fix: pair every tool call with a result before the model request A tool call and its result are persisted as two separate session events. If a turn ends between them, the session keeps a function call with no matching response, and replaying that history makes providers requiring strict pairing reject the whole conversation. Anthropic returns 400 "tool_use ids were found without tool_result blocks immediately after" on every later turn, leaving the session unusable until it is deleted. Repair the model request rather than the store: supply a placeholder result for a call that has none, and drop a result whose call is gone. Pairing is checked against the immediately following content, which is the invariant the provider enforces. Recorded history is untouched, so the UI still shows what happened and sessions already broken heal on their next turn without a migration. A conversation whose calls are all answered is left unchanged, and any conversation this does alter is one the provider would have rejected. Fixes #2277 Signed-off-by: QuentinBisson --- go/adk/pkg/agent/agent.go | 3 + go/adk/pkg/agent/tool_pairing.go | 165 ++++++++++++++ go/adk/pkg/agent/tool_pairing_test.go | 212 +++++++++++++++++ .../src/kagent/adk/_tool_pairing.py | 120 ++++++++++ .../kagent-adk/src/kagent/adk/types.py | 8 +- .../tests/unittests/test_tool_pairing.py | 215 ++++++++++++++++++ 6 files changed, 722 insertions(+), 1 deletion(-) create mode 100644 go/adk/pkg/agent/tool_pairing.go create mode 100644 go/adk/pkg/agent/tool_pairing_test.go create mode 100644 python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py create mode 100644 python/packages/kagent-adk/tests/unittests/test_tool_pairing.py diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index 6164f5d0d..7cc472cd6 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -109,6 +109,9 @@ func CreateGoogleADKAgent(ctx context.Context, agentConfig *adk.AgentConfig, age log.Info("Wiring MCP App model result callback", "toolCount", len(mcpAppToolNames)) beforeModelCallbacks = append(beforeModelCallbacks, MakeMCPAppModelResultCallback(mcpAppToolNames)) } + // Pairing repair runs last so it also covers anything the earlier callbacks + // leave unpaired. + beforeModelCallbacks = append(beforeModelCallbacks, MakeToolPairingCallback()) beforeToolCallbacks = append(beforeToolCallbacks, makeBeforeToolCallback(log)) llmAgentConfig := llmagent.Config{ diff --git a/go/adk/pkg/agent/tool_pairing.go b/go/adk/pkg/agent/tool_pairing.go new file mode 100644 index 000000000..fca72c95b --- /dev/null +++ b/go/adk/pkg/agent/tool_pairing.go @@ -0,0 +1,165 @@ +package agent + +import ( + "google.golang.org/adk/v2/agent" + "google.golang.org/adk/v2/agent/llmagent" + adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/genai" +) + +// missingToolResult stands in for a tool result that was never recorded. It +// deliberately states no cause: the call may have been interrupted, or may +// belong to a long-running tool that has not returned yet. It matches the +// placeholder the provider converters already use, so a request that reaches +// one of those unchanged behaves exactly as before. +const missingToolResult = "No response available for this function call." + +// MakeToolPairingCallback repairs tool call/response pairing in the model +// request. +// +// A tool call and its result are persisted as two separate session events. If a +// turn ends between them (process restart, OOM, client disconnect, cancellation, +// or a second message arriving while a slow tool is still running), the session +// keeps a function call with no matching response. History is replayed verbatim +// on every later turn, and providers that require strict pairing reject the +// whole conversation, leaving the session unusable until it is deleted. +// +// The repair runs against the request rather than the store, so recorded history +// stays intact for the UI and sessions already broken heal on their next turn. +// Pairing is checked positionally, against the immediately following content, +// because that is the invariant the provider enforces; a response elsewhere in +// the history does not satisfy it. +// +// A conversation whose calls are all answered is left untouched, and any +// conversation this does change is one the provider would have rejected. +func MakeToolPairingCallback() llmagent.BeforeModelCallback { + return func(_ agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { + if len(req.Contents) == 0 { + return nil, nil + } + req.Contents = synthesizeMissingResponses(dropOrphanedResponses(req.Contents)) + return nil, nil + } +} + +// dropOrphanedResponses removes function responses whose call is not present in +// the immediately preceding content, and drops any content left empty. +func dropOrphanedResponses(contents []*genai.Content) []*genai.Content { + kept := make([]*genai.Content, 0, len(contents)) + for index, content := range contents { + if content == nil { + continue + } + if !hasFunctionResponse(content) { + kept = append(kept, content) + continue + } + + var answerable map[string]bool + if index > 0 { + answerable = callIDs(contents[index-1]) + } + + parts := make([]*genai.Part, 0, len(content.Parts)) + for _, part := range content.Parts { + if part == nil { + continue + } + if part.FunctionResponse != nil && !answerable[part.FunctionResponse.ID] { + continue + } + parts = append(parts, part) + } + if len(parts) == 0 { + continue + } + content.Parts = parts + kept = append(kept, content) + } + return kept +} + +// synthesizeMissingResponses gives every function call a function response in +// the immediately following content. +func synthesizeMissingResponses(contents []*genai.Content) []*genai.Content { + repaired := make([]*genai.Content, 0, len(contents)) + for index, content := range contents { + repaired = append(repaired, content) + + calls := functionCalls(content) + if len(calls) == 0 { + continue + } + + var following *genai.Content + if index+1 < len(contents) { + following = contents[index+1] + } + answered := responseIDs(following) + + parts := make([]*genai.Part, 0, len(calls)) + for _, call := range calls { + if answered[call.ID] { + continue + } + parts = append(parts, &genai.Part{ + FunctionResponse: &genai.FunctionResponse{ + ID: call.ID, + Name: call.Name, + Response: map[string]any{"result": missingToolResult}, + }, + }) + } + if len(parts) == 0 { + continue + } + + // Join an existing response turn so the results stay in one message; + // otherwise the results need a turn of their own, before whatever + // currently follows. + if following != nil && len(answered) > 0 { + following.Parts = append(following.Parts, parts...) + continue + } + repaired = append(repaired, &genai.Content{Role: "user", Parts: parts}) + } + return repaired +} + +func functionCalls(content *genai.Content) []*genai.FunctionCall { + if content == nil { + return nil + } + var calls []*genai.FunctionCall + for _, part := range content.Parts { + if part != nil && part.FunctionCall != nil { + calls = append(calls, part.FunctionCall) + } + } + return calls +} + +func callIDs(content *genai.Content) map[string]bool { + ids := map[string]bool{} + for _, call := range functionCalls(content) { + ids[call.ID] = true + } + return ids +} + +func responseIDs(content *genai.Content) map[string]bool { + ids := map[string]bool{} + if content == nil { + return ids + } + for _, part := range content.Parts { + if part != nil && part.FunctionResponse != nil { + ids[part.FunctionResponse.ID] = true + } + } + return ids +} + +func hasFunctionResponse(content *genai.Content) bool { + return len(responseIDs(content)) > 0 +} diff --git a/go/adk/pkg/agent/tool_pairing_test.go b/go/adk/pkg/agent/tool_pairing_test.go new file mode 100644 index 000000000..d7e0d1f84 --- /dev/null +++ b/go/adk/pkg/agent/tool_pairing_test.go @@ -0,0 +1,212 @@ +package agent + +import ( + "testing" + + adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/genai" +) + +func callContent(ids ...string) *genai.Content { + parts := make([]*genai.Part, 0, len(ids)) + for _, id := range ids { + parts = append(parts, &genai.Part{FunctionCall: &genai.FunctionCall{ID: id, Name: "get_pods"}}) + } + return &genai.Content{Role: "model", Parts: parts} +} + +func responseContent(id string, result string) *genai.Content { + return &genai.Content{Role: "user", Parts: []*genai.Part{{ + FunctionResponse: &genai.FunctionResponse{ + ID: id, + Name: "get_pods", + Response: map[string]any{"result": result}, + }, + }}} +} + +func textContent(role, text string) *genai.Content { + return &genai.Content{Role: role, Parts: []*genai.Part{{Text: text}}} +} + +func repair(contents []*genai.Content) []*genai.Content { + req := &adkmodel.LLMRequest{Contents: contents} + if _, err := MakeToolPairingCallback()(nil, req); err != nil { + panic(err) + } + return req.Contents +} + +func responsesIn(content *genai.Content) []*genai.FunctionResponse { + var responses []*genai.FunctionResponse + for _, part := range content.Parts { + if part.FunctionResponse != nil { + responses = append(responses, part.FunctionResponse) + } + } + return responses +} + +func TestToolPairingSynthesizesResultForDanglingCall(t *testing.T) { + t.Parallel() + + contents := repair([]*genai.Content{textContent("user", "what's failing?"), callContent("abc")}) + + if len(contents) != 3 { + t.Fatalf("expected a synthesized response turn, got %d contents", len(contents)) + } + responses := responsesIn(contents[2]) + if len(responses) != 1 || responses[0].ID != "abc" { + t.Fatalf("expected one response for abc, got %+v", responses) + } + if got := responses[0].Response["result"]; got != missingToolResult { + t.Errorf("result = %v, want %q", got, missingToolResult) + } +} + +func TestToolPairingInsertsResultBeforeFollowingUserMessage(t *testing.T) { + t.Parallel() + + contents := repair([]*genai.Content{callContent("abc"), textContent("user", "second message")}) + + if len(contents) != 3 { + t.Fatalf("expected 3 contents, got %d", len(contents)) + } + if responses := responsesIn(contents[1]); len(responses) != 1 || responses[0].ID != "abc" { + t.Fatalf("expected the result between call and message, got %+v", responses) + } + if contents[2].Parts[0].Text != "second message" { + t.Errorf("user message was not preserved after the synthesized result") + } +} + +func TestToolPairingReusesSiblingResponseTurn(t *testing.T) { + t.Parallel() + + contents := repair([]*genai.Content{callContent("abc", "def"), responseContent("abc", "pod X running")}) + + if len(contents) != 2 { + t.Fatalf("expected the existing response turn to be reused, got %d contents", len(contents)) + } + responses := responsesIn(contents[1]) + if len(responses) != 2 { + t.Fatalf("expected 2 responses, got %d", len(responses)) + } + if responses[0].Response["result"] != "pod X running" { + t.Errorf("existing result was overwritten: %+v", responses[0]) + } + if responses[1].ID != "def" || responses[1].Response["result"] != missingToolResult { + t.Errorf("expected a placeholder for def, got %+v", responses[1]) + } +} + +func TestToolPairingRepairsNonAdjacentResponse(t *testing.T) { + t.Parallel() + + contents := repair([]*genai.Content{ + callContent("abc"), + textContent("user", "are you there?"), + responseContent("abc", "late result"), + }) + + responses := responsesIn(contents[1]) + if len(responses) != 1 || responses[0].Response["result"] != missingToolResult { + t.Fatalf("expected a placeholder adjacent to the call, got %+v", responses) + } +} + +func TestToolPairingDropsOrphanedResponse(t *testing.T) { + t.Parallel() + + contents := repair([]*genai.Content{textContent("user", "hello"), responseContent("abc", "stale")}) + + if len(contents) != 1 { + t.Fatalf("expected the orphaned response turn to be dropped, got %d contents", len(contents)) + } + if contents[0].Parts[0].Text != "hello" { + t.Errorf("wrong content survived: %+v", contents[0]) + } +} + +func TestToolPairingDropsOnlyTheOrphanedResponse(t *testing.T) { + t.Parallel() + + responseTurn := &genai.Content{Role: "user", Parts: []*genai.Part{ + {FunctionResponse: &genai.FunctionResponse{ID: "abc", Name: "get_pods", Response: map[string]any{"result": "ok"}}}, + {FunctionResponse: &genai.FunctionResponse{ID: "zzz", Name: "ghost", Response: map[string]any{"result": "stale"}}}, + }} + + contents := repair([]*genai.Content{callContent("abc"), responseTurn}) + + responses := responsesIn(contents[1]) + if len(responses) != 1 || responses[0].ID != "abc" { + t.Fatalf("expected only the orphan to be dropped, got %+v", responses) + } +} + +func TestToolPairingLeavesHealthyHistoryUntouched(t *testing.T) { + t.Parallel() + + original := []*genai.Content{ + textContent("user", "what's failing?"), + callContent("abc"), + responseContent("abc", "pod X running"), + textContent("model", "pod X is down"), + } + + contents := repair(original) + + if len(contents) != len(original) { + t.Fatalf("healthy history was modified: %d contents, want %d", len(contents), len(original)) + } + for i := range contents { + if contents[i] != original[i] { + t.Errorf("content %d was replaced", i) + } + } +} + +func TestToolPairingHandlesEmptyInput(t *testing.T) { + t.Parallel() + + if contents := repair(nil); len(contents) != 0 { + t.Errorf("expected no contents, got %d", len(contents)) + } + if contents := repair([]*genai.Content{nil, callContent("abc")}); len(contents) != 2 { + t.Errorf("expected the nil content to be dropped and the call answered, got %d", len(contents)) + } +} + +// TestToolPairingSatisfiesAnthropicAdjacency asserts the invariant the Anthropic +// API enforces: every tool_use must be answered in the immediately following +// message. +func TestToolPairingSatisfiesAnthropicAdjacency(t *testing.T) { + t.Parallel() + + cases := map[string][]*genai.Content{ + "interrupted turn": {textContent("user", "what's failing?"), callContent("abc")}, + "double message": {callContent("abc"), textContent("user", "second message")}, + "partial answer": {callContent("abc", "def"), responseContent("abc", "pod X running")}, + } + + for name, contents := range cases { + t.Run(name, func(t *testing.T) { + repaired := repair(contents) + for i, content := range repaired { + calls := functionCalls(content) + if len(calls) == 0 { + continue + } + if i+1 >= len(repaired) { + t.Fatalf("content %d ends the conversation with an unanswered call", i) + } + answered := responseIDs(repaired[i+1]) + for _, call := range calls { + if !answered[call.ID] { + t.Errorf("call %q has no result in the following content", call.ID) + } + } + } + }) + } +} diff --git a/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py b/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py new file mode 100644 index 000000000..b9e3cd89e --- /dev/null +++ b/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py @@ -0,0 +1,120 @@ +"""Tool call/response pairing repair for the model request. + +A tool call and its result are persisted as two separate session events. If a +turn ends between them (process restart, OOM, client disconnect, cancellation, +or a second message arriving while a slow tool is still running), the session +keeps a ``function_call`` with no matching ``function_response``. + +History is replayed verbatim on every later turn, and providers that require +strict pairing reject the whole conversation. Anthropic returns +``tool_use ids were found without tool_result blocks immediately after`` and the +session stays unusable until it is deleted. + +The repair runs against the model request, not the store: ADK builds the request +from deep copies of the session events (see ``flows/llm_flows/contents.py``), so +the recorded history stays intact for the UI and sessions already broken in the +field heal on their next turn without a migration. + +Pairing is checked positionally, against the immediately following content, +because that is the invariant the provider enforces. A response that exists +somewhere else in the history does not satisfy it. + +On a conversation whose calls are all answered this is a no-op, and any +conversation it does change is one the provider would have rejected outright. +""" + +from __future__ import annotations + +from google.adk.agents.callback_context import CallbackContext +from google.adk.models.llm_request import LlmRequest +from google.genai import types + +# Stands in for a result that was never recorded. Deliberately states no cause: +# the call may have been interrupted, or may belong to a long-running tool that +# has not returned yet. Matches the placeholder the OpenAI and Go converters +# already use, so providers that repair on their own keep the same behaviour. +MISSING_TOOL_RESULT = "No response available for this function call." + + +def _call_ids(content: types.Content | None) -> set[str | None]: + if content is None: + return set() + return {p.function_call.id for p in content.parts or [] if p.function_call} + + +def _response_ids(content: types.Content | None) -> set[str | None]: + if content is None: + return set() + return {p.function_response.id for p in content.parts or [] if p.function_response} + + +def _drop_orphaned_responses(contents: list[types.Content]) -> list[types.Content]: + """Remove function_response parts whose call is not in the preceding content.""" + kept: list[types.Content] = [] + for index, content in enumerate(contents): + responses = _response_ids(content) + if not responses: + kept.append(content) + continue + + answerable = _call_ids(contents[index - 1]) if index > 0 else set() + parts = [p for p in content.parts or [] if not p.function_response or p.function_response.id in answerable] + if not parts: + continue + content.parts = parts + kept.append(content) + return kept + + +def _synthesize_missing_responses(contents: list[types.Content]) -> list[types.Content]: + """Give every function_call a function_response in the immediately following content.""" + repaired: list[types.Content] = [] + for index, content in enumerate(contents): + repaired.append(content) + + calls = [p.function_call for p in content.parts or [] if p.function_call] + if not calls: + continue + + following = contents[index + 1] if index + 1 < len(contents) else None + answered = _response_ids(following) + missing = [call for call in calls if call.id not in answered] + if not missing: + continue + + parts = [ + types.Part( + function_response=types.FunctionResponse( + id=call.id, + name=call.name, + response={"result": MISSING_TOOL_RESULT}, + ) + ) + for call in missing + ] + + # Join an existing response turn so the results stay in one message; + # otherwise the results need a turn of their own, before whatever + # currently follows. + if following is not None and answered: + following.parts = list(following.parts or []) + parts + else: + repaired.append(types.Content(role="user", parts=parts)) + return repaired + + +def repair_tool_call_pairing_callback( + callback_context: CallbackContext, + llm_request: LlmRequest, +) -> None: + """Before-model callback that pairs every tool call with a tool result. + + Drops a result whose call is gone, then supplies a placeholder result for a + call that has none, so the request satisfies the strict call/result pairing + that Anthropic (and Bedrock) require. + """ + if not llm_request.contents: + return None + contents = _drop_orphaned_responses(list(llm_request.contents)) + llm_request.contents = _synthesize_missing_responses(contents) + return None diff --git a/python/packages/kagent-adk/src/kagent/adk/types.py b/python/packages/kagent-adk/src/kagent/adk/types.py index 03d64ab2b..1e16bfc10 100644 --- a/python/packages/kagent-adk/src/kagent/adk/types.py +++ b/python/packages/kagent-adk/src/kagent/adk/types.py @@ -17,6 +17,7 @@ from kagent.adk._mcp_apps import MCPAppToolNames, make_mcp_app_model_result_callback from kagent.adk._mcp_toolset import KAgentMcpToolset from kagent.adk._remote_a2a_tool import KAgentRemoteA2AToolset +from kagent.adk._tool_pairing import repair_tool_call_pairing_callback from kagent.adk.models._anthropic import KAgentAnthropicLlm from kagent.adk.models._bedrock import KAgentBedrockLlm from kagent.adk.models._gemini import KAgentGeminiLlm, KAgentGeminiVertexAILlm @@ -538,7 +539,12 @@ async def rewrite_url_to_proxy(request: httpx.Request) -> None: # Build before_tool_callback if any tools require approval before_tool_callback = make_approval_callback(tools_requiring_approval) if tools_requiring_approval else None # ADK 2.x filters its synthetic confirmation events before model calls. - before_model_callbacks = [make_mcp_app_model_result_callback(mcp_app_tool_names)] + # Pairing repair runs last so it also covers anything the earlier + # callbacks leave unpaired. + before_model_callbacks = [ + make_mcp_app_model_result_callback(mcp_app_tool_names), + repair_tool_call_pairing_callback, + ] # static_instruction is sent directly to the model without any placeholder processing agent = Agent( diff --git a/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py b/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py new file mode 100644 index 000000000..1bc7f3b42 --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py @@ -0,0 +1,215 @@ +"""Tests for tool call/response pairing repair.""" + +from google.adk.models.anthropic_llm import content_to_message_param +from google.adk.models.llm_request import LlmRequest +from google.genai import types + +from kagent.adk._tool_pairing import ( + MISSING_TOOL_RESULT, + repair_tool_call_pairing_callback, +) + + +def _call(call_id: str, name: str = "get_pods") -> types.Content: + return types.Content( + role="model", + parts=[types.Part(function_call=types.FunctionCall(id=call_id, name=name, args={"ns": "default"}))], + ) + + +def _response(call_id: str, name: str = "get_pods", text: str = "pod X running") -> types.Content: + return types.Content( + role="user", + parts=[types.Part(function_response=types.FunctionResponse(id=call_id, name=name, response={"result": text}))], + ) + + +def _text(role: str, text: str) -> types.Content: + return types.Content(role=role, parts=[types.Part(text=text)]) + + +def _repair(contents: list[types.Content]) -> list[types.Content]: + request = LlmRequest(contents=contents) + repair_tool_call_pairing_callback(callback_context=None, llm_request=request) + return request.contents + + +def _responses_in(content: types.Content) -> list[types.FunctionResponse]: + return [p.function_response for p in content.parts or [] if p.function_response] + + +class TestSynthesizeMissingResponses: + def test_dangling_call_at_end_of_history(self): + """The interrupted-turn case from the issue: nothing follows the call.""" + contents = _repair([_text("user", "what's failing?"), _call("abc")]) + + assert len(contents) == 3 + responses = _responses_in(contents[2]) + assert len(responses) == 1 + assert responses[0].id == "abc" + assert responses[0].name == "get_pods" + assert responses[0].response == {"result": MISSING_TOOL_RESULT} + assert contents[2].role == "user" + + def test_result_inserted_before_a_following_user_message(self): + """The double-message race: a second message arrives before the result.""" + contents = _repair([_call("abc"), _text("user", "second message")]) + + assert len(contents) == 3 + assert _responses_in(contents[1])[0].id == "abc" + assert contents[2].parts[0].text == "second message" + + def test_sibling_response_turn_is_reused(self): + """A partially answered call turn gets its gap filled, not a new turn.""" + call_turn = types.Content( + role="model", + parts=[ + types.Part(function_call=types.FunctionCall(id="abc", name="get_pods", args={})), + types.Part(function_call=types.FunctionCall(id="def", name="get_logs", args={})), + ], + ) + contents = _repair([call_turn, _response("abc")]) + + assert len(contents) == 2 + responses = _responses_in(contents[1]) + assert [r.id for r in responses] == ["abc", "def"] + assert responses[0].response == {"result": "pod X running"} + assert responses[1].response == {"result": MISSING_TOOL_RESULT} + + def test_only_the_unanswered_call_is_synthesized(self): + call_turn = types.Content( + role="model", + parts=[ + types.Part(function_call=types.FunctionCall(id="abc", name="get_pods", args={})), + types.Part(function_call=types.FunctionCall(id="def", name="get_logs", args={})), + ], + ) + contents = _repair([call_turn, _response("def", name="get_logs", text="log line")]) + + responses = _responses_in(contents[1]) + assert {r.id for r in responses} == {"abc", "def"} + by_id = {r.id: r.response for r in responses} + assert by_id["def"] == {"result": "log line"} + assert by_id["abc"] == {"result": MISSING_TOOL_RESULT} + + def test_non_adjacent_response_is_still_repaired(self): + """Positional pairing: a response elsewhere does not satisfy the provider.""" + contents = _repair([_call("abc"), _text("user", "are you there?"), _response("abc")]) + + assert _responses_in(contents[1])[0].response == {"result": MISSING_TOOL_RESULT} + + def test_call_without_id_still_gets_a_response(self): + contents = _repair( + [ + types.Content( + role="model", parts=[types.Part(function_call=types.FunctionCall(name="get_pods", args={}))] + ) + ] + ) + + assert len(contents) == 2 + assert _responses_in(contents[1])[0].id is None + + +class TestDropOrphanedResponses: + def test_response_without_a_call_is_dropped(self): + contents = _repair([_text("user", "hello"), _response("abc")]) + + assert len(contents) == 1 + assert contents[0].parts[0].text == "hello" + + def test_only_the_orphan_is_dropped(self): + response_turn = types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse(id="abc", name="get_pods", response={"result": "ok"}) + ), + types.Part( + function_response=types.FunctionResponse(id="zzz", name="ghost", response={"result": "stale"}) + ), + ], + ) + contents = _repair([_call("abc"), response_turn]) + + responses = _responses_in(contents[1]) + assert [r.id for r in responses] == ["abc"] + + def test_surrounding_parts_survive_an_orphan(self): + response_turn = types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse(id="zzz", name="ghost", response={"result": "stale"}) + ), + types.Part(text="and here is my question"), + ], + ) + contents = _repair([_text("user", "hello"), response_turn]) + + assert len(contents) == 2 + assert _responses_in(contents[1]) == [] + assert contents[1].parts[0].text == "and here is my question" + + +class TestDoesNotDisturbHealthyHistory: + def test_paired_history_is_unchanged(self): + original = [_text("user", "what's failing?"), _call("abc"), _response("abc"), _text("model", "pod X is down")] + before = [c.model_dump_json() for c in original] + + contents = _repair(original) + + assert [c.model_dump_json() for c in contents] == before + + def test_history_without_tools_is_unchanged(self): + original = [_text("user", "hi"), _text("model", "hello")] + before = [c.model_dump_json() for c in original] + + assert [c.model_dump_json() for c in _repair(original)] == before + + def test_empty_and_missing_parts_do_not_raise(self): + request = LlmRequest(contents=[]) + repair_tool_call_pairing_callback(callback_context=None, llm_request=request) + assert request.contents == [] + + contents = _repair([types.Content(role="user", parts=None), _call("abc")]) + assert _responses_in(contents[-1])[0].id == "abc" + + +class TestAnthropicPairingInvariant: + """The invariant the Anthropic API enforces, over repaired contents.""" + + @staticmethod + def _assert_paired(contents: list[types.Content]) -> None: + messages = [content_to_message_param(c) for c in contents] + for index, message in enumerate(messages): + call_ids = [b["id"] for b in message["content"] if b.get("type") == "tool_use"] + if not call_ids: + continue + assert index + 1 < len(messages), f"message {index} ends the conversation with an unanswered tool_use" + following = messages[index + 1] + answered = [b["tool_use_id"] for b in following["content"] if b.get("type") == "tool_result"] + assert set(call_ids) <= set(answered), f"message {index} has tool_use ids without tool_result: {call_ids}" + + def test_interrupted_turn_produces_a_valid_conversation(self): + self._assert_paired(_repair([_text("user", "what's failing?"), _call("abc")])) + + def test_double_message_race_produces_a_valid_conversation(self): + self._assert_paired(_repair([_call("abc"), _text("user", "second message")])) + + def test_partially_answered_turn_produces_a_valid_conversation(self): + call_turn = types.Content( + role="model", + parts=[ + types.Part(function_call=types.FunctionCall(id="abc", name="get_pods", args={})), + types.Part(function_call=types.FunctionCall(id="def", name="get_logs", args={})), + ], + ) + self._assert_paired(_repair([call_turn, _response("abc")])) + + def test_unrepaired_history_would_fail_the_invariant(self): + """Guards the tests themselves: the invariant must reject the broken input.""" + import pytest + + with pytest.raises(AssertionError): + self._assert_paired([_text("user", "what's failing?"), _call("abc")]) From b42e24a7856143cd6d9c0e7d41f7691cfb6db4bf Mon Sep 17 00:00:00 2001 From: QuentinBisson Date: Wed, 19 Aug 2026 01:38:29 +0200 Subject: [PATCH 2/2] fix: distinguish pending tool calls from lost ones when pairing Read long_running_tool_ids from the session so a call ADK is holding open for a human approval, ask_user, or a long-running tool is described as awaiting a response rather than as having returned nothing. Drop the orphaned-response pass. ADK removes an orphaned function response before the callback runs, and recovers the compaction case by re-injecting the missing call event, so the pass could only discard a real result. Match responses to calls by consuming ids one at a time so several calls with no id in one turn each get their own result. Log when a request had to be repaired. Signed-off-by: QuentinBisson --- go/adk/pkg/agent/agent.go | 2 +- go/adk/pkg/agent/tool_pairing.go | 197 +++++++++------ go/adk/pkg/agent/tool_pairing_test.go | 226 ++++++++++++++---- .../src/kagent/adk/_tool_pairing.py | 154 +++++++----- .../tests/unittests/test_tool_pairing.py | 163 ++++++++----- 5 files changed, 496 insertions(+), 246 deletions(-) diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index 7cc472cd6..69caa059d 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -111,7 +111,7 @@ func CreateGoogleADKAgent(ctx context.Context, agentConfig *adk.AgentConfig, age } // Pairing repair runs last so it also covers anything the earlier callbacks // leave unpaired. - beforeModelCallbacks = append(beforeModelCallbacks, MakeToolPairingCallback()) + beforeModelCallbacks = append(beforeModelCallbacks, MakeToolPairingCallback(log)) beforeToolCallbacks = append(beforeToolCallbacks, makeBeforeToolCallback(log)) llmAgentConfig := llmagent.Config{ diff --git a/go/adk/pkg/agent/tool_pairing.go b/go/adk/pkg/agent/tool_pairing.go index fca72c95b..e771cdf78 100644 --- a/go/adk/pkg/agent/tool_pairing.go +++ b/go/adk/pkg/agent/tool_pairing.go @@ -1,89 +1,113 @@ package agent import ( + "github.com/go-logr/logr" "google.golang.org/adk/v2/agent" "google.golang.org/adk/v2/agent/llmagent" adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/session" "google.golang.org/genai" ) -// missingToolResult stands in for a tool result that was never recorded. It -// deliberately states no cause: the call may have been interrupted, or may -// belong to a long-running tool that has not returned yet. It matches the -// placeholder the provider converters already use, so a request that reaches -// one of those unchanged behaves exactly as before. +// missingToolResult stands in for a result that was never recorded. It states +// no cause: the call may have been interrupted, or the process may have died +// between the two events. It matches the placeholder the Anthropic and OpenAI +// converters already use, so a request that reaches one of those unchanged +// behaves exactly as before. const missingToolResult = "No response available for this function call." +// pendingToolResult stands in for a call ADK is deliberately holding open: a +// long-running tool, a human approval, or ask_user. No function response exists +// for these until the answer arrives, so the call is pending rather than lost. +// The distinction matters: told a tool "returned nothing", a model reissues it +// or proceeds; told the call is still awaiting a response, it can wait. +const pendingToolResult = "This call is awaiting a response and has not completed yet." + // MakeToolPairingCallback repairs tool call/response pairing in the model // request. // // A tool call and its result are persisted as two separate session events. If a -// turn ends between them (process restart, OOM, client disconnect, cancellation, -// or a second message arriving while a slow tool is still running), the session -// keeps a function call with no matching response. History is replayed verbatim -// on every later turn, and providers that require strict pairing reject the -// whole conversation, leaving the session unusable until it is deleted. +// turn ends between them (process restart, OOM, client disconnect, +// cancellation), the session keeps a function call with no matching response. +// History is replayed on every later turn, and providers that require strict +// pairing reject the whole conversation, leaving the session unusable until it +// is deleted. +// +// The repair runs against the request rather than the store, so recorded +// history stays intact and sessions already broken heal on their next turn. It +// also sits above the session service, so it applies whichever store backs the +// session. +// +// ADK hands each request a deep copy of the content, so assigning Parts here +// cannot reach persisted history. // -// The repair runs against the request rather than the store, so recorded history -// stays intact for the UI and sessions already broken heal on their next turn. // Pairing is checked positionally, against the immediately following content, -// because that is the invariant the provider enforces; a response elsewhere in -// the history does not satisfy it. +// because that is the invariant the provider enforces. // // A conversation whose calls are all answered is left untouched, and any // conversation this does change is one the provider would have rejected. -func MakeToolPairingCallback() llmagent.BeforeModelCallback { - return func(_ agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { +func MakeToolPairingCallback(log logr.Logger) llmagent.BeforeModelCallback { + return func(ctx agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { if len(req.Contents) == 0 { return nil, nil } - req.Contents = synthesizeMissingResponses(dropOrphanedResponses(req.Contents)) + var source sessionSource + if ctx != nil { + source = ctx + } + repaired, synthesized := synthesizeMissingResponses(req.Contents, pendingCallIDs(source)) + if synthesized > 0 { + // Logged because a repaired request means a turn ended between a + // tool call and its result somewhere upstream. + log.Info("Paired unanswered tool calls before the model request", "count", synthesized) + } + req.Contents = repaired return nil, nil } } -// dropOrphanedResponses removes function responses whose call is not present in -// the immediately preceding content, and drops any content left empty. -func dropOrphanedResponses(contents []*genai.Content) []*genai.Content { - kept := make([]*genai.Content, 0, len(contents)) - for index, content := range contents { - if content == nil { - continue - } - if !hasFunctionResponse(content) { - kept = append(kept, content) - continue - } - - var answerable map[string]bool - if index > 0 { - answerable = callIDs(contents[index-1]) - } +// sessionSource is the part of agent.Context this file needs. +type sessionSource interface { + Session() session.Session +} - parts := make([]*genai.Part, 0, len(content.Parts)) - for _, part := range content.Parts { - if part == nil { - continue - } - if part.FunctionResponse != nil && !answerable[part.FunctionResponse.ID] { - continue - } - parts = append(parts, part) - } - if len(parts) == 0 { +// pendingCallIDs returns the ids of calls ADK is holding open for a +// long-running tool or an approval. They are read from the session events +// because LongRunningToolIDs lives on the event and does not survive the +// conversion to contents. +func pendingCallIDs(source sessionSource) map[string]bool { + pending := map[string]bool{} + if source == nil { + return pending + } + current := source.Session() + if current == nil { + return pending + } + events := current.Events() + if events == nil { + return pending + } + for event := range events.All() { + if event == nil { continue } - content.Parts = parts - kept = append(kept, content) + for _, id := range event.LongRunningToolIDs { + pending[id] = true + } } - return kept + return pending } // synthesizeMissingResponses gives every function call a function response in -// the immediately following content. -func synthesizeMissingResponses(contents []*genai.Content) []*genai.Content { +// the immediately following content, and reports how many it had to supply. +func synthesizeMissingResponses(contents []*genai.Content, pending map[string]bool) ([]*genai.Content, int) { + synthesized := 0 repaired := make([]*genai.Content, 0, len(contents)) for index, content := range contents { + if content == nil { + continue + } repaired = append(repaired, content) calls := functionCalls(content) @@ -96,23 +120,22 @@ func synthesizeMissingResponses(contents []*genai.Content) []*genai.Content { following = contents[index+1] } answered := responseIDs(following) + missing := unanswered(calls, answered) + if len(missing) == 0 { + continue + } - parts := make([]*genai.Part, 0, len(calls)) - for _, call := range calls { - if answered[call.ID] { - continue - } + synthesized += len(missing) + parts := make([]*genai.Part, 0, len(missing)) + for _, call := range missing { parts = append(parts, &genai.Part{ FunctionResponse: &genai.FunctionResponse{ ID: call.ID, Name: call.Name, - Response: map[string]any{"result": missingToolResult}, + Response: map[string]any{"result": placeholderFor(call, pending)}, }, }) } - if len(parts) == 0 { - continue - } // Join an existing response turn so the results stay in one message; // otherwise the results need a turn of their own, before whatever @@ -121,9 +144,41 @@ func synthesizeMissingResponses(contents []*genai.Content) []*genai.Content { following.Parts = append(following.Parts, parts...) continue } - repaired = append(repaired, &genai.Content{Role: "user", Parts: parts}) + repaired = append(repaired, &genai.Content{Role: genai.RoleUser, Parts: parts}) } - return repaired + return repaired, synthesized +} + +func placeholderFor(call *genai.FunctionCall, pending map[string]bool) string { + if call.ID != "" && pending[call.ID] { + return pendingToolResult + } + return missingToolResult +} + +// unanswered returns the calls with no response in answered. Ids are consumed +// one at a time rather than matched through a set, so a turn carrying several +// calls with no id (Gemini omits them, and ADK strips its own "adk-" ids before +// the request is built) does not have one response silently answer all of them. +func unanswered(calls []*genai.FunctionCall, answered []string) []*genai.FunctionCall { + remaining := make([]string, len(answered)) + copy(remaining, answered) + + var missing []*genai.FunctionCall + for _, call := range calls { + matched := false + for i, id := range remaining { + if id == call.ID { + remaining = append(remaining[:i], remaining[i+1:]...) + matched = true + break + } + } + if !matched { + missing = append(missing, call) + } + } + return missing } func functionCalls(content *genai.Content) []*genai.FunctionCall { @@ -139,27 +194,15 @@ func functionCalls(content *genai.Content) []*genai.FunctionCall { return calls } -func callIDs(content *genai.Content) map[string]bool { - ids := map[string]bool{} - for _, call := range functionCalls(content) { - ids[call.ID] = true - } - return ids -} - -func responseIDs(content *genai.Content) map[string]bool { - ids := map[string]bool{} +func responseIDs(content *genai.Content) []string { if content == nil { - return ids + return nil } + var ids []string for _, part := range content.Parts { if part != nil && part.FunctionResponse != nil { - ids[part.FunctionResponse.ID] = true + ids = append(ids, part.FunctionResponse.ID) } } return ids } - -func hasFunctionResponse(content *genai.Content) bool { - return len(responseIDs(content)) > 0 -} diff --git a/go/adk/pkg/agent/tool_pairing_test.go b/go/adk/pkg/agent/tool_pairing_test.go index d7e0d1f84..989ba1693 100644 --- a/go/adk/pkg/agent/tool_pairing_test.go +++ b/go/adk/pkg/agent/tool_pairing_test.go @@ -1,9 +1,14 @@ package agent import ( + "iter" "testing" + "time" + + "github.com/go-logr/logr" adkmodel "google.golang.org/adk/v2/model" + "google.golang.org/adk/v2/session" "google.golang.org/genai" ) @@ -16,7 +21,7 @@ func callContent(ids ...string) *genai.Content { } func responseContent(id string, result string) *genai.Content { - return &genai.Content{Role: "user", Parts: []*genai.Part{{ + return &genai.Content{Role: genai.RoleUser, Parts: []*genai.Part{{ FunctionResponse: &genai.FunctionResponse{ ID: id, Name: "get_pods", @@ -29,12 +34,52 @@ func textContent(role, text string) *genai.Content { return &genai.Content{Role: role, Parts: []*genai.Part{{Text: text}}} } -func repair(contents []*genai.Content) []*genai.Content { - req := &adkmodel.LLMRequest{Contents: contents} - if _, err := MakeToolPairingCallback()(nil, req); err != nil { - panic(err) +// fakeEvents / fakeSession / fakeContext stand in for the session an +// agent.Context exposes. Only Session() is consulted by the code under test. +type fakeEvents []*session.Event + +func (e fakeEvents) All() iter.Seq[*session.Event] { + return func(yield func(*session.Event) bool) { + for _, event := range e { + if !yield(event) { + return + } + } } - return req.Contents +} +func (e fakeEvents) Len() int { return len(e) } +func (e fakeEvents) At(i int) *session.Event { return e[i] } + +type fakeSession struct{ events fakeEvents } + +func (s fakeSession) ID() string { return "s" } +func (s fakeSession) AppName() string { return "app" } +func (s fakeSession) UserID() string { return "u" } +func (s fakeSession) State() session.State { return nil } +func (s fakeSession) Events() session.Events { return s.events } +func (s fakeSession) LastUpdateTime() time.Time { return time.Time{} } + +type fakeContext struct{ session session.Session } + +func (c fakeContext) Session() session.Session { return c.session } + +// pendingEvent is an event whose tool calls ADK is holding open (approval, +// ask_user, long-running tool). +func pendingEvent(ids ...string) *session.Event { + return &session.Event{LongRunningToolIDs: ids} +} + +func withPending(ids ...string) sessionSource { + return fakeContext{session: fakeSession{events: fakeEvents{pendingEvent(ids...)}}} +} + +func repair(contents []*genai.Content) []*genai.Content { + return repairWith(contents, nil) +} + +func repairWith(contents []*genai.Content, source sessionSource) []*genai.Content { + repaired, _ := synthesizeMissingResponses(contents, pendingCallIDs(source)) + return repaired } func responsesIn(content *genai.Content) []*genai.FunctionResponse { @@ -47,10 +92,22 @@ func responsesIn(content *genai.Content) []*genai.FunctionResponse { return responses } +func TestToolPairingCallbackWiring(t *testing.T) { + t.Parallel() + + req := &adkmodel.LLMRequest{Contents: []*genai.Content{callContent("abc")}} + if _, err := MakeToolPairingCallback(logr.Discard())(nil, req); err != nil { + t.Fatalf("callback returned error: %v", err) + } + if len(req.Contents) != 2 { + t.Fatalf("expected the callback to answer the call, got %d contents", len(req.Contents)) + } +} + func TestToolPairingSynthesizesResultForDanglingCall(t *testing.T) { t.Parallel() - contents := repair([]*genai.Content{textContent("user", "what's failing?"), callContent("abc")}) + contents := repair([]*genai.Content{textContent(genai.RoleUser, "what's failing?"), callContent("abc")}) if len(contents) != 3 { t.Fatalf("expected a synthesized response turn, got %d contents", len(contents)) @@ -67,7 +124,7 @@ func TestToolPairingSynthesizesResultForDanglingCall(t *testing.T) { func TestToolPairingInsertsResultBeforeFollowingUserMessage(t *testing.T) { t.Parallel() - contents := repair([]*genai.Content{callContent("abc"), textContent("user", "second message")}) + contents := repair([]*genai.Content{callContent("abc"), textContent(genai.RoleUser, "second message")}) if len(contents) != 3 { t.Fatalf("expected 3 contents, got %d", len(contents)) @@ -105,7 +162,7 @@ func TestToolPairingRepairsNonAdjacentResponse(t *testing.T) { contents := repair([]*genai.Content{ callContent("abc"), - textContent("user", "are you there?"), + textContent(genai.RoleUser, "are you there?"), responseContent("abc", "late result"), }) @@ -115,32 +172,63 @@ func TestToolPairingRepairsNonAdjacentResponse(t *testing.T) { } } -func TestToolPairingDropsOrphanedResponse(t *testing.T) { +// TestToolPairingIDLessSiblingsEachGetAResponse guards the id-matching rule: +// Gemini omits call ids and ADK strips its own "adk-" ids before the request is +// built, so one response must not answer every id-less call in the turn. +func TestToolPairingIDLessSiblingsEachGetAResponse(t *testing.T) { t.Parallel() - contents := repair([]*genai.Content{textContent("user", "hello"), responseContent("abc", "stale")}) + callTurn := &genai.Content{Role: "model", Parts: []*genai.Part{ + {FunctionCall: &genai.FunctionCall{Name: "get_pods"}}, + {FunctionCall: &genai.FunctionCall{Name: "get_logs"}}, + }} + answered := &genai.Content{Role: genai.RoleUser, Parts: []*genai.Part{ + {FunctionResponse: &genai.FunctionResponse{Name: "get_pods", Response: map[string]any{"result": "ok"}}}, + }} + + contents := repair([]*genai.Content{callTurn, answered}) - if len(contents) != 1 { - t.Fatalf("expected the orphaned response turn to be dropped, got %d contents", len(contents)) + responses := responsesIn(contents[1]) + if len(responses) != 2 { + t.Fatalf("expected both id-less calls to be answered, got %d responses", len(responses)) } - if contents[0].Parts[0].Text != "hello" { - t.Errorf("wrong content survived: %+v", contents[0]) + if responses[0].Response["result"] != "ok" { + t.Errorf("existing result was overwritten: %+v", responses[0]) + } + if responses[1].Response["result"] != missingToolResult { + t.Errorf("expected a placeholder for the second call, got %+v", responses[1]) } } -func TestToolPairingDropsOnlyTheOrphanedResponse(t *testing.T) { +// TestToolPairingPendingApprovalKeepsPendingWording covers a call ADK is holding +// open. A plain message sent while an approval is pending reaches this callback +// as an ordinary turn, and reporting an empty result would invite the model to +// reissue the approval or proceed without it. +func TestToolPairingPendingApprovalKeepsPendingWording(t *testing.T) { t.Parallel() - responseTurn := &genai.Content{Role: "user", Parts: []*genai.Part{ - {FunctionResponse: &genai.FunctionResponse{ID: "abc", Name: "get_pods", Response: map[string]any{"result": "ok"}}}, - {FunctionResponse: &genai.FunctionResponse{ID: "zzz", Name: "ghost", Response: map[string]any{"result": "stale"}}}, - }} + contents := repairWith( + []*genai.Content{callContent("call-1"), textContent(genai.RoleUser, "never mind, what about X?")}, + withPending("call-1"), + ) + + responses := responsesIn(contents[1]) + if len(responses) != 1 || responses[0].Response["result"] != pendingToolResult { + t.Fatalf("expected the pending placeholder, got %+v", responses) + } +} + +func TestToolPairingInterruptedCallKeepsMissingWording(t *testing.T) { + t.Parallel() - contents := repair([]*genai.Content{callContent("abc"), responseTurn}) + contents := repairWith( + []*genai.Content{callContent("other"), textContent(genai.RoleUser, "still there?")}, + withPending("call-1"), + ) responses := responsesIn(contents[1]) - if len(responses) != 1 || responses[0].ID != "abc" { - t.Fatalf("expected only the orphan to be dropped, got %+v", responses) + if len(responses) != 1 || responses[0].Response["result"] != missingToolResult { + t.Fatalf("expected the neutral placeholder, got %+v", responses) } } @@ -148,7 +236,7 @@ func TestToolPairingLeavesHealthyHistoryUntouched(t *testing.T) { t.Parallel() original := []*genai.Content{ - textContent("user", "what's failing?"), + textContent(genai.RoleUser, "what's failing?"), callContent("abc"), responseContent("abc", "pod X running"), textContent("model", "pod X is down"), @@ -166,6 +254,21 @@ func TestToolPairingLeavesHealthyHistoryUntouched(t *testing.T) { } } +func TestToolPairingLeavesAnsweredPendingCallUntouched(t *testing.T) { + t.Parallel() + + original := []*genai.Content{callContent("call-1"), responseContent("call-1", "approved")} + + contents := repairWith(original, withPending("call-1")) + + if len(contents) != 2 { + t.Fatalf("expected the answered pending call to be left alone, got %d contents", len(contents)) + } + if responsesIn(contents[1])[0].Response["result"] != "approved" { + t.Errorf("real result was replaced: %+v", responsesIn(contents[1])[0]) + } +} + func TestToolPairingHandlesEmptyInput(t *testing.T) { t.Parallel() @@ -177,36 +280,65 @@ func TestToolPairingHandlesEmptyInput(t *testing.T) { } } -// TestToolPairingSatisfiesAnthropicAdjacency asserts the invariant the Anthropic -// API enforces: every tool_use must be answered in the immediately following -// message. -func TestToolPairingSatisfiesAnthropicAdjacency(t *testing.T) { +// assertPaired walks the contents independently of the production helpers and +// asserts the invariant the Anthropic and Bedrock APIs enforce: every tool call +// must be answered in the immediately following content. +func assertPaired(t *testing.T, contents []*genai.Content) { + t.Helper() + for i, content := range contents { + var ids []string + for _, part := range content.Parts { + if part != nil && part.FunctionCall != nil { + ids = append(ids, part.FunctionCall.ID) + } + } + if len(ids) == 0 { + continue + } + if i+1 >= len(contents) { + t.Errorf("content %d ends the conversation with an unanswered call", i) + continue + } + answered := map[string]int{} + for _, part := range contents[i+1].Parts { + if part != nil && part.FunctionResponse != nil { + answered[part.FunctionResponse.ID]++ + } + } + for _, id := range ids { + if answered[id] == 0 { + t.Errorf("call %q has no result in the following content", id) + continue + } + answered[id]-- + } + } +} + +func TestToolPairingSatisfiesProviderAdjacency(t *testing.T) { t.Parallel() cases := map[string][]*genai.Content{ - "interrupted turn": {textContent("user", "what's failing?"), callContent("abc")}, - "double message": {callContent("abc"), textContent("user", "second message")}, - "partial answer": {callContent("abc", "def"), responseContent("abc", "pod X running")}, + "interrupted turn": {textContent(genai.RoleUser, "what's failing?"), callContent("abc")}, + "message during running tool": {callContent("abc"), textContent(genai.RoleUser, "second message")}, + "partial answer": {callContent("abc", "def"), responseContent("abc", "pod X running")}, } for name, contents := range cases { t.Run(name, func(t *testing.T) { - repaired := repair(contents) - for i, content := range repaired { - calls := functionCalls(content) - if len(calls) == 0 { - continue - } - if i+1 >= len(repaired) { - t.Fatalf("content %d ends the conversation with an unanswered call", i) - } - answered := responseIDs(repaired[i+1]) - for _, call := range calls { - if !answered[call.ID] { - t.Errorf("call %q has no result in the following content", call.ID) - } - } - } + assertPaired(t, repair(contents)) }) } } + +// TestUnrepairedHistoryFailsAdjacency guards the assertion above: it must +// reject the broken input, otherwise it would pass on everything. +func TestUnrepairedHistoryFailsAdjacency(t *testing.T) { + t.Parallel() + + fake := &testing.T{} + assertPaired(fake, []*genai.Content{textContent(genai.RoleUser, "what's failing?"), callContent("abc")}) + if !fake.Failed() { + t.Error("adjacency assertion accepted an unrepaired conversation") + } +} diff --git a/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py b/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py index b9e3cd89e..467119978 100644 --- a/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py +++ b/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py @@ -1,23 +1,29 @@ """Tool call/response pairing repair for the model request. A tool call and its result are persisted as two separate session events. If a -turn ends between them (process restart, OOM, client disconnect, cancellation, -or a second message arriving while a slow tool is still running), the session -keeps a ``function_call`` with no matching ``function_response``. - -History is replayed verbatim on every later turn, and providers that require -strict pairing reject the whole conversation. Anthropic returns -``tool_use ids were found without tool_result blocks immediately after`` and the -session stays unusable until it is deleted. - -The repair runs against the model request, not the store: ADK builds the request -from deep copies of the session events (see ``flows/llm_flows/contents.py``), so -the recorded history stays intact for the UI and sessions already broken in the -field heal on their next turn without a migration. +turn ends between them (process restart, OOM, client disconnect, cancellation), +the session keeps a ``function_call`` with no matching ``function_response``. + +History is replayed on every later turn, and providers that require strict +pairing reject the whole conversation. Anthropic returns ``tool_use ids were +found without tool_result blocks immediately after`` and the session stays +unusable until it is deleted. + +The repair runs against the model request rather than the store, so recorded +history stays intact and sessions already broken in the field heal on their next +turn without a migration. It also sits above the session service, so it applies +whichever store backs the session. + +ADK hands each request session-isolated copies of the content +(``_copy_content_for_request`` in ``flows/llm_flows/contents.py``). Those copies +are **shallow**: ``Content`` and every ``Part`` are copied, but nested payloads +(``function_call.args``, ``function_response.response``, ``inline_data.data``) +are shared with the session events. This module may therefore only assign +``Content.parts`` or append new ``Part`` objects. Mutating a nested field in +place would corrupt persisted history. Pairing is checked positionally, against the immediately following content, -because that is the invariant the provider enforces. A response that exists -somewhere else in the history does not satisfy it. +because that is the invariant the provider enforces. On a conversation whose calls are all answered this is a no-op, and any conversation it does change is one the provider would have rejected outright. @@ -25,49 +31,91 @@ from __future__ import annotations +import logging + from google.adk.agents.callback_context import CallbackContext from google.adk.models.llm_request import LlmRequest from google.genai import types -# Stands in for a result that was never recorded. Deliberately states no cause: -# the call may have been interrupted, or may belong to a long-running tool that -# has not returned yet. Matches the placeholder the OpenAI and Go converters +logger = logging.getLogger(__name__) + +# Stands in for a result that was never recorded. States no cause, because the +# call may have been interrupted or the process may simply have died between the +# two events. Matches the placeholder the Anthropic and OpenAI converters # already use, so providers that repair on their own keep the same behaviour. MISSING_TOOL_RESULT = "No response available for this function call." +# Stands in for a call that ADK is deliberately holding open: a long-running +# tool, a human approval, or ask_user. ADK creates no function_response for +# these until the answer arrives, so the call is pending rather than lost. The +# distinction matters: told a tool "returned nothing", a model reissues it or +# proceeds; told the call is still awaiting a response, it can wait. +PENDING_TOOL_RESULT = "This call is awaiting a response and has not completed yet." -def _call_ids(content: types.Content | None) -> set[str | None]: - if content is None: + +def _pending_call_ids(callback_context: CallbackContext | None) -> set[str]: + """Ids of calls ADK is holding open for a long-running tool or an approval. + + Read from the session events rather than the request, because + ``long_running_tool_ids`` lives on the event and does not survive the + conversion to contents. + """ + session = getattr(callback_context, "session", None) + if session is None: return set() - return {p.function_call.id for p in content.parts or [] if p.function_call} + pending: set[str] = set() + for event in session.events or []: + if event.long_running_tool_ids: + pending.update(event.long_running_tool_ids) + return pending -def _response_ids(content: types.Content | None) -> set[str | None]: +def _response_ids(content: types.Content | None) -> list[str | None]: if content is None: - return set() - return {p.function_response.id for p in content.parts or [] if p.function_response} + return [] + return [p.function_response.id for p in content.parts or [] if p.function_response] -def _drop_orphaned_responses(contents: list[types.Content]) -> list[types.Content]: - """Remove function_response parts whose call is not in the preceding content.""" - kept: list[types.Content] = [] - for index, content in enumerate(contents): - responses = _response_ids(content) - if not responses: - kept.append(content) - continue +def _placeholder_for(call: types.FunctionCall, pending_ids: set[str]) -> str: + if call.id and call.id in pending_ids: + return PENDING_TOOL_RESULT + return MISSING_TOOL_RESULT + - answerable = _call_ids(contents[index - 1]) if index > 0 else set() - parts = [p for p in content.parts or [] if not p.function_response or p.function_response.id in answerable] - if not parts: +def _unanswered(calls: list[types.FunctionCall], answered: list[str | None]) -> list[types.FunctionCall]: + """Calls with no response in ``answered``. + + Ids are matched by consuming them one at a time rather than through a set, + so a turn carrying several calls with no id (Gemini omits them, and ADK + strips its own ``adk-`` ids before the request is built) does not have one + response silently answer all of them. + """ + remaining = list(answered) + missing = [] + for call in calls: + if call.id in remaining: + remaining.remove(call.id) continue - content.parts = parts - kept.append(content) - return kept + missing.append(call) + return missing + + +def repair_tool_call_pairing_callback( + callback_context: CallbackContext, + llm_request: LlmRequest, +) -> None: + """Before-model callback that pairs every tool call with a tool result. + Supplies a placeholder result for any call the immediately following content + does not answer, so the request satisfies the strict call/result pairing that + Anthropic and Bedrock require. + """ + contents = llm_request.contents + if not contents: + return None -def _synthesize_missing_responses(contents: list[types.Content]) -> list[types.Content]: - """Give every function_call a function_response in the immediately following content.""" + pending_ids = _pending_call_ids(callback_context) + synthesized = 0 repaired: list[types.Content] = [] for index, content in enumerate(contents): repaired.append(content) @@ -78,16 +126,17 @@ def _synthesize_missing_responses(contents: list[types.Content]) -> list[types.C following = contents[index + 1] if index + 1 < len(contents) else None answered = _response_ids(following) - missing = [call for call in calls if call.id not in answered] + missing = _unanswered(calls, answered) if not missing: continue + synthesized += len(missing) parts = [ types.Part( function_response=types.FunctionResponse( id=call.id, name=call.name, - response={"result": MISSING_TOOL_RESULT}, + response={"result": _placeholder_for(call, pending_ids)}, ) ) for call in missing @@ -100,21 +149,10 @@ def _synthesize_missing_responses(contents: list[types.Content]) -> list[types.C following.parts = list(following.parts or []) + parts else: repaired.append(types.Content(role="user", parts=parts)) - return repaired - -def repair_tool_call_pairing_callback( - callback_context: CallbackContext, - llm_request: LlmRequest, -) -> None: - """Before-model callback that pairs every tool call with a tool result. - - Drops a result whose call is gone, then supplies a placeholder result for a - call that has none, so the request satisfies the strict call/result pairing - that Anthropic (and Bedrock) require. - """ - if not llm_request.contents: - return None - contents = _drop_orphaned_responses(list(llm_request.contents)) - llm_request.contents = _synthesize_missing_responses(contents) + llm_request.contents = repaired + if synthesized: + # Logged because a repaired request means a turn ended between a tool + # call and its result somewhere upstream. + logger.info("Paired %d unanswered tool call(s) before the model request", synthesized) return None diff --git a/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py b/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py index 1bc7f3b42..0943fbd2b 100644 --- a/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py +++ b/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py @@ -1,15 +1,30 @@ """Tests for tool call/response pairing repair.""" +import pytest +from google.adk.events.event import Event from google.adk.models.anthropic_llm import content_to_message_param from google.adk.models.llm_request import LlmRequest from google.genai import types from kagent.adk._tool_pairing import ( MISSING_TOOL_RESULT, + PENDING_TOOL_RESULT, repair_tool_call_pairing_callback, ) +class _Session: + def __init__(self, events): + self.events = events + + +class _Context: + """Stands in for ADK's CallbackContext, which exposes `session`.""" + + def __init__(self, events=None): + self.session = _Session(events or []) + + def _call(call_id: str, name: str = "get_pods") -> types.Content: return types.Content( role="model", @@ -17,6 +32,13 @@ def _call(call_id: str, name: str = "get_pods") -> types.Content: ) +def _calls(*ids: str) -> types.Content: + return types.Content( + role="model", + parts=[types.Part(function_call=types.FunctionCall(id=i, name="get_pods", args={})) for i in ids], + ) + + def _response(call_id: str, name: str = "get_pods", text: str = "pod X running") -> types.Content: return types.Content( role="user", @@ -28,9 +50,22 @@ def _text(role: str, text: str) -> types.Content: return types.Content(role=role, parts=[types.Part(text=text)]) -def _repair(contents: list[types.Content]) -> list[types.Content]: +def _pending_event(*call_ids: str) -> Event: + """An event whose tool calls ADK is holding open (approval, ask_user, long-running).""" + return Event( + author="agent", + invocation_id="i", + content=types.Content( + role="model", + parts=[types.Part(function_call=types.FunctionCall(id=i, name="approve", args={})) for i in call_ids], + ), + long_running_tool_ids=set(call_ids), + ) + + +def _repair(contents: list[types.Content], context: _Context | None = None) -> list[types.Content]: request = LlmRequest(contents=contents) - repair_tool_call_pairing_callback(callback_context=None, llm_request=request) + repair_tool_call_pairing_callback(callback_context=context or _Context(), llm_request=request) return request.contents @@ -52,7 +87,7 @@ def test_dangling_call_at_end_of_history(self): assert contents[2].role == "user" def test_result_inserted_before_a_following_user_message(self): - """The double-message race: a second message arrives before the result.""" + """A second message reached the history before the result did.""" contents = _repair([_call("abc"), _text("user", "second message")]) assert len(contents) == 3 @@ -61,14 +96,7 @@ def test_result_inserted_before_a_following_user_message(self): def test_sibling_response_turn_is_reused(self): """A partially answered call turn gets its gap filled, not a new turn.""" - call_turn = types.Content( - role="model", - parts=[ - types.Part(function_call=types.FunctionCall(id="abc", name="get_pods", args={})), - types.Part(function_call=types.FunctionCall(id="def", name="get_logs", args={})), - ], - ) - contents = _repair([call_turn, _response("abc")]) + contents = _repair([_calls("abc", "def"), _response("abc")]) assert len(contents) == 2 responses = _responses_in(contents[1]) @@ -77,14 +105,7 @@ def test_sibling_response_turn_is_reused(self): assert responses[1].response == {"result": MISSING_TOOL_RESULT} def test_only_the_unanswered_call_is_synthesized(self): - call_turn = types.Content( - role="model", - parts=[ - types.Part(function_call=types.FunctionCall(id="abc", name="get_pods", args={})), - types.Part(function_call=types.FunctionCall(id="def", name="get_logs", args={})), - ], - ) - contents = _repair([call_turn, _response("def", name="get_logs", text="log line")]) + contents = _repair([_calls("abc", "def"), _response("def", text="log line")]) responses = _responses_in(contents[1]) assert {r.id for r in responses} == {"abc", "def"} @@ -110,46 +131,59 @@ def test_call_without_id_still_gets_a_response(self): assert len(contents) == 2 assert _responses_in(contents[1])[0].id is None + def test_two_calls_without_ids_each_get_their_own_response(self): + """One response must not answer every id-less call in the turn. -class TestDropOrphanedResponses: - def test_response_without_a_call_is_dropped(self): - contents = _repair([_text("user", "hello"), _response("abc")]) - - assert len(contents) == 1 - assert contents[0].parts[0].text == "hello" - - def test_only_the_orphan_is_dropped(self): - response_turn = types.Content( - role="user", + Gemini omits call ids, and ADK strips its own `adk-` ids before the + request is built, so id-less siblings are not hypothetical. + """ + call_turn = types.Content( + role="model", parts=[ - types.Part( - function_response=types.FunctionResponse(id="abc", name="get_pods", response={"result": "ok"}) - ), - types.Part( - function_response=types.FunctionResponse(id="zzz", name="ghost", response={"result": "stale"}) - ), + types.Part(function_call=types.FunctionCall(name="get_pods", args={})), + types.Part(function_call=types.FunctionCall(name="get_logs", args={})), ], ) - contents = _repair([_call("abc"), response_turn]) + answered = types.Content( + role="user", + parts=[types.Part(function_response=types.FunctionResponse(name="get_pods", response={"result": "ok"}))], + ) + + contents = _repair([call_turn, answered]) responses = _responses_in(contents[1]) - assert [r.id for r in responses] == ["abc"] + assert len(responses) == 2 + assert [r.response for r in responses] == [{"result": "ok"}, {"result": MISSING_TOOL_RESULT}] - def test_surrounding_parts_survive_an_orphan(self): - response_turn = types.Content( - role="user", - parts=[ - types.Part( - function_response=types.FunctionResponse(id="zzz", name="ghost", response={"result": "stale"}) - ), - types.Part(text="and here is my question"), - ], - ) - contents = _repair([_text("user", "hello"), response_turn]) - assert len(contents) == 2 - assert _responses_in(contents[1]) == [] - assert contents[1].parts[0].text == "and here is my question" +class TestPendingLongRunningCalls: + """Calls ADK holds open on purpose must not be reported as empty results.""" + + def test_pending_approval_gets_the_pending_placeholder(self): + """A plain message sent while an approval is pending reaches this callback. + + `_translate_hitl_response` only rewrites messages carrying a HITL payload, + so an ordinary chat message during a pending confirmation arrives as a + normal turn with the approval call still unanswered. + """ + context = _Context([_pending_event("call-1")]) + + contents = _repair([_call("call-1", name="approve"), _text("user", "never mind, what about X?")], context) + + assert _responses_in(contents[1])[0].response == {"result": PENDING_TOOL_RESULT} + + def test_interrupted_call_in_the_same_session_keeps_the_missing_placeholder(self): + context = _Context([_pending_event("call-1")]) + + contents = _repair([_call("other"), _text("user", "still there?")], context) + + assert _responses_in(contents[1])[0].response == {"result": MISSING_TOOL_RESULT} + + def test_missing_context_falls_back_to_the_neutral_placeholder(self): + request = LlmRequest(contents=[_call("abc")]) + repair_tool_call_pairing_callback(callback_context=None, llm_request=request) + + assert _responses_in(request.contents[1])[0].response == {"result": MISSING_TOOL_RESULT} class TestDoesNotDisturbHealthyHistory: @@ -167,9 +201,17 @@ def test_history_without_tools_is_unchanged(self): assert [c.model_dump_json() for c in _repair(original)] == before + def test_answered_pending_call_is_left_alone(self): + """Once the human answers, the real response is present and wins.""" + context = _Context([_pending_event("call-1")]) + original = [_call("call-1", name="approve"), _response("call-1", name="approve", text="approved")] + before = [c.model_dump_json() for c in original] + + assert [c.model_dump_json() for c in _repair(original, context)] == before + def test_empty_and_missing_parts_do_not_raise(self): request = LlmRequest(contents=[]) - repair_tool_call_pairing_callback(callback_context=None, llm_request=request) + repair_tool_call_pairing_callback(callback_context=_Context(), llm_request=request) assert request.contents == [] contents = _repair([types.Content(role="user", parts=None), _call("abc")]) @@ -194,22 +236,17 @@ def _assert_paired(contents: list[types.Content]) -> None: def test_interrupted_turn_produces_a_valid_conversation(self): self._assert_paired(_repair([_text("user", "what's failing?"), _call("abc")])) - def test_double_message_race_produces_a_valid_conversation(self): + def test_message_during_a_running_tool_produces_a_valid_conversation(self): self._assert_paired(_repair([_call("abc"), _text("user", "second message")])) def test_partially_answered_turn_produces_a_valid_conversation(self): - call_turn = types.Content( - role="model", - parts=[ - types.Part(function_call=types.FunctionCall(id="abc", name="get_pods", args={})), - types.Part(function_call=types.FunctionCall(id="def", name="get_logs", args={})), - ], - ) - self._assert_paired(_repair([call_turn, _response("abc")])) + self._assert_paired(_repair([_calls("abc", "def"), _response("abc")])) + + def test_pending_approval_produces_a_valid_conversation(self): + context = _Context([_pending_event("call-1")]) + self._assert_paired(_repair([_call("call-1", name="approve"), _text("user", "what about X?")], context)) def test_unrepaired_history_would_fail_the_invariant(self): """Guards the tests themselves: the invariant must reject the broken input.""" - import pytest - with pytest.raises(AssertionError): self._assert_paired([_text("user", "what's failing?"), _call("abc")])