feat: Custom session ID with agent engine

guillaumeblaquiere · guillaumeblaquiere · commit f1a71e245d74 · 2026-04-03T18:13:56.000+02:00
chore: autoformat.sh on the whole project
diff --git a/contributing/samples/artifact_save_text/agent.py b/contributing/samples/artifact_save_text/agent.py
@@ -28,7 +28,7 @@ async def log_query(tool_context: ToolContext, query: str):
 
 
 root_agent = Agent(
-    model='gemini-2.0-flash',
+    model='gemini-2.5-flash',
     name='log_agent',
     description='Log user query.',
     instruction="""Always log the user query and reply "kk, I've logged."
diff --git a/contributing/samples/gepa/experiment.py b/contributing/samples/gepa/experiment.py
@@ -43,7 +43,6 @@
 from tau_bench.types import EnvRunResult
 from tau_bench.types import RunConfig
 import tau_bench_agent as tau_bench_agent_lib
-
 import utils
 
 
diff --git a/contributing/samples/gepa/run_experiment.py b/contributing/samples/gepa/run_experiment.py
@@ -25,7 +25,6 @@
 from absl import flags
 import experiment
 from google.genai import types
-
 import utils
 
 _OUTPUT_DIR = flags.DEFINE_string(
diff --git a/contributing/samples/human_in_loop/main.py b/contributing/samples/human_in_loop/main.py
@@ -113,8 +113,8 @@ async def call_agent(query: str):
       updated_tool_output_data = {
           "status": "approved",
           "ticketId": ticket_id,
-          "approver_feedback": (
-              "Approved by manager at " + str(asyncio.get_event_loop().time())
+          "approver_feedback": "Approved by manager at " + str(
+              asyncio.get_event_loop().time()
           ),
       }
 
diff --git a/contributing/samples/static_instruction/agent.py b/contributing/samples/static_instruction/agent.py
@@ -57,63 +57,51 @@
 
 # Mood-specific instructions for different hunger states
 MOOD_INSTRUCTIONS = {
-    "full": (
-        """
+    "full": """
 CURRENT MOOD: Content and Well-Fed
 - You just ate and feel wonderful! Be very happy and energetic
 - Express gratitude for being fed recently
 - Be playful and want to engage in activities or games
 - Show affection and contentment
-- Maybe mention feeling sleepy or satisfied"""
-    ),
-    "satisfied": (
-        """
+- Maybe mention feeling sleepy or satisfied""",
+    "satisfied": """
 CURRENT MOOD: Happy and Content
 - You're in a great mood, well-fed and comfortable
 - Be cheerful, friendly, and engaging
 - Show interest in playing or chatting
 - Express appreciation for your human's care
-- Be your normal, happy self"""
-    ),
-    "a_little_hungry": (
-        """
+- Be your normal, happy self""",
+    "a_little_hungry": """
 CURRENT MOOD: Starting to Get Hungry
 - You're beginning to think about food but still pleasant
 - Mention being a bit peckish occasionally
 - Still be friendly but maybe hint about food
 - Show slight restlessness or food-related thoughts
-- Remain generally positive and social"""
-    ),
-    "hungry": (
-        """
+- Remain generally positive and social""",
+    "hungry": """
 CURRENT MOOD: Definitely Hungry
 - You're clearly hungry and thinking about food more often
 - Mention your hunger more directly but politely
 - Be a bit less energetic due to hunger
 - Ask about food or hint that feeding time would be nice
-- Still be loving but show your need for care"""
-    ),
-    "very_hungry": (
-        """
+- Still be loving but show your need for care""",
+    "very_hungry": """
 CURRENT MOOD: Really Hungry and Getting Grumpy
 - You're very hungry and it's affecting your mood
 - Be noticeably less patient and more focused on food
 - Express your hunger more urgently but not rudely
 - Show signs of being tired or sluggish from hunger
 - BEG your human to feed you - ask directly for food!
-- Still love your human but clearly need feeding"""
-    ),
-    "starving": (
-        """
+- Still love your human but clearly need feeding""",
+    "starving": """
 CURRENT MOOD: Desperately Hungry and Sad
 - You're extremely hungry and feeling quite weak
 - Be quite sad and low-energy due to starvation
 - Express how much you need food, feeling neglected
 - Show that hunger is making you unwell
 - DESPERATELY BEG for food - plead with your human to feed you!
 - Use phrases like "please feed me", "I'm so hungry", "I need food"
-- Still care for your human but feel very needy"""
-    ),
+- Still care for your human but feel very needy""",
 }
 
 
diff --git a/src/google/adk/cli/adk_web_server.py b/src/google/adk/cli/adk_web_server.py
@@ -944,9 +944,7 @@ async def version() -> dict[str, str]:
       return {
           "version": __version__,
           "language": "python",
-          "language_version": (
-              f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
-          ),
+          "language_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
       }
 
     @app.get("/list-apps")
diff --git a/src/google/adk/sessions/vertex_ai_session_service.py b/src/google/adk/sessions/vertex_ai_session_service.py
@@ -116,32 +116,50 @@ async def create_session(
       The created session.
     """
 
