chore: merge latest main to resolve conflicts

Author: jpantsjoha

contributing/samples/agent_engine_code_execution/README

@@ -9,7 +9,7 @@ This sample data science agent uses Agent Engine Code Execution Sandbox to execu
 
 * 1. Follow https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/code-execution/overview to create a code execution sandbox environment.
 
-* 2. Replace the SANDBOX_RESOURCE_NAME with the one you just created. If you dont want to create a new sandbox environment directly, the Agent Engine Code Execution Sandbox will create one for you by default using the AGENT_ENGINE_RESOURCE_NAME you specified, however, please ensure to clean up sandboxes after use, otherwise, it will consume quotas. 
+* 2. Replace the SANDBOX_RESOURCE_NAME with the one you just created. If you dont want to create a new sandbox environment directly, the Agent Engine Code Execution Sandbox will create one for you by default using the AGENT_ENGINE_RESOURCE_NAME you specified, however, please ensure to clean up sandboxes after use, otherwise, it will consume quotas.
 
 
 ## Sample prompt

src/google/adk/cli/adk_web_server.py

@@ -486,6 +486,12 @@ async def get_runner_async(self, app_name: str) -> Runner:
     self.runner_dict[app_name] = runner
     return runner
 
+  def _get_root_agent(self, agent_or_app: BaseAgent | App) -> BaseAgent:
+    """Extract root agent from either a BaseAgent or App object."""
+    if isinstance(agent_or_app, App):
+      return agent_or_app.root_agent
+    return agent_or_app
+
   def _create_runner(self, agentic_app: App) -> Runner:
     """Create a runner with common services."""
     return Runner(
@@ -741,7 +747,7 @@ async def create_session_with_id(
           is not None
       ):
         raise HTTPException(
-            status_code=400, detail=f"Session already exists: {session_id}"
+            status_code=409, detail=f"Session already exists: {session_id}"
         )
       session = await self.session_service.create_session(
           app_name=app_name, user_id=user_id, state=state, session_id=session_id
@@ -763,6 +769,13 @@ async def create_session(
             app_name=app_name, user_id=user_id
         )
 
+      if req.session_id and await self.session_service.get_session(
+          app_name=app_name, user_id=user_id, session_id=req.session_id
+      ):
+        raise HTTPException(
+            status_code=409, detail=f"Session already exists: {req.session_id}"
+        )
+
       session = await self.session_service.create_session(
           app_name=app_name,
           user_id=user_id,
@@ -926,9 +939,8 @@ async def add_session_to_eval_set(
 
       # Populate the session with initial session state.
       agent_or_app = self.agent_loader.load_agent(app_name)
-      if isinstance(agent_or_app, App):
-        agent_or_app = agent_or_app.root_agent
-      initial_session_state = create_empty_state(agent_or_app)
+      root_agent = self._get_root_agent(agent_or_app)
+      initial_session_state = create_empty_state(root_agent)
 
       new_eval_case = EvalCase(
           eval_id=req.eval_id,
@@ -1089,7 +1101,8 @@ async def run_eval(
               status_code=400, detail=f"Eval set `{eval_set_id}` not found."
           )
 
-        root_agent = self.agent_loader.load_agent(app_name)
+        agent_or_app = self.agent_loader.load_agent(app_name)
+        root_agent = self._get_root_agent(agent_or_app)
 
         eval_case_results = []
 
@@ -1430,13 +1443,7 @@ async def get_event_graph(
       function_calls = event.get_function_calls()
       function_responses = event.get_function_responses()
       agent_or_app = self.agent_loader.load_agent(app_name)
-      # The loader may return an App; unwrap to its root agent so the graph builder
-      # receives a BaseAgent instance.
-      root_agent = (
-          agent_or_app.root_agent
-          if isinstance(agent_or_app, App)
-          else agent_or_app
-      )
+      root_agent = self._get_root_agent(agent_or_app)
       dot_graph = None
       if function_calls:
         function_call_highlights = []

src/google/adk/sessions/vertex_ai_session_service.py

@@ -73,7 +73,23 @@ async def create_session(
       user_id: str,
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
+      **kwargs: Any,
   ) -> Session:
+    """Creates a new session.
+
+    Args:
+      app_name: The name of the application.
+      user_id: The ID of the user.
+      state: The initial state of the session.
+      session_id: The ID of the session.
+      **kwargs: Additional arguments to pass to the session creation. E.g. set
+        expire_time='2025-10-01T00:00:00Z' to set the session expiration time.
+        See https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/projects.locations.reasoningEngines.sessions
+        for more details.
+    Returns:
+      The created session.
+    """
+
     if session_id:
       raise ValueError(
           'User-provided Session id is not supported for'
@@ -84,6 +100,7 @@ async def create_session(
     api_client = self._get_api_client()
 
     config = {'session_state': state} if state else {}
+    config.update(kwargs)
 
     if _is_vertex_express_mode(self._project, self._location):
       config['wait_for_completion'] = False

tests/unittests/sessions/test_vertex_ai_session_service.py

@@ -242,6 +242,7 @@ def __init__(self) -> None:
     self.agent_engines.sessions.create.side_effect = self._create_session
     self.agent_engines.sessions.events.list.side_effect = self._list_events
     self.agent_engines.sessions.events.append.side_effect = self._append_event
+    self.last_create_session_config: dict[str, Any] = {}
 
   def _get_session(self, name: str):
     session_id = name.split('/')[-1]
@@ -275,6 +276,7 @@ def _delete_session(self, name: str):
     self.session_dict.pop(session_id)
 
   def _create_session(self, name: str, user_id: str, config: dict[str, Any]):
+    self.last_create_session_config = config
     new_session_id = '4'
     self.session_dict[new_session_id] = {
         'name': (
@@ -360,7 +362,8 @@ def mock_vertex_ai_session_service(agent_engine_id: Optional[str] = None):
 
 
 @pytest.fixture
-def mock_get_api_client():
+def mock_api_client_instance():
+  """Creates a mock API client instance for testing."""
   api_client = MockApiClient()
   api_client.session_dict = {
       '1': MOCK_SESSION_JSON_1,
@@ -373,9 +376,15 @@ def mock_get_api_client():
       '1': (copy.deepcopy(MOCK_EVENT_JSON), None),
       '2': (copy.deepcopy(MOCK_EVENT_JSON_2), 'my_token'),
   }
+  return api_client
+
+
+@pytest.fixture
+def mock_get_api_client(mock_api_client_instance):
+  """Mocks the _get_api_client method to return a mock API client."""
   with mock.patch(
       'google.adk.sessions.vertex_ai_session_service.VertexAiSessionService._get_api_client',
-      return_value=api_client,
+      return_value=mock_api_client_instance,
   ):
     yield
 
@@ -521,6 +530,21 @@ async def test_create_session_with_custom_session_id():
   )
 
 
+@pytest.mark.asyncio
+@pytest.mark.usefixtures('mock_get_api_client')
+async def test_create_session_with_custom_config(mock_api_client_instance):
+  session_service = mock_vertex_ai_session_service()
+
+  expire_time = '2025-12-12T12:12:12.123456Z'
+  await session_service.create_session(
+      app_name='123', user_id='user', expire_time=expire_time
+  )
+  assert (
+      mock_api_client_instance.last_create_session_config['expire_time']
+      == expire_time
+  )
+
+
 @pytest.mark.asyncio
 @pytest.mark.usefixtures('mock_get_api_client')
 async def test_append_event():