Skip to content

Commit 45c306c

Browse files
damianmomotgooglecopybara-github
damianmomotgoogle
authored and
copybara-github
committed
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: 946860751
1 parent 410ff81 commit 45c306c

4 files changed

Lines changed: 311 additions & 8 deletions

File tree

‎core/src/main/java/com/google/adk/flows/llmflows/Functions.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -415,10 +415,12 @@ public static Set<String> getLongRunningFunctionCalls(
415415
List<FunctionCall> functionCalls, Map<String, BaseTool> tools) {
416416
Set<String> longRunningFunctionCalls = new HashSet<>();
417417
for (FunctionCall functionCall : functionCalls) {
418-
if (!tools.containsKey(functionCall.name().get())) {
418+
// Streamed function-call chunks may carry no name; skip them.
419+
String name = functionCall.name().orElse(null);
420+
if (name == null || !tools.containsKey(name)) {
419421
continue;
420422
}
421-
BaseTool tool = tools.get(functionCall.name().get());
423+
BaseTool tool = tools.get(name);
422424
if (tool != null && tool.longRunning()) {
423425
longRunningFunctionCalls.add(functionCall.id().orElse(""));
424426
}

‎core/src/main/java/com/google/adk/models/Gemini.java‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -416,19 +416,27 @@ private boolean accumulateParts(List<Part> parts) {
416416
*/
417417
private void processFunctionCallPart(Part part) {
418418
FunctionCall fc = part.functionCall().get();
419-
boolean streaming =
419+
boolean hasName = fc.name().filter(name -> !name.isEmpty()).isPresent();
420+
// A streamed call: it has partialArgs or willContinue, or is the nameless terminal
421+
// marker of an in-progress call. Gemini may end a call with a separate empty
422+
// willContinue=false part, so that marker completes it.
423+
boolean streamedPart =
420424
fc.partialArgs().map(args -> !args.isEmpty()).orElse(false)
421-
|| fc.willContinue().orElse(false);
422-
if (streaming) {
425+
|| fc.willContinue().orElse(false)
426+
|| (currentFcName != null && !hasName);
427+
if (streamedPart) {
423428
// Capture the thought signature from the first chunk that carries one.
424429
if (part.thoughtSignature().isPresent() && currentThoughtSignature == null) {
425430
currentThoughtSignature = part.thoughtSignature().get();
426431
}
427432
processStreamingFunctionCall(fc);
428-
} else if (fc.name().filter(name -> !name.isEmpty()).isPresent()) {
429-
// Complete function call. Skip empty calls, which are only streaming end markers. The part
430-
// already has an ID assigned by ensureFunctionCallIds.
433+
} else if (hasName) {
434+
// Complete (non-streamed) call. Safety guard: the model should terminate a streamed call
435+
// with willContinue=false before starting a new one; flush any still-in-progress call so it
436+
// is neither dropped nor merged. The part already has an ID assigned by
437+
// ensureFunctionCallIds.
431438
flushTextBufferToSequence();
439+
flushFunctionCallToSequence();
432440
accumulatedSequence.add(part);
433441
}
434442
}

‎core/src/test/java/com/google/adk/models/GeminiTest.java‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -371,6 +371,201 @@ public void processRawResponses_twoStreamingFunctionCalls_keepArgsSeparate() {
371371
assertThat(second.args().get()).containsExactly("b", "2");
372372
}
373373

374+
// The last partialArgs chunk keeps willContinue=true; completion arrives on a separate empty
375+
// willContinue=false marker, then trailing text follows. The marker must flush the call so it
376+
// precedes the text. Without handling the marker, close() flushes the call after the text,
377+
// reversing their order (a single call alone would be masked by that end-of-stream flush).
378+
@Test
379+
public void processRawResponses_streamedCallEndedByEmptyMarker_flushesCallBeforeTrailingText() {
380+
GenerateContentResponse name =
381+
toResponse(
382+
functionCallPart(FunctionCall.builder().name("bookFlight").willContinue(true).build()));
383+
GenerateContentResponse origin1 =
384+
toResponse(
385+
functionCallPart(
386+
FunctionCall.builder()
387+
.partialArgs(
388+
PartialArg.builder().jsonPath("$.origin").stringValue("Krak").build())
389+
.willContinue(true)
390+
.build()));
391+
GenerateContentResponse origin2 =
392+
toResponse(
393+
functionCallPart(
394+
FunctionCall.builder()
395+
.partialArgs(
396+
PartialArg.builder().jsonPath("$.origin").stringValue("ow").build())
397+
.willContinue(true)
398+
.build()));
399+
GenerateContentResponse destination =
400+
toResponse(
401+
functionCallPart(
402+
FunctionCall.builder()
403+
.partialArgs(
404+
PartialArg.builder()
405+
.jsonPath("$.destination")
406+
.stringValue("Warsaw")
407+
.build())
408+
.willContinue(true)
409+
.build()));
410+
GenerateContentResponse endMarker =
411+
toResponse(functionCallPart(FunctionCall.builder().willContinue(false).build()));
412+
GenerateContentResponse trailingText = toResponseWithText("Booked.", FinishReason.Known.STOP);
413+
414+
ImmutableList<LlmResponse> responses =
415+
ImmutableList.copyOf(
416+
Gemini.processRawResponses(
417+
Flowable.just(name, origin1, origin2, destination, endMarker, trailingText))
418+
.blockingIterable());
419+
420+
LlmResponse finalResponse = Iterables.getLast(responses);
421+
assertThat(finalResponse.content().get().parts().get()).hasSize(2);
422+
FunctionCall finalCall =
423+
finalResponse.content().get().parts().get().get(0).functionCall().get();
424+
assertThat(finalCall.name()).hasValue("bookFlight");
425+
assertThat(finalCall.args().get()).containsExactly("origin", "Krakow", "destination", "Warsaw");
426+
assertThat(finalResponse.content().get().parts().get().get(1).text()).hasValue("Booked.");
427+
}
428+
429+
// Two multi-arg streamed calls each ended by an empty willContinue=false marker must not drop the
430+
// first call nor bleed its args into the second.
431+
@Test
432+
public void processRawResponses_twoStreamedCallsEndedByEmptyMarkers_keepArgsSeparate() {
433+
GenerateContentResponse call1Name =
434+
toResponse(
435+
functionCallPart(
436+
FunctionCall.builder().name("getTemperature").willContinue(true).build()));
437+
GenerateContentResponse call1City =
438+
toResponse(
439+
functionCallPart(
440+
FunctionCall.builder()
441+
.partialArgs(
442+
PartialArg.builder().jsonPath("$.city").stringValue("Krakow").build())
443+
.willContinue(true)
444+
.build()));
445+
GenerateContentResponse call1Unit =
446+
toResponse(
447+
functionCallPart(
448+
FunctionCall.builder()
449+
.partialArgs(PartialArg.builder().jsonPath("$.unit").stringValue("C").build())
450+
.willContinue(true)
451+
.build()));
452+
GenerateContentResponse marker1 =
453+
toResponse(functionCallPart(FunctionCall.builder().willContinue(false).build()));
454+
GenerateContentResponse call2Name =
455+
toResponse(
456+
functionCallPart(
457+
FunctionCall.builder().name("getCondition").willContinue(true).build()));
458+
GenerateContentResponse call2City =
459+
toResponse(
460+
functionCallPart(
461+
FunctionCall.builder()
462+
.partialArgs(
463+
PartialArg.builder().jsonPath("$.city").stringValue("Warsaw").build())
464+
.willContinue(true)
465+
.build()));
466+
GenerateContentResponse call2Unit =
467+
toResponse(
468+
functionCallPart(
469+
FunctionCall.builder()
470+
.partialArgs(PartialArg.builder().jsonPath("$.unit").stringValue("F").build())
471+
.willContinue(true)
472+
.build()));
473+
GenerateContentResponse marker2 =
474+
toResponse(
475+
Candidate.builder()
476+
.content(
477+
Content.builder()
478+
.parts(functionCallPart(FunctionCall.builder().willContinue(false).build()))
479+
.build())
480+
.finishReason(new FinishReason(FinishReason.Known.STOP))
481+
.build());
482+
483+
ImmutableList<LlmResponse> responses =
484+
ImmutableList.copyOf(
485+
Gemini.processRawResponses(
486+
Flowable.just(
487+
call1Name, call1City, call1Unit, marker1, call2Name, call2City, call2Unit,
488+
marker2))
489+
.blockingIterable());
490+
491+
LlmResponse finalResponse = Iterables.getLast(responses);
492+
assertThat(finalResponse.content().get().parts().get()).hasSize(2);
493+
FunctionCall first = finalResponse.content().get().parts().get().get(0).functionCall().get();
494+
FunctionCall second = finalResponse.content().get().parts().get().get(1).functionCall().get();
495+
assertThat(first.name()).hasValue("getTemperature");
496+
assertThat(first.args().get()).containsExactly("city", "Krakow", "unit", "C");
497+
assertThat(second.name()).hasValue("getCondition");
498+
assertThat(second.args().get()).containsExactly("city", "Warsaw", "unit", "F");
499+
}
500+
501+
// Safety guard for non-conforming output: a streamed call still in progress (the model should
502+
// have terminated it with willContinue=false) is followed by a complete non-streaming call. The
503+
// in-progress call is flushed before appending, so neither is dropped nor merged.
504+
@Test
505+
public void processRawResponses_streamedCallFollowedByCompleteCall_flushesInProgressFirst() {
506+
GenerateContentResponse streamedName =
507+
toResponse(
508+
functionCallPart(
509+
FunctionCall.builder().name("stream_call").willContinue(true).build()));
510+
GenerateContentResponse streamedArg =
511+
toResponse(
512+
functionCallPart(
513+
FunctionCall.builder()
514+
.partialArgs(PartialArg.builder().jsonPath("$.a").stringValue("1").build())
515+
.willContinue(true)
516+
.build()));
517+
GenerateContentResponse completeCall =
518+
toResponse(
519+
Candidate.builder()
520+
.content(
521+
Content.builder()
522+
.parts(
523+
functionCallPart(
524+
FunctionCall.builder()
525+
.name("plain_call")
526+
.args(ImmutableMap.of("b", "2"))
527+
.build()))
528+
.build())
529+
.finishReason(new FinishReason(FinishReason.Known.STOP))
530+
.build());
531+
532+
ImmutableList<LlmResponse> responses =
533+
ImmutableList.copyOf(
534+
Gemini.processRawResponses(Flowable.just(streamedName, streamedArg, completeCall))
535+
.blockingIterable());
536+
537+
LlmResponse finalResponse = Iterables.getLast(responses);
538+
assertThat(finalResponse.content().get().parts().get()).hasSize(2);
539+
FunctionCall first = finalResponse.content().get().parts().get().get(0).functionCall().get();
540+
FunctionCall second = finalResponse.content().get().parts().get().get(1).functionCall().get();
541+
assertThat(first.name()).hasValue("stream_call");
542+
assertThat(first.args().get()).containsExactly("a", "1");
543+
assertThat(second.name()).hasValue("plain_call");
544+
assertThat(second.args().get()).containsExactly("b", "2");
545+
}
546+
547+
// A stray nameless willContinue=false marker with no call in progress must be a safe no-op (the
548+
// currentFcName != null half of the guard): it must not add a function call nor split the
549+
// surrounding text. Without that half it would be treated as a streamed part and prematurely
550+
// flush the text buffer, splitting "Hello world" into two parts.
551+
@Test
552+
public void processRawResponses_strayNamelessMarker_isNoOpAndDoesNotSplitText() {
553+
GenerateContentResponse hello = toResponseWithText("Hello ");
554+
GenerateContentResponse strayMarker =
555+
toResponse(functionCallPart(FunctionCall.builder().willContinue(false).build()));
556+
GenerateContentResponse world = toResponseWithText("world", FinishReason.Known.STOP);
557+
558+
ImmutableList<LlmResponse> responses =
559+
ImmutableList.copyOf(
560+
Gemini.processRawResponses(Flowable.just(hello, strayMarker, world))
561+
.blockingIterable());
562+
563+
LlmResponse finalResponse = Iterables.getLast(responses);
564+
assertThat(finalResponse.content().get().parts().get()).hasSize(1);
565+
assertThat(finalResponse.content().get().parts().get().get(0).text()).hasValue("Hello world");
566+
assertThat(finalResponse.content().get().parts().get().get(0).functionCall()).isEmpty();
567+
}
568+
374569
@Test
375570
public void processRawResponses_imageOnlyWithStop_emitsFinalImagePart() {
376571
Part imagePart = Part.fromBytes(new byte[] {1, 2, 3}, "image/png");

‎core/src/test/java/com/google/adk/runner/RunnerTest.java‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@
7979
import com.google.genai.types.FunctionDeclaration;
8080
import com.google.genai.types.FunctionResponse;
8181
import com.google.genai.types.Part;
82+
import com.google.genai.types.PartialArg;
8283
import io.opentelemetry.api.trace.Tracer;
8384
import io.opentelemetry.context.Context;
8485
import io.opentelemetry.context.ContextKey;
@@ -1378,6 +1379,103 @@ public void runAsync_withSessionKey_success() {
13781379
assertThat(simplifyEvents(events)).containsExactly("test agent: from llm");
13791380
}
13801381

1382+
// Runner-level regression for streamed function-call arguments: a multi-arg call whose args
1383+
// arrive
1384+
// across partial events (nameless continuation chunks, with one value split across chunks) must
1385+
// not crash and must execute the tool exactly once with the reassembled args.
1386+
@Test
1387+
public void runAsync_streamedFunctionCallArgs_reassembledAndToolExecuted() {
1388+
// Turn 1: SSE stream mimicking the post-aggregator shape - a named chunk with willContinue,
1389+
// then
1390+
// nameless continuation chunks carrying partialArgs (origin split across two), then the
1391+
// aggregated complete call.
1392+
LlmResponse namedChunk =
1393+
partialFcResponse(
1394+
FunctionCall.builder().id("fc-1").name(echoTool.name()).willContinue(true).build());
1395+
LlmResponse originChunk1 =
1396+
partialFcResponse(
1397+
FunctionCall.builder()
1398+
.partialArgs(PartialArg.builder().jsonPath("$.origin").stringValue("Krak").build())
1399+
.willContinue(true)
1400+
.build());
1401+
LlmResponse originChunk2 =
1402+
partialFcResponse(
1403+
FunctionCall.builder()
1404+
.partialArgs(PartialArg.builder().jsonPath("$.origin").stringValue("ow").build())
1405+
.willContinue(true)
1406+
.build());
1407+
LlmResponse destinationChunk =
1408+
partialFcResponse(
1409+
FunctionCall.builder()
1410+
.partialArgs(
1411+
PartialArg.builder().jsonPath("$.destination").stringValue("Warsaw").build())
1412+
.willContinue(true)
1413+
.build());
1414+
LlmResponse aggregatedCall =
1415+
createFunctionCallLlmResponse(
1416+
"fc-1", echoTool.name(), ImmutableMap.of("origin", "Krakow", "destination", "Warsaw"));
1417+
1418+
TestLlm streamingLlm =
1419+
createTestLlm(
1420+
Flowable.just(namedChunk, originChunk1, originChunk2, destinationChunk, aggregatedCall),
1421+
Flowable.just(createTextLlmResponse("done")));
1422+
LlmAgent streamingAgent =
1423+
createTestAgentBuilder(streamingLlm).tools(ImmutableList.of(new EchoTool())).build();
1424+
Runner streamingRunner =
1425+
Runner.builder().app(App.builder().name("test").rootAgent(streamingAgent).build()).build();
1426+
Session streamingSession =
1427+
streamingRunner.sessionService().createSession("test", "user").blockingGet();
1428+
1429+
List<Event> events =
1430+
streamingRunner
1431+
.runAsync(
1432+
"user",
1433+
streamingSession.id(),
1434+
createContent("book a flight"),
1435+
RunConfig.builder().setStreamingMode(RunConfig.StreamingMode.SSE).build())
1436+
.toList()
1437+
.blockingGet();
1438+
1439+
int toolResponses = 0;
1440+
boolean sawPartialFcChunk = false;
1441+
boolean sawFinalText = false;
1442+
Map<String, Object> executedArgs = null;
1443+
for (Event e : events) {
1444+
toolResponses += e.functionResponses().size();
1445+
if (e.partial().orElse(false) && !e.functionCalls().isEmpty()) {
1446+
sawPartialFcChunk = true;
1447+
}
1448+
if (!e.partial().orElse(false) && !e.functionCalls().isEmpty()) {
1449+
executedArgs = e.functionCalls().get(0).args().orElse(ImmutableMap.of());
1450+
}
1451+
boolean hasDone =
1452+
e.content()
1453+
.flatMap(Content::parts)
1454+
.map(
1455+
parts ->
1456+
parts.stream()
1457+
.anyMatch(p -> p.text().map(t -> t.contains("done")).orElse(false)))
1458+
.orElse(false);
1459+
sawFinalText |= hasDone;
1460+
}
1461+
1462+
// The streamed (incl. nameless) chunks flowed through the runner without crashing; the tool ran
1463+
// exactly once (only the aggregated non-partial call triggers execution) with both reassembled
1464+
// args; and the final text was produced.
1465+
assertThat(sawPartialFcChunk).isTrue();
1466+
assertThat(toolResponses).isEqualTo(1);
1467+
assertThat(executedArgs).containsExactly("origin", "Krakow", "destination", "Warsaw");
1468+
assertThat(sawFinalText).isTrue();
1469+
}
1470+
1471+
private static LlmResponse partialFcResponse(FunctionCall fc) {
1472+
return LlmResponse.builder()
1473+
.content(
1474+
Content.builder().role("model").parts(Part.builder().functionCall(fc).build()).build())
1475+
.partial(true)
1476+
.build();
1477+
}
1478+
13811479
@Test
13821480
public void runAsync_withStateDelta_mergesStateIntoSession() {
13831481
ImmutableMap<String, Object> stateDelta = ImmutableMap.of("key1", "value1", "key2", 42);

0 commit comments

Comments
 (0)