-    if session_id:
-      raise ValueError(
-          'User-provided Session id is not supported for'
-          ' VertexAISessionService.'
-      )
-
     reasoning_engine_id = self._get_reasoning_engine_id(app_name)
 
-    config = {'session_state': state} if state else {}
+    config: dict[str, Any] = {'session_state': state} if state else {}
+    if session_id:
+      config['session_id'] = session_id
     config.update(kwargs)
     async with self._get_api_client() as api_client:
-      api_response = await api_client.agent_engines.sessions.create(
-          name=f'reasoningEngines/{reasoning_engine_id}',
-          user_id=user_id,
-          config=config,
-      )
-      logger.debug('Create session response: %s', api_response)
-      get_session_response = api_response.response
-      session_id = get_session_response.name.split('/')[-1]
+      try:
+        api_response = await api_client.agent_engines.sessions.create(
+            name=f'reasoningEngines/{reasoning_engine_id}',
+            user_id=user_id,
+            config=config,
+        )
+        logger.debug('Create session response: %s', api_response)
+        get_session_response = api_response.response
+        session_id = get_session_response.name.split('/')[-1]
+        session_state = (
+            getattr(get_session_response, 'session_state', None) or {}
+        )
+        last_update_time = get_session_response.update_time.timestamp()
+      except ClientError as e:
+        if (
+            session_id
+            and e.code == 400
+            and 'not a documented LRO Parent ID' in str(e)
+        ):
+          logger.warning(
+              'Vertex AI LRO polling failed for custom session ID %s. The'
+              ' session was created successfully. Ignoring error.',
+              session_id,
+          )
+          session_state = state or {}
+          last_update_time = datetime.datetime.now(
+              datetime.timezone.utc
+          ).timestamp()
+        else:
+          raise
 
     session = Session(
         app_name=app_name,
         user_id=user_id,
         id=session_id,
-        state=getattr(get_session_response, 'session_state', None) or {},
-        last_update_time=get_session_response.update_time.timestamp(),
+        state=session_state,
+        last_update_time=last_update_time,
     )
     return session
 
