diff --git a/go/adk/pkg/agent/agent.go b/go/adk/pkg/agent/agent.go index 1becb1c00..96dd17125 100644 --- a/go/adk/pkg/agent/agent.go +++ b/go/adk/pkg/agent/agent.go @@ -148,6 +148,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(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 new file mode 100644 index 000000000..e771cdf78 --- /dev/null +++ b/go/adk/pkg/agent/tool_pairing.go @@ -0,0 +1,208 @@ +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 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), 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. +// +// Pairing is checked positionally, against the immediately following content, +// 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(log logr.Logger) llmagent.BeforeModelCallback { + return func(ctx agent.Context, req *adkmodel.LLMRequest) (*adkmodel.LLMResponse, error) { + if len(req.Contents) == 0 { + return nil, nil + } + 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 + } +} + +// sessionSource is the part of agent.Context this file needs. +type sessionSource interface { + Session() session.Session +} + +// 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 + } + for _, id := range event.LongRunningToolIDs { + pending[id] = true + } + } + return pending +} + +// synthesizeMissingResponses gives every function call a function response in +// 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) + if len(calls) == 0 { + continue + } + + var following *genai.Content + if index+1 < len(contents) { + following = contents[index+1] + } + answered := responseIDs(following) + missing := unanswered(calls, answered) + if len(missing) == 0 { + 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": placeholderFor(call, pending)}, + }, + }) + } + + // 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: genai.RoleUser, Parts: parts}) + } + 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 { + 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 responseIDs(content *genai.Content) []string { + if content == nil { + return nil + } + var ids []string + for _, part := range content.Parts { + if part != nil && part.FunctionResponse != nil { + ids = append(ids, part.FunctionResponse.ID) + } + } + return ids +} 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..989ba1693 --- /dev/null +++ b/go/adk/pkg/agent/tool_pairing_test.go @@ -0,0 +1,344 @@ +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" +) + +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: genai.RoleUser, 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}}} +} + +// 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 + } + } + } +} +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 { + var responses []*genai.FunctionResponse + for _, part := range content.Parts { + if part.FunctionResponse != nil { + responses = append(responses, part.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(genai.RoleUser, "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(genai.RoleUser, "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(genai.RoleUser, "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) + } +} + +// 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() + + 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}) + + responses := responsesIn(contents[1]) + if len(responses) != 2 { + t.Fatalf("expected both id-less calls to be answered, got %d responses", len(responses)) + } + 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]) + } +} + +// 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() + + 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 := repairWith( + []*genai.Content{callContent("other"), textContent(genai.RoleUser, "still there?")}, + withPending("call-1"), + ) + + responses := responsesIn(contents[1]) + if len(responses) != 1 || responses[0].Response["result"] != missingToolResult { + t.Fatalf("expected the neutral placeholder, got %+v", responses) + } +} + +func TestToolPairingLeavesHealthyHistoryUntouched(t *testing.T) { + t.Parallel() + + original := []*genai.Content{ + textContent(genai.RoleUser, "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 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() + + 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)) + } +} + +// 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(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) { + 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 new file mode 100644 index 000000000..467119978 --- /dev/null +++ b/python/packages/kagent-adk/src/kagent/adk/_tool_pairing.py @@ -0,0 +1,158 @@ +"""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), +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. + +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 + +import logging + +from google.adk.agents.callback_context import CallbackContext +from google.adk.models.llm_request import LlmRequest +from google.genai import types + +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 _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() + 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) -> list[str | None]: + if content is None: + return [] + return [p.function_response.id for p in content.parts or [] if p.function_response] + + +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 + + +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 + 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 + + pending_ids = _pending_call_ids(callback_context) + synthesized = 0 + 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 = _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": _placeholder_for(call, pending_ids)}, + ) + ) + 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)) + + 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/src/kagent/adk/types.py b/python/packages/kagent-adk/src/kagent/adk/types.py index 9d4a25b28..813d922c6 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 @@ -536,7 +537,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..0943fbd2b --- /dev/null +++ b/python/packages/kagent-adk/tests/unittests/test_tool_pairing.py @@ -0,0 +1,252 @@ +"""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", + parts=[types.Part(function_call=types.FunctionCall(id=call_id, name=name, args={"ns": "default"}))], + ) + + +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", + 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 _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=context or _Context(), 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): + """A second message reached the history before the result did.""" + 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.""" + contents = _repair([_calls("abc", "def"), _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): + contents = _repair([_calls("abc", "def"), _response("def", 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 + + def test_two_calls_without_ids_each_get_their_own_response(self): + """One response must not answer every id-less call in the turn. + + 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_call=types.FunctionCall(name="get_pods", args={})), + types.Part(function_call=types.FunctionCall(name="get_logs", args={})), + ], + ) + 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 len(responses) == 2 + assert [r.response for r in responses] == [{"result": "ok"}, {"result": MISSING_TOOL_RESULT}] + + +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: + 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_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=_Context(), 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_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): + 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.""" + with pytest.raises(AssertionError): + self._assert_paired([_text("user", "what's failing?"), _call("abc")])