Skip to content

Commit 20f9fc9

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, and skip nameless streamed chunks when computing long-running tool ids (previously threw). PiperOrigin-RevId: 946860751
1 parent 29c01fd commit 20f9fc9

4 files changed

Lines changed: 298 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: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -416,19 +416,29 @@ 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+
// Part of a streamed call: has partialArgs/willContinue, or is the nameless terminal marker
421+
// of
422+
// an in-progress call. Gemini may keep willContinue=true on the last partialArgs chunk and
423+
// end
424+
// the call with a separate empty willContinue=false part, so that marker completes the call.
425+
boolean streamedPart =
420426
fc.partialArgs().map(args -> !args.isEmpty()).orElse(false)
421-
|| fc.willContinue().orElse(false);
422-
if (streaming) {
427+
|| fc.willContinue().orElse(false)
428+
|| (currentFcName != null && !hasName);
429+
if (streamedPart) {
423430
// Capture the thought signature from the first chunk that carries one.
424431
if (part.thoughtSignature().isPresent() && currentThoughtSignature == null) {
425432
currentThoughtSignature = part.thoughtSignature().get();
426433
}
427434
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.
435+
} else if (hasName) {
436+
// Complete (non-streamed) call. Safety guard: the model should terminate a streamed call
437+
// with willContinue=false before starting a new one; flush any still-in-progress call so it
438+
// is neither dropped nor merged. The part already has an ID assigned by
439+
// ensureFunctionCallIds.
431440
flushTextBufferToSequence();
441+
flushFunctionCallToSequence();
432442
accumulatedSequence.add(part);
433443
}
434444
}
@@ -454,6 +464,8 @@ private void processStreamingFunctionCall(FunctionCall fc) {
454464
}
455465
applyPartialArg(partialArg, jsonPath);
456466
}
467+
// willContinue=false marks the call's last part (final partialArgs chunk or a separate empty
468+
// marker).
457469
if (!fc.willContinue().orElse(false)) {
458470
flushTextBufferToSequence();
459471
flushFunctionCallToSequence();

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

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

374+
// Scenario B: the last partialArgs chunk keeps willContinue=true and completion arrives on a
375+
// separate empty willContinue=false part. Multi-arg call with one value ("origin") split across
376+
// chunks, mirroring the real streaming shape.
377+
@Test
378+
public void processRawResponses_streamedCallEndedByEmptyMarker_flushesCall() {
379+
GenerateContentResponse name =
380+
toResponse(
381+
functionCallPart(FunctionCall.builder().name("bookFlight").willContinue(true).build()));
382+
GenerateContentResponse origin1 =
383+
toResponse(
384+
functionCallPart(
385+
FunctionCall.builder()
386+
.partialArgs(
387+
PartialArg.builder().jsonPath("$.origin").stringValue("Krak").build())
388+
.willContinue(true)
389+
.build()));
390+
GenerateContentResponse origin2 =
391+
toResponse(
392+
functionCallPart(
393+
FunctionCall.builder()
394+
.partialArgs(
395+
PartialArg.builder().jsonPath("$.origin").stringValue("ow").build())
396+
.willContinue(true)
397+
.build()));
398+
GenerateContentResponse destination =
399+
toResponse(
400+
functionCallPart(
401+
FunctionCall.builder()
402+
.partialArgs(
403+
PartialArg.builder()
404+
.jsonPath("$.destination")
405+
.stringValue("Warsaw")
406+
.build())
407+
.willContinue(true)
408+
.build()));
409+
GenerateContentResponse endMarker =
410+
toResponse(
411+
Candidate.builder()
412+
.content(
413+
Content.builder()
414+
.parts(functionCallPart(FunctionCall.builder().willContinue(false).build()))
415+
.build())
416+
.finishReason(new FinishReason(FinishReason.Known.STOP))
417+
.build());
418+
419+
ImmutableList<LlmResponse> responses =
420+
ImmutableList.copyOf(
421+
Gemini.processRawResponses(
422+
Flowable.just(name, origin1, origin2, destination, endMarker))
423+
.blockingIterable());
424+
425+
LlmResponse finalResponse = Iterables.getLast(responses);
426+
assertThat(finalResponse.content().get().parts().get()).hasSize(1);
427+
FunctionCall finalCall =
428+
finalResponse.content().get().parts().get().get(0).functionCall().get();
429+
assertThat(finalCall.name()).hasValue("bookFlight");
430+
assertThat(finalCall.args().get()).containsExactly("origin", "Krakow", "destination", "Warsaw");
431+
}
432+
433+
// Two multi-arg streamed calls each ended by an empty willContinue=false marker must not drop the
434+
// first call nor bleed its args into the second.
435+
@Test
436+
public void processRawResponses_twoStreamedCallsEndedByEmptyMarkers_keepArgsSeparate() {
437+
GenerateContentResponse call1Name =
438+
toResponse(
439+
functionCallPart(
440+
FunctionCall.builder().name("getTemperature").willContinue(true).build()));
441+
GenerateContentResponse call1City =
442+
toResponse(
443+
functionCallPart(
444+
FunctionCall.builder()
445+
.partialArgs(
446+
PartialArg.builder().jsonPath("$.city").stringValue("Krakow").build())
447+
.willContinue(true)
448+
.build()));
449+
GenerateContentResponse call1Unit =
450+
toResponse(
451+
functionCallPart(
452+
FunctionCall.builder()
453+
.partialArgs(PartialArg.builder().jsonPath("$.unit").stringValue("C").build())
454+
.willContinue(true)
455+
.build()));
456+
GenerateContentResponse marker1 =
457+
toResponse(functionCallPart(FunctionCall.builder().willContinue(false).build()));
458+
GenerateContentResponse call2Name =
459+
toResponse(
460+
functionCallPart(
461+
FunctionCall.builder().name("getCondition").willContinue(true).build()));
462+
GenerateContentResponse call2City =
463+
toResponse(
464+
functionCallPart(
465+
FunctionCall.builder()
466+
.partialArgs(
467+
PartialArg.builder().jsonPath("$.city").stringValue("Warsaw").build())
468+
.willContinue(true)
469+
.build()));
470+
GenerateContentResponse call2Unit =
471+
toResponse(
472+
functionCallPart(
473+
FunctionCall.builder()
474+
.partialArgs(PartialArg.builder().jsonPath("$.unit").stringValue("F").build())
475+
.willContinue(true)
476+
.build()));
477+
GenerateContentResponse marker2 =
478+
toResponse(
479+
Candidate.builder()
480+
.content(
481+
Content.builder()
482+
.parts(functionCallPart(FunctionCall.builder().willContinue(false).build()))
483+
.build())
484+
.finishReason(new FinishReason(FinishReason.Known.STOP))
485+
.build());
486+
487+
ImmutableList<LlmResponse> responses =
488+
ImmutableList.copyOf(
489+
Gemini.processRawResponses(
490+
Flowable.just(
491+
call1Name, call1City, call1Unit, marker1, call2Name, call2City, call2Unit,
492+
marker2))
493+
.blockingIterable());
494+
495+
LlmResponse finalResponse = Iterables.getLast(responses);
496+
assertThat(finalResponse.content().get().parts().get()).hasSize(2);
497+
FunctionCall first = finalResponse.content().get().parts().get().get(0).functionCall().get();
498+
FunctionCall second = finalResponse.content().get().parts().get().get(1).functionCall().get();
499+
assertThat(first.name()).hasValue("getTemperature");
500+
assertThat(first.args().get()).containsExactly("city", "Krakow", "unit", "C");
501+
assertThat(second.name()).hasValue("getCondition");
502+
assertThat(second.args().get()).containsExactly("city", "Warsaw", "unit", "F");
503+
}
504+
505+
// Safety guard for non-conforming output: a streamed call still in progress (the model should
506+
// have
507+
// terminated it with willContinue=false) is followed by a complete non-streaming call. The
508+
// in-progress call is flushed before appending, so neither is dropped nor merged.
509+
@Test
510+
public void processRawResponses_streamedCallFollowedByCompleteCall_flushesInProgressFirst() {
511+
GenerateContentResponse streamedName =
512+
toResponse(
513+
functionCallPart(
514+
FunctionCall.builder().name("stream_call").willContinue(true).build()));
515+
GenerateContentResponse streamedArg =
516+
toResponse(
517+
functionCallPart(
518+
FunctionCall.builder()
519+
.partialArgs(PartialArg.builder().jsonPath("$.a").stringValue("1").build())
520+
.willContinue(true)
521+
.build()));
522+
GenerateContentResponse completeCall =
523+
toResponse(
524+
Candidate.builder()
525+
.content(
526+
Content.builder()
527+
.parts(
528+
functionCallPart(
529+
FunctionCall.builder()
530+
.name("plain_call")
531+
.args(ImmutableMap.of("b", "2"))
532+
.build()))
533+
.build())
534+
.finishReason(new FinishReason(FinishReason.Known.STOP))
535+
.build());
536+
537+
ImmutableList<LlmResponse> responses =
538+
ImmutableList.copyOf(
539+
Gemini.processRawResponses(Flowable.just(streamedName, streamedArg, completeCall))
540+
.blockingIterable());
541+
542+
LlmResponse finalResponse = Iterables.getLast(responses);
543+
assertThat(finalResponse.content().get().parts().get()).hasSize(2);
544+
FunctionCall first = finalResponse.content().get().parts().get().get(0).functionCall().get();
545+
FunctionCall second = finalResponse.content().get().parts().get().get(1).functionCall().get();
546+
assertThat(first.name()).hasValue("stream_call");
547+
assertThat(first.args().get()).containsExactly("a", "1");
548+
assertThat(second.name()).hasValue("plain_call");
549+
assertThat(second.args().get()).containsExactly("b", "2");
550+
}
551+
374552
@Test
375553
public void processRawResponses_imageOnlyWithStop_emitsFinalImagePart() {
376554
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)