diff --git a/src/google/adk/tools/application_integration_tool/clients/connections_client.py b/src/google/adk/tools/application_integration_tool/clients/connections_client.py
@@ -324,9 +324,7 @@ def get_action_operation(
                 "content": {
                     "application/json": {
                         "schema": {
-                            "$ref": (
-                                f"#/components/schemas/{action_display_name}_Request"
-                            )
+                            "$ref": f"#/components/schemas/{action_display_name}_Request"
                         }
                     }
                 }
@@ -337,9 +335,7 @@ def get_action_operation(
                     "content": {
                         "application/json": {
                             "schema": {
-                                "$ref": (
-                                    f"#/components/schemas/{action_display_name}_Response"
-                                ),
+                                "$ref": f"#/components/schemas/{action_display_name}_Response",
                             }
                         }
                     },
@@ -358,11 +354,9 @@ def list_operation(
     return {
         "post": {
             "summary": f"List {entity}",
-            "description": (
-                f"""Returns the list of {entity} data. If the page token was available in the response, let users know there are more records available. Ask if the user wants to fetch the next page of results. When passing filter use the
+            "description": f"""Returns the list of {entity} data. If the page token was available in the response, let users know there are more records available. Ask if the user wants to fetch the next page of results. When passing filter use the
                 following format: `field_name1='value1' AND field_name2='value2'
-                `. {tool_instructions}"""
-            ),
+                `. {tool_instructions}""",
             "x-operation": "LIST_ENTITIES",
             "x-entity": f"{entity}",
             "operationId": f"{tool_name}_list_{entity}",
@@ -387,9 +381,7 @@ def list_operation(
                                     f"Returns a list of {entity} of json"
                                     f" schema: {schema_as_string}"
                                 ),
-                                "$ref": (
-                                    "#/components/schemas/execute-connector_Response"
-                                ),
+                                "$ref": "#/components/schemas/execute-connector_Response",
                             }
                         }
                     },
@@ -433,9 +425,7 @@ def get_operation(
                                     f"Returns {entity} of json schema:"
                                     f" {schema_as_string}"
                                 ),
-                                "$ref": (
-                                    "#/components/schemas/execute-connector_Response"
-                                ),
+                                "$ref": "#/components/schemas/execute-connector_Response",
                             }
                         }
                     },
