2323import json
2424import logging
2525import typing
26+ from unittest .mock import AsyncMock
2627from unittest .mock import MagicMock
2728
29+ from google .adk .agents ._streaming_mode import StreamingMode
30+ from google .adk .agents .llm_agent import LlmAgent
31+ from google .adk .agents .run_config import RunConfig
2832from google .adk .models import interactions_utils
33+ from google .adk .models .google_llm import Gemini
2934from google .adk .models .llm_request import LlmRequest
3035from google .adk .models .llm_response import LlmResponse
36+ from google .adk .runners import InMemoryRunner
37+ from google .genai import Client
3138from google .genai import interactions
3239from google .genai import types
3340from google .genai .interactions import CodeExecutionResultStep
3744from google .genai .interactions import Interaction
3845from google .genai .interactions import InteractionCompletedEvent
3946from google .genai .interactions import InteractionCreatedEvent
47+ from google .genai .interactions import InteractionSSEEvent
4048from google .genai .interactions import InteractionSseEventInteraction
4149from google .genai .interactions import ModelOutputStep
4250from google .genai .interactions import StepDelta
@@ -1911,8 +1919,15 @@ def test_thought_summary_delta_ignores_non_text_content(self):
19111919 assert result is None
19121920 assert state .parts == []
19131921
1914- def test_thought_signature_delta_attaches_to_last_thought (self ):
1915- """thought_signature mutates the last thought part and emits no event."""
1922+ def test_thought_signature_delta_after_a_summary_gets_its_own_part (
1923+ self ,
1924+ ) -> None :
1925+ """The signature follows the summary on a text-free part of its own.
1926+
1927+ A part that also carries text is re-encoded as a text step, which has no
1928+ signature field, so a signature written onto the summary part would be
1929+ dropped on the next request.
1930+ """
19161931 state = interactions_utils ._StreamState ()
19171932 interactions_utils .convert_interaction_event_to_llm_response (
19181933 StepDelta (
@@ -1936,8 +1951,14 @@ def test_thought_signature_delta_attaches_to_last_thought(self):
19361951 state ,
19371952 interaction_id = 'int_ts' ,
19381953 )
1954+
19391955 assert result is None
1940- assert state .parts [- 1 ].thought_signature == b'sig-bytes'
1956+ assert len (state .parts ) == 2
1957+ assert state .parts [0 ].text == 'reasoning'
1958+ assert state .parts [0 ].thought_signature is None
1959+ assert state .parts [1 ].thought is True
1960+ assert state .parts [1 ].text is None
1961+ assert state .parts [1 ].thought_signature == b'sig-bytes'
19411962
19421963 def test_thought_signature_delta_alone_becomes_its_own_part (self ):
19431964 """A signature with no preceding thought summary still reaches the stream.
@@ -2012,6 +2033,67 @@ def test_thought_signature_delta_without_signature_adds_no_part(self):
20122033 assert result is None
20132034 assert state .parts == []
20142035
2036+ def test_thought_signature_deltas_of_two_thought_steps_are_both_kept (
2037+ self ,
2038+ ) -> None :
2039+ """Each thought step keeps its own signature, in step order.
2040+
2041+ The non-streaming conversion gives every thought step a part of its own,
2042+ so a later step's signature must not overwrite an earlier one's.
2043+ """
2044+ state = interactions_utils ._StreamState ()
2045+ events : list [InteractionSSEEvent ] = [
2046+ StepDelta (
2047+ event_type = 'step.delta' ,
2048+ index = 0 ,
2049+ delta = {
2050+ 'type' : 'thought_signature' ,
2051+ 'signature' : base64 .b64encode (b'first-sig' ).decode ('utf-8' ),
2052+ },
2053+ ),
2054+ StepStart (
2055+ event_type = 'step.start' ,
2056+ index = 1 ,
2057+ step = FunctionCallStep (
2058+ type = 'function_call' ,
2059+ id = 'call_1' ,
2060+ name = 'get_weather' ,
2061+ arguments = {},
2062+ ),
2063+ ),
2064+ StepStop (event_type = 'step.stop' , index = 1 ),
2065+ StepDelta (
2066+ event_type = 'step.delta' ,
2067+ index = 2 ,
2068+ delta = {
2069+ 'type' : 'thought_signature' ,
2070+ 'signature' : base64 .b64encode (b'second-sig' ).decode ('utf-8' ),
2071+ },
2072+ ),
2073+ StepStart (
2074+ event_type = 'step.start' ,
2075+ index = 3 ,
2076+ step = FunctionCallStep (
2077+ type = 'function_call' ,
2078+ id = 'call_2' ,
2079+ name = 'get_time' ,
2080+ arguments = {},
2081+ ),
2082+ ),
2083+ StepStop (event_type = 'step.stop' , index = 3 ),
2084+ ]
2085+ for event in events :
2086+ interactions_utils .convert_interaction_event_to_llm_response (
2087+ event , state , interaction_id = 'int_ts'
2088+ )
2089+
2090+ assert [part .thought_signature for part in state .parts ] == [
2091+ b'first-sig' ,
2092+ None ,
2093+ b'second-sig' ,
2094+ None ,
2095+ ]
2096+
20152097 def test_audio_delta_with_data (self ):
20162098 """audio delta becomes an inline_data part via the shared media handler."""
20172099 event = StepDelta (
@@ -2931,6 +3013,193 @@ def test_arguments_delta_for_finalized_call_logs_warning(self, caplog):
29313013 assert state .parts [0 ].function_call .args == {'city' : 'Paris' }
29323014
29333015
3016+ class TestThoughtSignatureRoundTrip :
3017+ """A tool-calling turn's thought signature reaches the request that follows.
3018+
3019+ With thought summaries on, the summary streams ahead of the signature. Kept
3020+ on the summary part, the signature is lost on the ``previous_interaction_id``
3021+ path: the summary is echoed back as ``model_output`` text and the signature
3022+ is dropped. A stateless Gemini 3 request carrying the tool result without
3023+ the signature is rejected.
3024+ """
3025+
3026+ _SIGNATURE = base64 .b64encode (b'sig-bytes' ).decode ('utf-8' )
3027+ _SUMMARY = 'I should check the weather.'
3028+
3029+ def _streamed_events (self ) -> list [InteractionSSEEvent ]:
3030+ return [
3031+ StepDelta (
3032+ event_type = 'step.delta' ,
3033+ index = 0 ,
3034+ delta = {
3035+ 'type' : 'thought_summary' ,
3036+ 'content' : {'type' : 'text' , 'text' : self ._SUMMARY },
3037+ },
3038+ ),
3039+ StepDelta (
3040+ event_type = 'step.delta' ,
3041+ index = 0 ,
3042+ delta = {'type' : 'thought_signature' , 'signature' : self ._SIGNATURE },
3043+ ),
3044+ StepStart (
3045+ event_type = 'step.start' ,
3046+ index = 1 ,
3047+ step = FunctionCallStep (
3048+ type = 'function_call' ,
3049+ id = 'call_1' ,
3050+ name = 'get_weather' ,
3051+ arguments = {},
3052+ ),
3053+ ),
3054+ StepDelta (
3055+ event_type = 'step.delta' ,
3056+ index = 1 ,
3057+ delta = {'type' : 'arguments_delta' , 'arguments' : '{"city": "Tokyo"}' },
3058+ ),
3059+ StepStop (event_type = 'step.stop' , index = 1 ),
3060+ InteractionCompletedEvent (
3061+ event_type = 'interaction.completed' ,
3062+ interaction = InteractionSseEventInteraction (
3063+ id = 'int_1' , status = 'requires_action' , steps = []
3064+ ),
3065+ ),
3066+ ]
3067+
3068+ def _streamed_turn (self ) -> types .Content :
3069+ state = interactions_utils ._StreamState ()
3070+ responses = [
3071+ interactions_utils .convert_interaction_event_to_llm_response (
3072+ event , state , interaction_id = 'int_1'
3073+ )
3074+ for event in self ._streamed_events ()
3075+ ]
3076+ final_response = responses [- 1 ]
3077+ assert final_response is not None
3078+ assert final_response .content is not None
3079+ return final_response .content
3080+
3081+ def _unary_turn (self ) -> types .Content :
3082+ now = datetime .now (timezone .utc ).isoformat ()
3083+ interaction = Interaction (
3084+ id = 'int_1' ,
3085+ created = now ,
3086+ updated = now ,
3087+ status = 'requires_action' ,
3088+ steps = [
3089+ ThoughtStep (
3090+ type = 'thought' ,
3091+ signature = self ._SIGNATURE ,
3092+ summary = [TextContent (type = 'text' , text = self ._SUMMARY )],
3093+ ),
3094+ FunctionCallStep (
3095+ type = 'function_call' ,
3096+ id = 'call_1' ,
3097+ name = 'get_weather' ,
3098+ arguments = {'city' : 'Tokyo' },
3099+ ),
3100+ ],
3101+ )
3102+ response = interactions_utils .convert_interaction_to_llm_response (
3103+ interaction
3104+ )
3105+ assert response .content is not None
3106+ return response .content
3107+
3108+ @pytest .mark .parametrize ('streamed' , [True , False ], ids = ['streamed' , 'unary' ])
3109+ def test_signature_is_sent_back_with_the_tool_result (
3110+ self , streamed : bool
3111+ ) -> None :
3112+ """The follow-up request echoes the signature and not the summary."""
3113+ model_turn = self ._streamed_turn () if streamed else self ._unary_turn ()
3114+ contents = [
3115+ types .Content (role = 'user' , parts = [types .Part (text = 'Weather?' )]),
3116+ model_turn ,
3117+ types .Content (
3118+ role = 'user' ,
3119+ parts = [
3120+ types .Part (
3121+ function_response = types .FunctionResponse (
3122+ id = 'call_1' ,
3123+ name = 'get_weather' ,
3124+ response = {'weather' : 'sunny' },
3125+ )
3126+ )
3127+ ],
3128+ ),
3129+ ]
3130+
3131+ steps = interactions_utils ._convert_contents_to_steps (
3132+ interactions_utils ._get_latest_user_contents (contents )
3133+ )
3134+
3135+ assert [step ['type' ] for step in steps ] == ['thought' , 'function_result' ]
3136+ assert steps [0 ] == {'type' : 'thought' , 'signature' : self ._SIGNATURE }
3137+
3138+ async def test_runner_sends_the_streamed_signature_back (self ) -> None :
3139+ """An agent streaming over the interactions API echoes the signature."""
3140+ client = MagicMock (spec = Client )
3141+ client .vertexai = False
3142+ client .aio .interactions .create = AsyncMock (
3143+ side_effect = [
3144+ _MockAsyncIterator ([* self ._streamed_events ()]),
3145+ _MockAsyncIterator ([
3146+ StepDelta (
3147+ event_type = 'step.delta' ,
3148+ index = 0 ,
3149+ delta = {'type' : 'text' , 'text' : 'It is sunny.' },
3150+ ),
3151+ InteractionCompletedEvent (
3152+ event_type = 'interaction.completed' ,
3153+ interaction = InteractionSseEventInteraction (
3154+ id = 'int_2' , status = 'completed' , steps = []
3155+ ),
3156+ ),
3157+ ]),
3158+ ]
3159+ )
3160+
3161+ def get_weather (city : str ) -> dict [str , str ]:
3162+ """Returns the weather in a city."""
3163+ del city
3164+ return {'weather' : 'sunny' }
3165+
3166+ runner = InMemoryRunner (
3167+ agent = LlmAgent (
3168+ name = 'weather_agent' ,
3169+ model = Gemini (
3170+ model = 'gemini-3-flash-preview' ,
3171+ use_interactions_api = True ,
3172+ client = client ,
3173+ ),
3174+ tools = [get_weather ],
3175+ ),
3176+ app_name = 'weather_app' ,
3177+ )
3178+ session = await runner .session_service .create_session (
3179+ app_name = 'weather_app' , user_id = 'user'
3180+ )
3181+ async for _ in runner .run_async (
3182+ user_id = 'user' ,
3183+ session_id = session .id ,
3184+ new_message = types .Content (
3185+ role = 'user' , parts = [types .Part (text = 'Weather in Tokyo?' )]
3186+ ),
3187+ run_config = RunConfig (streaming_mode = StreamingMode .SSE ),
3188+ ):
3189+ pass
3190+
3191+ follow_up = client .aio .interactions .create .call_args_list [1 ].kwargs
3192+ assert follow_up ['previous_interaction_id' ] == 'int_1'
3193+ assert [step ['type' ] for step in follow_up ['input' ]] == [
3194+ 'thought' ,
3195+ 'function_result' ,
3196+ ]
3197+ assert follow_up ['input' ][0 ] == {
3198+ 'type' : 'thought' ,
3199+ 'signature' : self ._SIGNATURE ,
3200+ }
3201+
3202+
29343203@pytest .mark .parametrize (
29353204 ('streamed_events_factory' , 'expected_ids' ),
29363205 [
0 commit comments