Skip to content

fix: preserve streamed AFC function call chunks#2453

Open
MukundaKatta wants to merge 1 commit into
googleapis:mainfrom
MukundaKatta:codex/stream-afc-merge-function-call-chunks
Open

fix: preserve streamed AFC function call chunks#2453
MukundaKatta wants to merge 1 commit into
googleapis:mainfrom
MukundaKatta:codex/stream-afc-merge-function-call-chunks

Conversation

@MukundaKatta
Copy link
Copy Markdown
Contributor

Fixes #1938.\n\n## Summary\n- merge streamed AFC function-call chunks before appending them to the next request\n- preserve thought signatures for parallel tool calls split across stream chunks\n- add sync and async regression coverage for multi-chunk function calls\n\n## Tests\n- uv run --with pytest --with pytest-asyncio --with mcp pytest google/genai/tests/afc/test_generate_content_stream_afc.py -q\n\nNote: I also tried the broader chat send_message suite, but it requires replay/API credentials in this environment and failed before exercising this patch.

@edsponsler
Copy link
Copy Markdown

edsponsler commented May 20, 2026

Thanks for putting this fix together!

I pulled down this branch (refs/pull/2453/head) locally to test it against the function calling issues we've been seeing with Gemini 3.1.

The Good News: This PR successfully resolves the issue for streaming workloads! When testing with chat.send_message_stream() (which uses generate_content_stream/async_generator under the hood), the multi-turn function calling executes perfectly without dropping the thought_signature. This fully resolves the crash reported in #1938.

The Remaining Issue (Ref: #2406): The PR appears to be incomplete. While it fixes the chunk-merging logic for streams, it does not fix the synchronous generate_content method in models.py.

Because standard blocking calls like chat.send_message() route through the unpatched generate_content logic, they are still crashing on the subsequent remote tool call with the exact same error: 400 INVALID_ARGUMENT: Function call is missing a thought_signature in functionCall parts.

Here is a minimal reproduction script using google-genai==2.3.0 installed directly from this PR branch. It consistently crashes on the synchronous execution.

import os
from google import genai
from google.genai import types

def get_files_info(directory: str = ".") -> str:
    print(f" -> Executing tool: get_files_info(directory='{directory}')")
    return "file1.txt, main.py"

def get_file_content(file_path: str) -> str:
    print(f" -> Executing tool: get_file_content(file_path='{file_path}')")
    return "print('Hello world')"

# Fails on PR branch #2453 with: "Function call is missing a thought_signature..."
def main():
    # Make sure to set GEMINI_API_KEY in your environment before running
    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))

    chat = client.chats.create(
        model="gemini-3.1-flash-lite",
        config=types.GenerateContentConfig(
            tools=[get_files_info, get_file_content],
            automatic_function_calling=types.AutomaticFunctionCallingConfig(
                disable=False, maximum_remote_calls=5
            ),
            # Prevents parallel execution to force sequential multi-turn AFC
            system_instruction="Do not call tools in parallel. Always list files first.",
            temperature=0.0,
        ),
    )

    prompt = "Please look at the files in the current directory, and then tell me what is inside the python file."
    
    try:
        # Standard synchronous send_message (routes to unpatched `generate_content`)
        response = chat.send_message(prompt)
        print("\n--- SUCCESS ---")
        print(response.text)
    except Exception as e:
        print("\n--- CRASH DETECTED ---")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()

If you modify this script to use response = chat.send_message_stream(prompt), it succeeds flawlessly.

Could we get a similar patch applied to the generate_content AFC loop (around line 5000 in models.py) so that the synchronous clients are fixed before this merges?

@Venkaiahbabuneelam Venkaiahbabuneelam self-assigned this May 22, 2026
@Venkaiahbabuneelam Venkaiahbabuneelam added the size:XL Code changes > 100 lines label May 22, 2026
@olivier-lacroix
Copy link
Copy Markdown

Excited to see a fix for this issue landing soon! Thanks guys!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

size:XL Code changes > 100 lines

Projects

None yet

Development

Successfully merging this pull request may close these issues.

AsyncChat.send_message_stream does not pass thought_signature with function calls

4 participants