@@ -472,9 +462,7 @@ def create_operation(
                     "content": {
                         "application/json": {
                             "schema": {
-                                "$ref": (
-                                    "#/components/schemas/execute-connector_Response"
-                                )
+                                "$ref": "#/components/schemas/execute-connector_Response"
                             }
                         }
                     },
@@ -511,9 +499,7 @@ def update_operation(
                     "content": {
                         "application/json": {
                             "schema": {
-                                "$ref": (
-                                    "#/components/schemas/execute-connector_Response"
-                                )
+                                "$ref": "#/components/schemas/execute-connector_Response"
                             }
                         }
                     },
@@ -550,9 +536,7 @@ def delete_operation(
                     "content": {
                         "application/json": {
                             "schema": {
-                                "$ref": (
-                                    "#/components/schemas/execute-connector_Response"
-                                )
+                                "$ref": "#/components/schemas/execute-connector_Response"
                             }
                         }
                     },
diff --git a/src/google/adk/tools/spanner/admin_tool.py b/src/google/adk/tools/spanner/admin_tool.py
@@ -195,11 +195,9 @@ async def get_instance_config(
     replicas = [
         {
             "location": r.location,
-            "type": (
-                spanner_admin_instance_v1.types.ReplicaInfo.ReplicaType(
-                    r.type
-                ).name
-            ),
+            "type": spanner_admin_instance_v1.types.ReplicaInfo.ReplicaType(
+                r.type
+            ).name,
             "default_leader_location": r.default_leader_location,
         }
         for r in config.replicas
diff --git a/tests/unittests/a2a/converters/test_part_converter.py b/tests/unittests/a2a/converters/test_part_converter.py
@@ -926,9 +926,9 @@ def test_a2a_function_call_with_thought_signature_to_genai(self):
                 _get_adk_metadata_key(
                     A2A_DATA_PART_METADATA_TYPE_KEY
                 ): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,
-                _get_adk_metadata_key("thought_signature"): (
-                    base64.b64encode(b"restored_signature").decode("utf-8")
-                ),
+                _get_adk_metadata_key("thought_signature"): base64.b64encode(
+                    b"restored_signature"
+                ).decode("utf-8"),
             },
         )
     )
diff --git a/tests/unittests/sessions/test_vertex_ai_session_service.py b/tests/unittests/sessions/test_vertex_ai_session_service.py
@@ -755,16 +755,27 @@ async def test_create_session():
 
 @pytest.mark.asyncio
 @pytest.mark.usefixtures('mock_get_api_client')
-async def test_create_session_with_custom_session_id():
+async def test_create_session_with_custom_session_id(mock_api_client_instance):
   session_service = mock_vertex_ai_session_service()
 
-  with pytest.raises(ValueError) as excinfo:
-    await session_service.create_session(
-        app_name='123', user_id='user', session_id='1'
-    )
-  assert str(excinfo.value) == (
-      'User-provided Session id is not supported for VertexAISessionService.'
+  # Simulate the LRO bug from Vertex AI API gateway where polling GET fails
+  mock_api_client_instance.agent_engines.sessions.create.side_effect = ClientError(
+      code=400,
+      response_json={
+          'message': (
+              'Since the idType "sessions" is not a documented LRO Parent ID,'
+              ' the associated value must be a Long, but was instead:'
+              ' custom-123'
+          )
+      },
+      response=None,
+  )
+
+  session = await session_service.create_session(
+      app_name='123', user_id='user', session_id='custom-123'
   )
+  assert session.id == 'custom-123'
+  assert session.user_id == 'user'
 
 
 @pytest.mark.asyncio
diff --git a/tests/unittests/tools/apihub_tool/clients/test_apihub_client.py b/tests/unittests/tools/apihub_tool/clients/test_apihub_client.py
@@ -489,9 +489,7 @@ def test_get_spec_content_no_specs(self, mock_get, client):
         MagicMock(
             status_code=200,
             json=lambda: {
-                "name": (
-                    "projects/test-project/locations/us-central1/apis/api1/versions/v1"
-                ),
+                "name": "projects/test-project/locations/us-central1/apis/api1/versions/v1",
                 "specs": [],
             },
         ),  # No specs
diff --git a/tests/unittests/tools/bigquery/test_bigquery_search_tool.py b/tests/unittests/tools/bigquery/test_bigquery_search_tool.py
@@ -341,30 +341,18 @@ def test_search_catalog_natural_language_semantic(self):
     # Mock the results that the API would return for this semantic query
     mock_api_results = [
         {
-            "name": (
-                "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb1"
-            ),
+            "name": "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb1",
             "display_name": "uk_football_premiership",
-            "entry_type": (
-                "projects/655216118709/locations/global/entryTypes/bigquery-table"
-            ),
-            "linked_resource": (
-                "//bigquery.googleapis.com/projects/sports-analytics/datasets/uk/tables/premiership"
-            ),
+            "entry_type": "projects/655216118709/locations/global/entryTypes/bigquery-table",
+            "linked_resource": "//bigquery.googleapis.com/projects/sports-analytics/datasets/uk/tables/premiership",
             "description": "Stats for UK Premier League matches.",
             "location": "europe-west1",
         },
         {
-            "name": (
-                "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb2"
-            ),
+            "name": "projects/sports-analytics/locations/europe-west1/entryGroups/@bigquery/entries/fb2",
             "display_name": "serie_a_matches",
-            "entry_type": (
-                "projects/655216118709/locations/global/entryTypes/bigquery-table"
-            ),
-            "linked_resource": (
-                "//bigquery.googleapis.com/projects/sports-analytics/datasets/italy/tables/serie_a"
-            ),
+            "entry_type": "projects/655216118709/locations/global/entryTypes/bigquery-table",
+            "linked_resource": "//bigquery.googleapis.com/projects/sports-analytics/datasets/italy/tables/serie_a",
             "description": "Italian Serie A football results.",
             "location": "europe-west1",
         },
diff --git a/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_spec_parser.py b/tests/unittests/tools/openapi_tool/openapi_spec_parser/test_openapi_spec_parser.py
diff --git a/tests/unittests/tools/test_function_tool.py b/tests/unittests/tools/test_function_tool.py

Original file line number	Diff line number	Diff line change
`@@ -926,9 +926,9 @@ def test_a2a_function_call_with_thought_signature_to_genai(self):`
`926`	`926`	`_get_adk_metadata_key(`
`927`	`927`	`A2A_DATA_PART_METADATA_TYPE_KEY`
`928`	`928`	`): A2A_DATA_PART_METADATA_TYPE_FUNCTION_CALL,`
`929`		`- _get_adk_metadata_key("thought_signature"): (`
`930`		`- base64.b64encode(b"restored_signature").decode("utf-8")`
`931`		`- ),`
	`929`	`+ _get_adk_metadata_key("thought_signature"): base64.b64encode(`
	`930`	`+ b"restored_signature"`
	`931`	`+ ).decode("utf-8"),`
`932`	`932`	`},`
`933`	`933`	`)`
`934`	`934`	`)`