From 6bae658b0592aa936e1b48e96ff9f995593ba086 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Mon, 13 Jul 2026 07:28:48 -0700 Subject: [PATCH] fix: correctly reassemble streamed function-call arguments in Gemini streaming Gemini can keep willContinue=true on the last partialArgs chunk and end a function call with a separate empty part. Flush the buffered call on that end marker so streamed calls are not dropped or merged. PiperOrigin-RevId: 947029146 --- .../google/adk/flows/llmflows/Functions.java | 6 +- .../java/com/google/adk/models/Gemini.java | 20 +- .../com/google/adk/models/GeminiTest.java | 195 ++++++++++++++++++ .../com/google/adk/runner/RunnerTest.java | 98 +++++++++ 4 files changed, 311 insertions(+), 8 deletions(-) diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index f531e214d..0653052f2 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -415,10 +415,12 @@ public static Set getLongRunningFunctionCalls( List functionCalls, Map tools) { Set longRunningFunctionCalls = new HashSet<>(); for (FunctionCall functionCall : functionCalls) { - if (!tools.containsKey(functionCall.name().get())) { + // Streamed function-call chunks may carry no name; skip them. + String name = functionCall.name().orElse(null); + if (name == null || !tools.containsKey(name)) { continue; } - BaseTool tool = tools.get(functionCall.name().get()); + BaseTool tool = tools.get(name); if (tool != null && tool.longRunning()) { longRunningFunctionCalls.add(functionCall.id().orElse("")); } diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index 61e87e66f..e13305540 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -416,19 +416,27 @@ private boolean accumulateParts(List parts) { */ private void processFunctionCallPart(Part part) { FunctionCall fc = part.functionCall().get(); - boolean streaming = + boolean hasName = fc.name().filter(name -> !name.isEmpty()).isPresent(); + // A streamed call: it has partialArgs or willContinue, or is the nameless terminal + // marker of an in-progress call. Gemini may end a call with a separate empty + // willContinue=false part, so that marker completes it. + boolean streamedPart = fc.partialArgs().map(args -> !args.isEmpty()).orElse(false) - || fc.willContinue().orElse(false); - if (streaming) { + || fc.willContinue().orElse(false) + || (currentFcName != null && !hasName); + if (streamedPart) { // Capture the thought signature from the first chunk that carries one. if (part.thoughtSignature().isPresent() && currentThoughtSignature == null) { currentThoughtSignature = part.thoughtSignature().get(); } processStreamingFunctionCall(fc); - } else if (fc.name().filter(name -> !name.isEmpty()).isPresent()) { - // Complete function call. Skip empty calls, which are only streaming end markers. The part - // already has an ID assigned by ensureFunctionCallIds. + } else if (hasName) { + // Complete (non-streamed) call. Safety guard: the model should terminate a streamed call + // with willContinue=false before starting a new one; flush any still-in-progress call so it + // is neither dropped nor merged. The part already has an ID assigned by + // ensureFunctionCallIds. flushTextBufferToSequence(); + flushFunctionCallToSequence(); accumulatedSequence.add(part); } } diff --git a/core/src/test/java/com/google/adk/models/GeminiTest.java b/core/src/test/java/com/google/adk/models/GeminiTest.java index 9ecdcec28..a965e5b68 100644 --- a/core/src/test/java/com/google/adk/models/GeminiTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiTest.java @@ -371,6 +371,201 @@ public void processRawResponses_twoStreamingFunctionCalls_keepArgsSeparate() { assertThat(second.args().get()).containsExactly("b", "2"); } + // The last partialArgs chunk keeps willContinue=true; completion arrives on a separate empty + // willContinue=false marker, then trailing text follows. The marker must flush the call so it + // precedes the text. Without handling the marker, close() flushes the call after the text, + // reversing their order (a single call alone would be masked by that end-of-stream flush). + @Test + public void processRawResponses_streamedCallEndedByEmptyMarker_flushesCallBeforeTrailingText() { + GenerateContentResponse name = + toResponse( + functionCallPart(FunctionCall.builder().name("bookFlight").willContinue(true).build())); + GenerateContentResponse origin1 = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder().jsonPath("$.origin").stringValue("Krak").build()) + .willContinue(true) + .build())); + GenerateContentResponse origin2 = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder().jsonPath("$.origin").stringValue("ow").build()) + .willContinue(true) + .build())); + GenerateContentResponse destination = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder() + .jsonPath("$.destination") + .stringValue("Warsaw") + .build()) + .willContinue(true) + .build())); + GenerateContentResponse endMarker = + toResponse(functionCallPart(FunctionCall.builder().willContinue(false).build())); + GenerateContentResponse trailingText = toResponseWithText("Booked.", FinishReason.Known.STOP); + + ImmutableList responses = + ImmutableList.copyOf( + Gemini.processRawResponses( + Flowable.just(name, origin1, origin2, destination, endMarker, trailingText)) + .blockingIterable()); + + LlmResponse finalResponse = Iterables.getLast(responses); + assertThat(finalResponse.content().get().parts().get()).hasSize(2); + FunctionCall finalCall = + finalResponse.content().get().parts().get().get(0).functionCall().get(); + assertThat(finalCall.name()).hasValue("bookFlight"); + assertThat(finalCall.args().get()).containsExactly("origin", "Krakow", "destination", "Warsaw"); + assertThat(finalResponse.content().get().parts().get().get(1).text()).hasValue("Booked."); + } + + // Two multi-arg streamed calls each ended by an empty willContinue=false marker must not drop the + // first call nor bleed its args into the second. + @Test + public void processRawResponses_twoStreamedCallsEndedByEmptyMarkers_keepArgsSeparate() { + GenerateContentResponse call1Name = + toResponse( + functionCallPart( + FunctionCall.builder().name("getTemperature").willContinue(true).build())); + GenerateContentResponse call1City = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder().jsonPath("$.city").stringValue("Krakow").build()) + .willContinue(true) + .build())); + GenerateContentResponse call1Unit = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.unit").stringValue("C").build()) + .willContinue(true) + .build())); + GenerateContentResponse marker1 = + toResponse(functionCallPart(FunctionCall.builder().willContinue(false).build())); + GenerateContentResponse call2Name = + toResponse( + functionCallPart( + FunctionCall.builder().name("getCondition").willContinue(true).build())); + GenerateContentResponse call2City = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder().jsonPath("$.city").stringValue("Warsaw").build()) + .willContinue(true) + .build())); + GenerateContentResponse call2Unit = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.unit").stringValue("F").build()) + .willContinue(true) + .build())); + GenerateContentResponse marker2 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts(functionCallPart(FunctionCall.builder().willContinue(false).build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + ImmutableList responses = + ImmutableList.copyOf( + Gemini.processRawResponses( + Flowable.just( + call1Name, call1City, call1Unit, marker1, call2Name, call2City, call2Unit, + marker2)) + .blockingIterable()); + + LlmResponse finalResponse = Iterables.getLast(responses); + assertThat(finalResponse.content().get().parts().get()).hasSize(2); + FunctionCall first = finalResponse.content().get().parts().get().get(0).functionCall().get(); + FunctionCall second = finalResponse.content().get().parts().get().get(1).functionCall().get(); + assertThat(first.name()).hasValue("getTemperature"); + assertThat(first.args().get()).containsExactly("city", "Krakow", "unit", "C"); + assertThat(second.name()).hasValue("getCondition"); + assertThat(second.args().get()).containsExactly("city", "Warsaw", "unit", "F"); + } + + // Safety guard for non-conforming output: a streamed call still in progress (the model should + // have terminated it with willContinue=false) is followed by a complete non-streaming call. The + // in-progress call is flushed before appending, so neither is dropped nor merged. + @Test + public void processRawResponses_streamedCallFollowedByCompleteCall_flushesInProgressFirst() { + GenerateContentResponse streamedName = + toResponse( + functionCallPart( + FunctionCall.builder().name("stream_call").willContinue(true).build())); + GenerateContentResponse streamedArg = + toResponse( + functionCallPart( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.a").stringValue("1").build()) + .willContinue(true) + .build())); + GenerateContentResponse completeCall = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts( + functionCallPart( + FunctionCall.builder() + .name("plain_call") + .args(ImmutableMap.of("b", "2")) + .build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + ImmutableList responses = + ImmutableList.copyOf( + Gemini.processRawResponses(Flowable.just(streamedName, streamedArg, completeCall)) + .blockingIterable()); + + LlmResponse finalResponse = Iterables.getLast(responses); + assertThat(finalResponse.content().get().parts().get()).hasSize(2); + FunctionCall first = finalResponse.content().get().parts().get().get(0).functionCall().get(); + FunctionCall second = finalResponse.content().get().parts().get().get(1).functionCall().get(); + assertThat(first.name()).hasValue("stream_call"); + assertThat(first.args().get()).containsExactly("a", "1"); + assertThat(second.name()).hasValue("plain_call"); + assertThat(second.args().get()).containsExactly("b", "2"); + } + + // A stray nameless willContinue=false marker with no call in progress must be a safe no-op (the + // currentFcName != null half of the guard): it must not add a function call nor split the + // surrounding text. Without that half it would be treated as a streamed part and prematurely + // flush the text buffer, splitting "Hello world" into two parts. + @Test + public void processRawResponses_strayNamelessMarker_isNoOpAndDoesNotSplitText() { + GenerateContentResponse hello = toResponseWithText("Hello "); + GenerateContentResponse strayMarker = + toResponse(functionCallPart(FunctionCall.builder().willContinue(false).build())); + GenerateContentResponse world = toResponseWithText("world", FinishReason.Known.STOP); + + ImmutableList responses = + ImmutableList.copyOf( + Gemini.processRawResponses(Flowable.just(hello, strayMarker, world)) + .blockingIterable()); + + LlmResponse finalResponse = Iterables.getLast(responses); + assertThat(finalResponse.content().get().parts().get()).hasSize(1); + assertThat(finalResponse.content().get().parts().get().get(0).text()).hasValue("Hello world"); + assertThat(finalResponse.content().get().parts().get().get(0).functionCall()).isEmpty(); + } + @Test public void processRawResponses_imageOnlyWithStop_emitsFinalImagePart() { Part imagePart = Part.fromBytes(new byte[] {1, 2, 3}, "image/png"); diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 42afda4a0..38485a5a7 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -79,6 +79,7 @@ import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; +import com.google.genai.types.PartialArg; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Context; import io.opentelemetry.context.ContextKey; @@ -1378,6 +1379,103 @@ public void runAsync_withSessionKey_success() { assertThat(simplifyEvents(events)).containsExactly("test agent: from llm"); } + // Runner-level regression for streamed function-call arguments: a multi-arg call whose args + // arrive + // across partial events (nameless continuation chunks, with one value split across chunks) must + // not crash and must execute the tool exactly once with the reassembled args. + @Test + public void runAsync_streamedFunctionCallArgs_reassembledAndToolExecuted() { + // Turn 1: SSE stream mimicking the post-aggregator shape - a named chunk with willContinue, + // then + // nameless continuation chunks carrying partialArgs (origin split across two), then the + // aggregated complete call. + LlmResponse namedChunk = + partialFcResponse( + FunctionCall.builder().id("fc-1").name(echoTool.name()).willContinue(true).build()); + LlmResponse originChunk1 = + partialFcResponse( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.origin").stringValue("Krak").build()) + .willContinue(true) + .build()); + LlmResponse originChunk2 = + partialFcResponse( + FunctionCall.builder() + .partialArgs(PartialArg.builder().jsonPath("$.origin").stringValue("ow").build()) + .willContinue(true) + .build()); + LlmResponse destinationChunk = + partialFcResponse( + FunctionCall.builder() + .partialArgs( + PartialArg.builder().jsonPath("$.destination").stringValue("Warsaw").build()) + .willContinue(true) + .build()); + LlmResponse aggregatedCall = + createFunctionCallLlmResponse( + "fc-1", echoTool.name(), ImmutableMap.of("origin", "Krakow", "destination", "Warsaw")); + + TestLlm streamingLlm = + createTestLlm( + Flowable.just(namedChunk, originChunk1, originChunk2, destinationChunk, aggregatedCall), + Flowable.just(createTextLlmResponse("done"))); + LlmAgent streamingAgent = + createTestAgentBuilder(streamingLlm).tools(ImmutableList.of(new EchoTool())).build(); + Runner streamingRunner = + Runner.builder().app(App.builder().name("test").rootAgent(streamingAgent).build()).build(); + Session streamingSession = + streamingRunner.sessionService().createSession("test", "user").blockingGet(); + + List events = + streamingRunner + .runAsync( + "user", + streamingSession.id(), + createContent("book a flight"), + RunConfig.builder().setStreamingMode(RunConfig.StreamingMode.SSE).build()) + .toList() + .blockingGet(); + + int toolResponses = 0; + boolean sawPartialFcChunk = false; + boolean sawFinalText = false; + Map executedArgs = null; + for (Event e : events) { + toolResponses += e.functionResponses().size(); + if (e.partial().orElse(false) && !e.functionCalls().isEmpty()) { + sawPartialFcChunk = true; + } + if (!e.partial().orElse(false) && !e.functionCalls().isEmpty()) { + executedArgs = e.functionCalls().get(0).args().orElse(ImmutableMap.of()); + } + boolean hasDone = + e.content() + .flatMap(Content::parts) + .map( + parts -> + parts.stream() + .anyMatch(p -> p.text().map(t -> t.contains("done")).orElse(false))) + .orElse(false); + sawFinalText |= hasDone; + } + + // The streamed (incl. nameless) chunks flowed through the runner without crashing; the tool ran + // exactly once (only the aggregated non-partial call triggers execution) with both reassembled + // args; and the final text was produced. + assertThat(sawPartialFcChunk).isTrue(); + assertThat(toolResponses).isEqualTo(1); + assertThat(executedArgs).containsExactly("origin", "Krakow", "destination", "Warsaw"); + assertThat(sawFinalText).isTrue(); + } + + private static LlmResponse partialFcResponse(FunctionCall fc) { + return LlmResponse.builder() + .content( + Content.builder().role("model").parts(Part.builder().functionCall(fc).build()).build()) + .partial(true) + .build(); + } + @Test public void runAsync_withStateDelta_mergesStateIntoSession() { ImmutableMap stateDelta = ImmutableMap.of("key1", "value1", "key2", 42);