From 59ddffe83e23be032892c0c49515ccf94b96f01d Mon Sep 17 00:00:00 2001 From: Senthurapandi <127392902+senthurapandi2000@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:51:39 -0700 Subject: [PATCH] fix(a2a): preserve Vertex Agent Engine sessions across messages --- .../a2a_protocol/providers/config_manager.py | 7 + .../providers/vertex_agent_engine/backend.py | 74 +++ .../providers/vertex_agent_engine/config.py | 175 ++++++ .../providers/vertex_agent_engine/session.py | 78 +++ .../test_vertex_agent_engine_provider.py | 524 ++++++++++++++++++ 5 files changed, 858 insertions(+) create mode 100644 litellm/a2a_protocol/providers/vertex_agent_engine/backend.py create mode 100644 litellm/a2a_protocol/providers/vertex_agent_engine/config.py create mode 100644 litellm/a2a_protocol/providers/vertex_agent_engine/session.py create mode 100644 tests/unit/a2a_protocol/providers/test_vertex_agent_engine_provider.py diff --git a/litellm/a2a_protocol/providers/config_manager.py b/litellm/a2a_protocol/providers/config_manager.py index 2eab2adb1ba..0bd50a2f3a3 100644 --- a/litellm/a2a_protocol/providers/config_manager.py +++ b/litellm/a2a_protocol/providers/config_manager.py @@ -46,6 +46,13 @@ class A2AProviderConfigManager: return BedrockAgentCoreA2AConfig() + if custom_llm_provider == "vertex_ai" and model and model.startswith("agent_engine/"): + from litellm.a2a_protocol.providers.vertex_agent_engine.config import ( + VertexAgentEngineA2AConfig, + ) + + return VertexAgentEngineA2AConfig() + if custom_llm_provider == "langflow": from litellm.a2a_protocol.providers.langflow.config import LangFlowA2AConfig diff --git a/litellm/a2a_protocol/providers/vertex_agent_engine/backend.py b/litellm/a2a_protocol/providers/vertex_agent_engine/backend.py new file mode 100644 index 00000000000..2b242452387 --- /dev/null +++ b/litellm/a2a_protocol/providers/vertex_agent_engine/backend.py @@ -0,0 +1,74 @@ +"""Google Vertex AI Agent Engine session API adapter.""" + +from collections.abc import Mapping + +import httpx + + +class VertexAgentEngineSessionBackend: + """List and create sessions through the Agent Engine query API.""" + + def __init__( + self, + *, + client: httpx.AsyncClient, + query_url: str, + auth_headers: Mapping[str, str], + ) -> None: + self.client = client + self.query_url = query_url + self.auth_headers = dict(auth_headers) + + async def _query( + self, + method: str, + user_id: str, + ) -> dict: + response = await self.client.post( + self.query_url, + headers=self.auth_headers, + json={ + "class_method": method, + "input": {"user_id": user_id}, + }, + ) + response.raise_for_status() + + result = response.json() + if not isinstance(result, dict): + raise ValueError("Invalid Agent Engine response") + + output = result.get("output") + if not isinstance(output, dict): + raise ValueError("Missing Agent Engine output") + + return output + + async def list_sessions(self, *, user_id: str) -> list[str]: + output = await self._query("async_list_sessions", user_id) + sessions = output.get("sessions", []) + + if not isinstance(sessions, list): + raise ValueError("Invalid Agent Engine session list") + + session_ids = [] + for session in sessions: + if not isinstance(session, dict): + raise ValueError("Invalid Agent Engine session") + + session_id = session.get("id") + if not isinstance(session_id, str) or not session_id: + raise ValueError("Missing Agent Engine session ID") + + session_ids.append(session_id) + + return session_ids + + async def create_session(self, *, user_id: str) -> str: + output = await self._query("async_create_session", user_id) + + session_id = output.get("id") + if not isinstance(session_id, str) or not session_id: + raise ValueError("Missing newly created Agent Engine session ID") + + return session_id diff --git a/litellm/a2a_protocol/providers/vertex_agent_engine/config.py b/litellm/a2a_protocol/providers/vertex_agent_engine/config.py new file mode 100644 index 00000000000..792e23a87f9 --- /dev/null +++ b/litellm/a2a_protocol/providers/vertex_agent_engine/config.py @@ -0,0 +1,175 @@ +"""A2A completion bridge for Vertex AI Agent Engine sessions.""" + +from collections.abc import AsyncIterator +from uuid import uuid4 + +import httpx + +from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, + A2ACompletionBridgeHandler, +) +from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig +from litellm.a2a_protocol.providers.vertex_agent_engine.backend import ( + VertexAgentEngineSessionBackend, +) +from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + resolve_session_id, + session_user_id, +) +from litellm.a2a_protocol.utils import get_session_id_from_a2a_params +from litellm.llms.vertex_ai.agent_engine.transformation import ( + VertexAgentEngineConfig, +) + + +class VertexAgentEngineA2AConfig(BaseA2AProviderConfig): + """Preserve an isolated Agent Engine session per A2A conversation.""" + + @staticmethod + async def _resolve_session_params( + params: dict[str, object], + litellm_params: dict, + ) -> tuple[str, dict]: + key_hash = litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) + + if not isinstance(key_hash, str) or not key_hash: + raise ValueError("Authenticated calling key hash is required for Vertex Agent Engine A2A sessions") + + incoming_context = get_session_id_from_a2a_params(params) + context_id = incoming_context if isinstance(incoming_context, str) and incoming_context else uuid4().hex + + model = litellm_params.get("model") + if not isinstance(model, str) or not model: + raise ValueError("Vertex Agent Engine model is required") + + engine_config = VertexAgentEngineConfig() + + stream_url = engine_config.get_complete_url( + api_base=None, + api_key=None, + model=model, + optional_params={}, + litellm_params=litellm_params, + stream=True, + ) + + if not stream_url.endswith(":streamQuery"): + raise ValueError("Unexpected Vertex Agent Engine endpoint") + + query_url = stream_url.removesuffix(":streamQuery") + ":query" + + auth_headers = engine_config._get_auth_headers( + optional_params={}, + litellm_params=litellm_params, + ) + + async with httpx.AsyncClient(timeout=30.0) as client: + backend = VertexAgentEngineSessionBackend( + client=client, + query_url=query_url, + auth_headers=auth_headers, + ) + + session_id = await resolve_session_id( + backend, + calling_key_hash=key_hash, + context_id=context_id, + ) + + user_id = session_user_id(key_hash, context_id) + + resolved_params = { + **litellm_params, + "user": user_id, + "user_id": user_id, + "session_id": session_id, + } + + return context_id, resolved_params + + @staticmethod + def _set_stream_context( + event: dict[str, object], + context_id: str, + ) -> dict[str, object]: + """Preserve the A2A conversation ID in streaming events.""" + result = event.get("result") + + if not isinstance(result, dict): + return event + + result["contextId"] = context_id + + history = result.get("history") + if isinstance(history, list): + for message in history: + if isinstance(message, dict): + message["contextId"] = context_id + + status = result.get("status") + if isinstance(status, dict): + message = status.get("message") + if isinstance(message, dict): + message["contextId"] = context_id + + return event + + async def handle_non_streaming( + self, + request_id: str, + params: dict[str, object], + api_base: str | None = None, + **kwargs, + ) -> dict[str, object]: + litellm_params = kwargs.get("litellm_params") + + if not isinstance(litellm_params, dict): + raise ValueError("litellm_params is required") + + context_id, resolved_params = await self._resolve_session_params( + params, + litellm_params, + ) + + response = await A2ACompletionBridgeHandler.handle_non_streaming( + request_id=request_id, + params=params, + litellm_params=resolved_params, + api_base=api_base, + agent_extra_headers=kwargs.get("agent_extra_headers"), + _skip_a2a_provider_routing=True, + ) + + result = response.get("result") + if isinstance(result, dict): + result["contextId"] = context_id + + return response + + async def handle_streaming( + self, + request_id: str, + params: dict[str, object], + api_base: str | None = None, + **kwargs, + ) -> AsyncIterator[dict[str, object]]: + litellm_params = kwargs.get("litellm_params") + + if not isinstance(litellm_params, dict): + raise ValueError("litellm_params is required") + + context_id, resolved_params = await self._resolve_session_params( + params, + litellm_params, + ) + + async for event in A2ACompletionBridgeHandler.handle_streaming( + request_id=request_id, + params=params, + litellm_params=resolved_params, + api_base=api_base, + agent_extra_headers=kwargs.get("agent_extra_headers"), + _skip_a2a_provider_routing=True, + ): + yield self._set_stream_context(event, context_id) diff --git a/litellm/a2a_protocol/providers/vertex_agent_engine/session.py b/litellm/a2a_protocol/providers/vertex_agent_engine/session.py new file mode 100644 index 00000000000..7a31deba663 --- /dev/null +++ b/litellm/a2a_protocol/providers/vertex_agent_engine/session.py @@ -0,0 +1,78 @@ +"""Session helpers for Vertex AI Agent Engine A2A.""" + +import asyncio +import hashlib +import json +from typing import Protocol +from weakref import WeakKeyDictionary, WeakValueDictionary + + +def session_user_id(calling_key_hash: str, context_id: str) -> str: + """Create an isolated identity for one caller and conversation.""" + if not calling_key_hash: + raise ValueError("Authenticated calling key hash is required") + + if not context_id: + raise ValueError("A2A context ID is required") + + identity = json.dumps( + [calling_key_hash, context_id], + separators=(",", ":"), + ensure_ascii=False, + ) + + digest = hashlib.sha256(identity.encode("utf-8")).hexdigest() + return f"litellm-a2a-{digest}" + + +class AgentEngineSessionBackend(Protocol): + """Interface for Agent Engine session operations.""" + + async def list_sessions(self, *, user_id: str) -> list[str]: ... + + async def create_session(self, *, user_id: str) -> str: ... + + +# Keep locks separate across event loops and release unused locks. +_SESSION_LOCKS: WeakKeyDictionary[ + asyncio.AbstractEventLoop, + WeakValueDictionary[str, asyncio.Lock], +] = WeakKeyDictionary() + + +def _session_lock(user_id: str) -> asyncio.Lock: + """Return a shared lock for one conversation on this event loop.""" + loop = asyncio.get_running_loop() + locks = _SESSION_LOCKS.get(loop) + + if locks is None: + locks = WeakValueDictionary() + _SESSION_LOCKS[loop] = locks + + lock = locks.get(user_id) + + if lock is None: + lock = asyncio.Lock() + locks[user_id] = lock + + return lock + + +async def resolve_session_id( + backend: AgentEngineSessionBackend, + *, + calling_key_hash: str, + context_id: str, +) -> str: + """Find or create one session per caller and conversation.""" + + user_id = session_user_id(calling_key_hash, context_id) + + # Serialize lookup and creation for this conversation. + async with _session_lock(user_id): + existing_sessions = await backend.list_sessions(user_id=user_id) + + if existing_sessions: + return existing_sessions[0] + + return await backend.create_session(user_id=user_id) diff --git a/tests/unit/a2a_protocol/providers/test_vertex_agent_engine_provider.py b/tests/unit/a2a_protocol/providers/test_vertex_agent_engine_provider.py new file mode 100644 index 00000000000..940a5e56fb5 --- /dev/null +++ b/tests/unit/a2a_protocol/providers/test_vertex_agent_engine_provider.py @@ -0,0 +1,524 @@ +from litellm.a2a_protocol.providers.config_manager import ( + A2AProviderConfigManager, +) + + +def test_vertex_agent_engine_routes_to_a2a_provider(): + """Agent Engine should have its own A2A provider.""" + provider = A2AProviderConfigManager.get_provider_config( + custom_llm_provider="vertex_ai", + model="agent_engine/123456789", + ) + + assert provider is not None + + +def test_regular_vertex_model_uses_default_routing(): + """Ordinary Vertex AI models should not use Agent Engine routing.""" + provider = A2AProviderConfigManager.get_provider_config( + custom_llm_provider="vertex_ai", + model="gemini-2.5-flash", + ) + + assert provider is None + + +import pytest + + +def test_vertex_session_identity_is_stable_and_isolated(): + from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + session_user_id, + ) + + first = session_user_id("key-hash-A", "conv-1") + + # Same caller and context should produce the same identity. + assert first == session_user_id("key-hash-A", "conv-1") + + # Different callers must not share identities. + assert first != session_user_id("key-hash-B", "conv-1") + + # Different conversations must remain separate. + assert first != session_user_id("key-hash-A", "conv-2") + + # Raw identifiers must not be exposed. + assert "key-hash-A" not in first + assert "conv-1" not in first + + +def test_vertex_session_identity_requires_authenticated_key(): + from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + session_user_id, + ) + + with pytest.raises(ValueError, match="calling key"): + session_user_id("", "conv-1") + + +class FakeVertexSessionBackend: + """In-memory stand-in for Google Agent Engine session APIs.""" + + def __init__(self): + self.sessions: dict[str, list[str]] = {} + self.creations = 0 + + async def list_sessions(self, *, user_id: str) -> list[str]: + return list(self.sessions.get(user_id, [])) + + async def create_session(self, *, user_id: str) -> str: + self.creations += 1 + session_id = f"server-generated-{self.creations}" + self.sessions.setdefault(user_id, []).append(session_id) + return session_id + + +@pytest.mark.asyncio +async def test_same_conversation_reuses_vertex_session(): + from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + resolve_session_id, + ) + + backend = FakeVertexSessionBackend() + + first = await resolve_session_id( + backend, + calling_key_hash="key-A", + context_id="conv-1", + ) + + second = await resolve_session_id( + backend, + calling_key_hash="key-A", + context_id="conv-1", + ) + + assert first == second + assert backend.creations == 1 + + +@pytest.mark.asyncio +async def test_different_callers_and_contexts_have_separate_sessions(): + from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + resolve_session_id, + ) + + backend = FakeVertexSessionBackend() + + session_a = await resolve_session_id(backend, calling_key_hash="key-A", context_id="conv-1") + session_b = await resolve_session_id(backend, calling_key_hash="key-B", context_id="conv-1") + session_c = await resolve_session_id(backend, calling_key_hash="key-A", context_id="conv-2") + + assert len({session_a, session_b, session_c}) == 3 + assert backend.creations == 3 + + +@pytest.mark.asyncio +async def test_vertex_http_backend_reuses_sessions(): + """Verify Agent Engine session reuse through mocked Google HTTP calls.""" + import json + + import httpx + + from litellm.a2a_protocol.providers.vertex_agent_engine.backend import ( + VertexAgentEngineSessionBackend, + ) + from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + resolve_session_id, + ) + + sessions = {} + methods = [] + + def fake_google_api(request: httpx.Request) -> httpx.Response: + assert request.headers["authorization"] == "Bearer test-token" + assert request.url.path.endswith("/reasoningEngines/123:query") + + payload = json.loads(request.content) + method = payload["class_method"] + user_id = payload["input"]["user_id"] + methods.append(method) + + if method == "async_list_sessions": + existing = sessions.get(user_id) + items = [{"id": existing}] if existing else [] + return httpx.Response(200, json={"output": {"sessions": items}}) + + if method == "async_create_session": + session_id = f"vertex-session-{len(sessions) + 1}" + sessions[user_id] = session_id + return httpx.Response(200, json={"output": {"id": session_id}}) + + return httpx.Response(400, json={"error": "Unknown method"}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(fake_google_api)) as client: + backend = VertexAgentEngineSessionBackend( + client=client, + query_url=( + "https://us-central1-aiplatform.googleapis.com/" + "v1/projects/demo/locations/us-central1/" + "reasoningEngines/123:query" + ), + auth_headers={"Authorization": "Bearer test-token"}, + ) + + first = await resolve_session_id(backend, calling_key_hash="key-A", context_id="conv-1") + second = await resolve_session_id(backend, calling_key_hash="key-A", context_id="conv-1") + other = await resolve_session_id(backend, calling_key_hash="key-B", context_id="conv-1") + + assert first == second + assert first != other + assert len(sessions) == 2 + assert methods.count("async_create_session") == 2 + assert methods.count("async_list_sessions") == 3 + + +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.asyncio +async def test_vertex_a2a_bridge_forwards_session(monkeypatch, streaming): + """Both A2A paths must forward the resolved session to completion.""" + from types import SimpleNamespace + + import litellm + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + from litellm.a2a_protocol.providers.vertex_agent_engine.config import ( + VertexAgentEngineA2AConfig, + ) + + captured = [] + + async def fake_resolve(params, settings): + return "ctx-1", { + **settings, + "user": "isolated-user", + "user_id": "isolated-user", + "session_id": "server-session-42", + } + + async def fake_completion(completion_params): + captured.append(dict(completion_params)) + + if completion_params["stream"]: + + async def chunks(): + yield SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="pong"))]) + + return chunks() + + return litellm.ModelResponse( + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": "pong", + }, + "finish_reason": "stop", + } + ] + ) + + monkeypatch.setattr( + VertexAgentEngineA2AConfig, + "_resolve_session_params", + staticmethod(fake_resolve), + ) + monkeypatch.setattr( + A2ACompletionBridgeHandler, + "_acompletion", + staticmethod(fake_completion), + ) + + provider = VertexAgentEngineA2AConfig() + + params = { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "ping"}], + "messageId": "m-1", + "contextId": "ctx-1", + } + } + + settings = { + "custom_llm_provider": "vertex_ai", + "model": "agent_engine/123", + "litellm_a2a_user_api_key_hash": "hashed-key", + } + + if streaming: + events = [ + event + async for event in provider.handle_streaming( + request_id="req-1", + params=params, + litellm_params=settings, + ) + ] + assert len(events) == 4 + assert all(event["result"]["contextId"] == "ctx-1" for event in events) + else: + response = await provider.handle_non_streaming( + request_id="req-1", + params=params, + litellm_params=settings, + ) + assert response["result"]["contextId"] == "ctx-1" + assert response["result"]["parts"][0]["text"] == "pong" + + assert len(captured) == 1 + + completion = captured[0] + assert completion["model"] == "vertex_ai/agent_engine/123" + assert completion["user"] == "isolated-user" + assert completion["user_id"] == "isolated-user" + assert completion["session_id"] == "server-session-42" + assert completion["stream"] is streaming + + # Internal authentication metadata must not reach completion. + assert "litellm_a2a_user_api_key_hash" not in completion + + # Shared agent settings must remain unchanged. + assert "session_id" not in settings + assert "user_id" not in settings + + +def test_vertex_session_survives_completion_parameter_mapping(): + """Verify user_id and session_id reach the Vertex Agent Engine payload.""" + from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2ACompletionBridgeHandler, + ) + from litellm.llms.vertex_ai.agent_engine.transformation import ( + VertexAgentEngineConfig, + ) + from litellm.utils import get_optional_params + + a2a_params = { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Remember my name"}], + "messageId": "msg-1", + "contextId": "conversation-1", + } + } + + settings = { + "custom_llm_provider": "vertex_ai", + "model": "agent_engine/123", + "user": "scoped-user-123", + "user_id": "scoped-user-123", + "session_id": "vertex-generated-session-42", + "litellm_a2a_user_api_key_hash": "hashed-key", + } + + # Convert the A2A message into completion parameters. + completion = A2ACompletionBridgeHandler._build_completion_params( + params=a2a_params, + litellm_params=settings, + api_base=None, + agent_extra_headers=None, + stream=False, + ) + + # Simulate LiteLLM's provider parameter mapping. + optional = get_optional_params( + model="agent_engine/123", + custom_llm_provider="vertex_ai", + user=completion["user"], + user_id=completion["user_id"], + session_id=completion["session_id"], + ) + + # Build the payload that would be sent to Vertex Agent Engine. + payload = VertexAgentEngineConfig().transform_request( + model="agent_engine/123", + messages=completion["messages"], + optional_params=optional, + litellm_params={}, + headers={}, + ) + + # Verify the correct conversation identity and session are forwarded. + assert payload["input"]["user_id"] == "scoped-user-123" + assert payload["input"]["session_id"] == "vertex-generated-session-42" + assert payload["input"]["message"] == "Remember my name" + + # Internal authentication metadata must not leak to completion. + assert "litellm_a2a_user_api_key_hash" not in completion + + +@pytest.mark.asyncio +async def test_concurrent_requests_reuse_one_vertex_session(): + """Simultaneous first messages should not create duplicate sessions.""" + import asyncio + + from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + resolve_session_id, + ) + + class RacingBackend: + def __init__(self): + self.sessions = {} + self.creations = 0 + + async def list_sessions(self, *, user_id): + existing = list(self.sessions.get(user_id, [])) + + # Give the second request an opportunity to check + # the session list before either creates a session. + await asyncio.sleep(0) + return existing + + async def create_session(self, *, user_id): + self.creations += 1 + session_id = f"server-session-{self.creations}" + self.sessions.setdefault(user_id, []).append(session_id) + return session_id + + backend = RacingBackend() + + first, second = await asyncio.gather( + resolve_session_id( + backend, + calling_key_hash="same-key", + context_id="same-conversation", + ), + resolve_session_id( + backend, + calling_key_hash="same-key", + context_id="same-conversation", + ), + ) + + assert first == second + assert backend.creations == 1 + + +@pytest.mark.parametrize("status_code", [401, 403, 503]) +@pytest.mark.asyncio +async def test_vertex_session_lookup_error_does_not_create_session(status_code): + """A failed lookup must not silently create a replacement session.""" + import httpx + + from litellm.a2a_protocol.providers.vertex_agent_engine.backend import ( + VertexAgentEngineSessionBackend, + ) + from litellm.a2a_protocol.providers.vertex_agent_engine.session import ( + resolve_session_id, + ) + + requests = [] + + def fake_google_api(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + status_code, + json={"error": {"message": "Session service unavailable"}}, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(fake_google_api)) as client: + backend = VertexAgentEngineSessionBackend( + client=client, + query_url=( + "https://us-central1-aiplatform.googleapis.com/" + "v1/projects/demo/locations/us-central1/" + "reasoningEngines/123:query" + ), + auth_headers={"Authorization": "Bearer test-token"}, + ) + + with pytest.raises(httpx.HTTPStatusError): + await resolve_session_id( + backend, + calling_key_hash="key-A", + context_id="conversation-1", + ) + + # Only the failed list request should have been sent. + # A create_session request would be unsafe here. + assert len(requests) == 1 + assert b"async_list_sessions" in requests[0].content + + +@pytest.mark.asyncio +async def test_vertex_resolver_uses_google_session_api(monkeypatch): + """Test the real resolver with mocked Google HTTP responses.""" + import json + + import httpx + + from litellm.a2a_protocol.providers.vertex_agent_engine.config import ( + VertexAgentEngineA2AConfig, + ) + from litellm.llms.vertex_ai.agent_engine.transformation import ( + VertexAgentEngineConfig, + ) + + original_client = httpx.AsyncClient + requests_seen = [] + sessions = {} + + def fake_google_api(request): + assert request.headers["Authorization"] == "Bearer test-token" + assert request.url.path.endswith("/projects/test-project/locations/us-central1/reasoningEngines/123:query") + + payload = json.loads(request.content) + requests_seen.append(payload) + + method = payload["class_method"] + user_id = payload["input"]["user_id"] + + if method == "async_list_sessions": + existing = sessions.get(user_id) + return httpx.Response( + 200, + json={"output": {"sessions": [{"id": existing}] if existing else []}}, + ) + + if method == "async_create_session": + session_id = f"generated-{len(sessions) + 1}" + sessions[user_id] = session_id + return httpx.Response(200, json={"output": {"id": session_id}}) + + return httpx.Response(400) + + def mock_client(*args, **kwargs): + return original_client( + *args, + transport=httpx.MockTransport(fake_google_api), + **kwargs, + ) + + monkeypatch.setattr(httpx, "AsyncClient", mock_client) + monkeypatch.setattr( + VertexAgentEngineConfig, + "_get_auth_headers", + lambda self, optional_params, litellm_params: {"Authorization": "Bearer test-token"}, + ) + + params = {"message": {"contextId": "conversation-1"}} + settings = { + "custom_llm_provider": "vertex_ai", + "model": "agent_engine/123", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "litellm_a2a_user_api_key_hash": "hashed-key", + } + + first_context, first = await VertexAgentEngineA2AConfig._resolve_session_params(params, settings) + second_context, second = await VertexAgentEngineA2AConfig._resolve_session_params(params, settings) + + assert first_context == second_context == "conversation-1" + assert first["session_id"] == second["session_id"] + assert first["session_id"] == "generated-1" + assert first["user_id"] == second["user_id"] + assert first["user"] == first["user_id"] + assert "session_id" not in settings + + assert [request["class_method"] for request in requests_seen] == [ + "async_list_sessions", + "async_create_session", + "async_list_sessions", + ]