From 71da0f70b4efb52db32519a9a14c4f5a9d18d465 Mon Sep 17 00:00:00 2001 From: GordonWei Date: Sun, 1 Nov 2026 17:33:53 +0800 Subject: [PATCH 1/2] feat(a2a): continue Vertex AI Agent Engine sessions across A2A turns A2A calls to a Vertex AI Agent Engine agent went through the stateless completion bridge, so every turn reached Agent Engine with a new user_id and no session_id. Multi-turn agents never got past the first turn. Add a VertexAgentEngineA2AConfig provider (routed for vertex_ai models containing agent_engine) that, before calling the completion bridge: - derives the Agent Engine user_id from the calling key hash and the A2A contextId, so two keys sending the same contextId never share a session - looks up that user's latest Agent Engine session and passes its id; on the first turn there is none and Agent Engine creates one Agent Engine only accepts session ids it generated, so the contextId is never sent as session_id. Missing key hash or contextId, or a failed lookup, leaves litellm_params unchanged (current behaviour). An explicit user_id / session_id in the agent's litellm_params is respected. Closes #49451 --- .../a2a_protocol/providers/config_manager.py | 7 + .../providers/vertex_agent_engine/__init__.py | 6 + .../providers/vertex_agent_engine/config.py | 73 +++++ litellm/llms/vertex_ai/agent_engine/a2a.py | 162 +++++++++++ .../agent_engine/test_agent_engine_a2a.py | 262 ++++++++++++++++++ 5 files changed, 510 insertions(+) create mode 100644 litellm/a2a_protocol/providers/vertex_agent_engine/__init__.py create mode 100644 litellm/a2a_protocol/providers/vertex_agent_engine/config.py create mode 100644 litellm/llms/vertex_ai/agent_engine/a2a.py create mode 100644 tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py diff --git a/litellm/a2a_protocol/providers/config_manager.py b/litellm/a2a_protocol/providers/config_manager.py index 2eab2adb1ba..32aff201038 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 "agent_engine" in model: + 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/__init__.py b/litellm/a2a_protocol/providers/vertex_agent_engine/__init__.py new file mode 100644 index 00000000000..8d9069a56cd --- /dev/null +++ b/litellm/a2a_protocol/providers/vertex_agent_engine/__init__.py @@ -0,0 +1,6 @@ +""" +Vertex AI Agent Engine A2A provider. + +Maps the A2A contextId onto Agent Engine user/session ids so multi-turn agents +keep their session across A2A calls, then uses the completion bridge. +""" 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..62bc6e10b67 --- /dev/null +++ b/litellm/a2a_protocol/providers/vertex_agent_engine/config.py @@ -0,0 +1,73 @@ +from collections.abc import AsyncIterator, Mapping +from typing import Final + +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.llms.vertex_ai.agent_engine.a2a import merge_a2a_session_into_litellm_params + + +def _require_litellm_params(litellm_params: object) -> Mapping[str, object]: + if not isinstance(litellm_params, Mapping) or not litellm_params: + raise ValueError( + "litellm_params is required for VertexAgentEngineA2AConfig (must contain custom_llm_provider and model)" + ) + return {str(key): value for key, value in litellm_params.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # validated untyped kwargs boundary + + +def _principal(litellm_params: Mapping[str, object]) -> str | None: + principal: Final = litellm_params.get(A2A_USER_API_KEY_HASH_PARAM) + return principal if isinstance(principal, str) else None + + +def _agent_extra_headers( + value: object, +) -> dict[str, str] | None: # mutable-ok: the bridge handler signature takes a dict + if not isinstance(value, Mapping): + return None + return {str(k): str(v) for k, v in value.items()} # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # validated untyped kwargs boundary + + +class VertexAgentEngineA2AConfig(BaseA2AProviderConfig): + """A2A bridge for Vertex AI Agent Engine: continues the caller's Agent Engine + session for the same contextId (scoped to the authenticated key), then uses + completion.""" + + async def handle_non_streaming( + self, + request_id: str, + params: dict[str, object], # mutable-ok: signature fixed by BaseA2AProviderConfig + api_base: str | None = None, + **kwargs: object, # kwargs-ok: signature fixed by BaseA2AProviderConfig + ) -> dict[str, object]: # mutable-ok: signature fixed by BaseA2AProviderConfig + litellm_params: Final = _require_litellm_params(kwargs.get("litellm_params")) + merged: Final = await merge_a2a_session_into_litellm_params(litellm_params, params, _principal(litellm_params)) + return await A2ACompletionBridgeHandler.handle_non_streaming( + request_id=request_id, + params=params, + litellm_params=dict(merged), + api_base=api_base, + agent_extra_headers=_agent_extra_headers(kwargs.get("agent_extra_headers")), + _skip_a2a_provider_routing=True, + ) + + async def handle_streaming( + self, + request_id: str, + params: dict[str, object], # mutable-ok: signature fixed by BaseA2AProviderConfig + api_base: str | None = None, + **kwargs: object, # kwargs-ok: signature fixed by BaseA2AProviderConfig + ) -> AsyncIterator[dict[str, object]]: # mutable-ok: signature fixed by BaseA2AProviderConfig + litellm_params: Final = _require_litellm_params(kwargs.get("litellm_params")) + merged: Final = await merge_a2a_session_into_litellm_params(litellm_params, params, _principal(litellm_params)) + async for chunk in A2ACompletionBridgeHandler.handle_streaming( + request_id=request_id, + params=params, + litellm_params=dict(merged), + api_base=api_base, + agent_extra_headers=_agent_extra_headers(kwargs.get("agent_extra_headers")), + _skip_a2a_provider_routing=True, + ): + yield chunk diff --git a/litellm/llms/vertex_ai/agent_engine/a2a.py b/litellm/llms/vertex_ai/agent_engine/a2a.py new file mode 100644 index 00000000000..4a51735648d --- /dev/null +++ b/litellm/llms/vertex_ai/agent_engine/a2a.py @@ -0,0 +1,162 @@ +""" +A2A session bridging for Vertex AI Agent Engine. + +The completion bridge is stateless: every A2A call reaches Agent Engine with a +fresh ``user_id`` and no ``session_id``, so a multi-turn agent never gets past +the first turn. This module maps the A2A ``message.contextId`` onto Agent +Engine's ``user_id`` / ``session_id`` so turns with the same contextId continue +the same session. + +- ``user_id`` is derived from the calling key hash plus the contextId, so two + keys sending the same contextId never share a session (same principle as the + AgentCore and LangFlow bridges). +- ``session_id`` cannot be the contextId: Agent Engine only accepts session ids + it generated itself, and an unknown id fails silently (empty response). So we + look up the caller's most recent session for that ``user_id``; on the first + turn there is none and Agent Engine creates one. + +Anything missing or failing (no key hash, no contextId, lookup error) leaves +``litellm_params`` untouched, i.e. the current stateless behaviour. +""" + +import hashlib +from collections.abc import Mapping +from typing import Final + +from pydantic import BaseModel, ConfigDict, Field + +from litellm._logging import verbose_logger +from litellm.a2a_protocol.utils import get_session_id_from_a2a_params +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared helper is loosely typed upstream +) +from litellm.llms.vertex_ai.agent_engine.transformation import VertexAgentEngineConfig +from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.types.llms.custom_http import httpxSpecialProvider + +A2A_AGENT_ENGINE_USER_ID_PREFIX: Final = "litellm-a2a-" +_USER_ID_HASH_LENGTH: Final = 16 +_SESSION_LOOKUP_TIMEOUT_SECONDS: Final = 5.0 +_SESSION_LOOKUP_PAGE_SIZE: Final = 100 +_DEFAULT_LOCATION: Final = "us-central1" + + +class _AgentEngineSession(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + name: str = "" + user_id: str = Field(default="", alias="userId") + update_time: str = Field(default="", alias="updateTime") + + +class _ListSessionsResponse(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + sessions: tuple[_AgentEngineSession, ...] = () + + +def derive_agent_engine_user_id(context_id: str | None, principal: str | None) -> str | None: + """Stable, per-caller Agent Engine ``user_id`` for an A2A conversation. + + Returns None when either part is missing: without a trusted principal any + caller could guess another caller's contextId, and without a contextId + every request from a key would collapse into a single session. + """ + if not context_id or not principal: + return None + digest: Final = hashlib.sha256(f"{principal}:{context_id}".encode()).hexdigest() + return f"{A2A_AGENT_ENGINE_USER_ID_PREFIX}{digest[:_USER_ID_HASH_LENGTH]}" + + +def _str_param(litellm_params: Mapping[str, object], key: str) -> str | None: + value: Final = litellm_params.get(key) + return value if isinstance(value, str) and value else None + + +def _resolve_engine(litellm_params: Mapping[str, object]) -> tuple[str, str] | None: + """Return (reasoningEngines resource path, location) for the agent's model. + + Accepts the same model forms as ``VertexAgentEngineConfig``: + ``agent_engine/projects/

/locations//reasoningEngines/`` or + ``agent_engine/`` with ``vertex_project`` / ``vertex_location``. + """ + model: Final = _str_param(litellm_params, "model") + if model is None or "agent_engine/" not in model: + return None + engine: Final = model.split("agent_engine/", 1)[1] + if engine.startswith("projects/"): + parts: Final = engine.split("/") + if len(parts) != 6 or parts[2] != "locations" or parts[4] != "reasoningEngines": + return None + return engine, parts[3] + + if "/" in engine: + return None + vertex_project: Final = VertexAgentEngineConfig.safe_get_vertex_ai_project(dict(litellm_params)) # pyright: ignore[reportUnknownMemberType] # VertexBase helpers take an untyped dict + if not vertex_project: + return None + vertex_location: Final = VertexAgentEngineConfig.safe_get_vertex_ai_location(litellm_params) or _DEFAULT_LOCATION + return ( + f"projects/{vertex_project}/locations/{vertex_location}/reasoningEngines/{engine}", + vertex_location, + ) + + +async def lookup_latest_session_id(litellm_params: Mapping[str, object], user_id: str) -> str | None: + """Most recently updated Agent Engine session of ``user_id``, or None. + + None also covers lookup failures: the caller then lets Agent Engine create a + new session instead of blocking the request. + """ + engine: Final = _resolve_engine(litellm_params) + if engine is None: + return None + resource_path, location = engine + + try: + config: Final = VertexAgentEngineConfig() + access_token, _project = await config.get_access_token_async( + credentials=config.safe_get_vertex_ai_credentials(dict(litellm_params)), # pyright: ignore[reportUnknownMemberType] # VertexBase helpers take an untyped dict + project_id=config.safe_get_vertex_ai_project(dict(litellm_params)), # pyright: ignore[reportUnknownMemberType] # VertexBase helpers take an untyped dict + ) + client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.A2AProvider) + response: Final = await client.get( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.get takes untyped dicts + f"{get_vertex_base_url(location)}/v1beta1/{resource_path}/sessions", + headers={"Authorization": f"Bearer {access_token}"}, + params={"filter": f'user_id="{user_id}"', "pageSize": _SESSION_LOOKUP_PAGE_SIZE}, + timeout=_SESSION_LOOKUP_TIMEOUT_SECONDS, + ) + _ = response.raise_for_status() + listing: Final = _ListSessionsResponse.model_validate_json(response.text) + except Exception as e: + verbose_logger.warning( + "Vertex Agent Engine A2A: session lookup failed, a new session will be created: %s", + e, + ) + return None + + # The server-side filter is not trusted alone; re-check the owner here. + owned: Final = tuple(s for s in listing.sessions if s.user_id == user_id and s.name) + if not owned: + return None + latest: Final = max(owned, key=lambda s: s.update_time) + return latest.name.rsplit("/", 1)[-1] or None + + +async def merge_a2a_session_into_litellm_params( + litellm_params: Mapping[str, object], + params: Mapping[str, object], + principal: str | None = None, +) -> Mapping[str, object]: + """``litellm_params`` with Agent Engine ``user_id`` / ``session_id`` set from the A2A context.""" + if litellm_params.get("user_id") or litellm_params.get("user"): + return litellm_params + + user_id: Final = derive_agent_engine_user_id(get_session_id_from_a2a_params(params), principal) + if user_id is None: + return litellm_params + + existing_session: Final = _str_param(litellm_params, "session_id") + session_id: Final = existing_session or await lookup_latest_session_id(litellm_params, user_id) + session_update: Final[Mapping[str, object]] = {"session_id": session_id} if session_id else {} + return {**litellm_params, "user_id": user_id, **session_update} diff --git a/tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py b/tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py new file mode 100644 index 00000000000..7414c55be45 --- /dev/null +++ b/tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py @@ -0,0 +1,262 @@ +""" +Tests for A2A session bridging to Vertex AI Agent Engine. + +No real API calls: the Agent Engine sessions endpoint and the access token are mocked. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from litellm.a2a_protocol.litellm_completion_bridge.handler import ( + A2A_USER_API_KEY_HASH_PARAM, +) +from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager +from litellm.a2a_protocol.providers.vertex_agent_engine.config import ( + VertexAgentEngineA2AConfig, +) +from litellm.llms.vertex_ai.agent_engine import a2a as agent_engine_a2a +from litellm.llms.vertex_ai.agent_engine.a2a import ( + A2A_AGENT_ENGINE_USER_ID_PREFIX, + derive_agent_engine_user_id, + lookup_latest_session_id, + merge_a2a_session_into_litellm_params, +) + +ENGINE_PATH = "projects/test-project/locations/test-region/reasoningEngines/123" +LITELLM_PARAMS = { + "custom_llm_provider": "vertex_ai", + "model": f"agent_engine/{ENGINE_PATH}", +} + + +def _params(context_id: str | None = "ctx-1") -> dict: + message: dict = {"role": "user", "parts": [{"kind": "text", "text": "hi"}]} + if context_id is not None: + message["contextId"] = context_id + return {"message": message} + + +def _mock_sessions_api(sessions: list[dict] | None = None, error: Exception | None = None): + """Patch the token and the HTTP client used by the session lookup.""" + response = httpx.Response( + 200, + json={"sessions": sessions or []}, + request=httpx.Request("GET", "https://test-region-aiplatform.googleapis.com"), + ) + client = MagicMock() + client.get = AsyncMock(side_effect=error) if error else AsyncMock(return_value=response) + return ( + patch.object( + agent_engine_a2a.VertexAgentEngineConfig, + "get_access_token_async", + AsyncMock(return_value=("token", "test-project")), + ), + patch.object(agent_engine_a2a, "get_async_httpx_client", return_value=client), + client, + ) + + +class TestDeriveUserId: + def test_same_key_and_context_is_stable(self): + assert derive_agent_engine_user_id("ctx-1", "hash-a") == derive_agent_engine_user_id("ctx-1", "hash-a") + + def test_same_key_different_context_is_isolated(self): + assert derive_agent_engine_user_id("ctx-1", "hash-a") != derive_agent_engine_user_id("ctx-2", "hash-a") + + def test_different_keys_same_context_are_isolated(self): + """Two keys reusing a contextId must not share an Agent Engine session.""" + assert derive_agent_engine_user_id("ctx-1", "hash-a") != derive_agent_engine_user_id("ctx-1", "hash-b") + + def test_principal_and_context_are_not_sent_verbatim(self): + user_id = derive_agent_engine_user_id("ctx-1", "hash-a") + assert user_id is not None + assert user_id.startswith(A2A_AGENT_ENGINE_USER_ID_PREFIX) + assert "hash-a" not in user_id + assert "ctx-1" not in user_id + + @pytest.mark.parametrize("context_id,principal", [(None, "hash-a"), ("", "hash-a"), ("ctx-1", None), ("ctx-1", "")]) + def test_missing_part_returns_none(self, context_id, principal): + assert derive_agent_engine_user_id(context_id, principal) is None + + +class TestLookupLatestSessionId: + @pytest.mark.asyncio + async def test_returns_latest_session_of_user(self): + user_id = derive_agent_engine_user_id("ctx-1", "hash-a") + sessions = [ + {"name": f"{ENGINE_PATH}/sessions/111", "userId": user_id, "updateTime": "2026-01-01T00:00:00Z"}, + {"name": f"{ENGINE_PATH}/sessions/222", "userId": user_id, "updateTime": "2026-01-02T00:00:00Z"}, + ] + token_patch, client_patch, client = _mock_sessions_api(sessions) + with token_patch, client_patch: + assert await lookup_latest_session_id(LITELLM_PARAMS, user_id) == "222" + + url = client.get.call_args.args[0] + assert url == f"https://test-region-aiplatform.googleapis.com/v1beta1/{ENGINE_PATH}/sessions" + assert client.get.call_args.kwargs["params"]["filter"] == f'user_id="{user_id}"' + + @pytest.mark.asyncio + async def test_ignores_sessions_of_other_users(self): + user_id = derive_agent_engine_user_id("ctx-1", "hash-a") + other = derive_agent_engine_user_id("ctx-1", "hash-b") + sessions = [{"name": f"{ENGINE_PATH}/sessions/999", "userId": other, "updateTime": "2026-01-03T00:00:00Z"}] + token_patch, client_patch, _ = _mock_sessions_api(sessions) + with token_patch, client_patch: + assert await lookup_latest_session_id(LITELLM_PARAMS, user_id) is None + + @pytest.mark.asyncio + async def test_lookup_error_returns_none(self): + token_patch, client_patch, _ = _mock_sessions_api(error=RuntimeError("boom")) + with token_patch, client_patch: + assert await lookup_latest_session_id(LITELLM_PARAMS, "litellm-a2a-x") is None + + @pytest.mark.asyncio + async def test_engine_id_only_model_uses_vertex_project_and_location(self): + params = { + "custom_llm_provider": "vertex_ai", + "model": "agent_engine/456", + "vertex_project": "other-project", + "vertex_location": "other-region", + } + token_patch, client_patch, client = _mock_sessions_api([]) + with token_patch, client_patch: + await lookup_latest_session_id(params, "litellm-a2a-x") + assert client.get.call_args.args[0] == ( + "https://other-region-aiplatform.googleapis.com/v1beta1/" + "projects/other-project/locations/other-region/reasoningEngines/456/sessions" + ) + + @pytest.mark.asyncio + async def test_non_agent_engine_model_skips_lookup(self): + token_patch, client_patch, client = _mock_sessions_api([]) + with token_patch, client_patch: + result = await lookup_latest_session_id( + {"custom_llm_provider": "vertex_ai", "model": "gemini-pro"}, "litellm-a2a-x" + ) + assert result is None + client.get.assert_not_called() + + +class TestMergeSessionIntoLitellmParams: + @pytest.mark.asyncio + async def test_first_turn_sets_user_id_only(self): + token_patch, client_patch, _ = _mock_sessions_api([]) + with token_patch, client_patch: + merged = await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(), "hash-a") + assert merged["user_id"] == derive_agent_engine_user_id("ctx-1", "hash-a") + assert "session_id" not in merged + + @pytest.mark.asyncio + async def test_next_turn_reuses_existing_session(self): + user_id = derive_agent_engine_user_id("ctx-1", "hash-a") + sessions = [{"name": f"{ENGINE_PATH}/sessions/777", "userId": user_id, "updateTime": "2026-01-01T00:00:00Z"}] + token_patch, client_patch, _ = _mock_sessions_api(sessions) + with token_patch, client_patch: + merged = await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(), "hash-a") + assert merged["user_id"] == user_id + assert merged["session_id"] == "777" + + @pytest.mark.asyncio + async def test_lookup_failure_falls_back_to_new_session(self): + token_patch, client_patch, _ = _mock_sessions_api(error=RuntimeError("boom")) + with token_patch, client_patch: + merged = await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(), "hash-a") + assert "session_id" not in merged + + @pytest.mark.asyncio + @pytest.mark.parametrize("context_id,principal", [(None, "hash-a"), ("ctx-1", None)]) + async def test_without_context_or_principal_is_noop(self, context_id, principal): + token_patch, client_patch, client = _mock_sessions_api([]) + with token_patch, client_patch: + merged = await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(context_id), principal) + assert merged == LITELLM_PARAMS + client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_explicit_user_id_is_respected(self): + token_patch, client_patch, client = _mock_sessions_api([]) + with token_patch, client_patch: + merged = await merge_a2a_session_into_litellm_params( + {**LITELLM_PARAMS, "user_id": "configured-user"}, _params(), "hash-a" + ) + assert merged["user_id"] == "configured-user" + client.get.assert_not_called() + + @pytest.mark.asyncio + async def test_input_params_are_not_mutated(self): + original = dict(LITELLM_PARAMS) + token_patch, client_patch, _ = _mock_sessions_api([]) + with token_patch, client_patch: + await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(), "hash-a") + assert LITELLM_PARAMS == original + + +class TestVertexAgentEngineA2AConfig: + def test_registered_for_agent_engine_models(self): + cfg = A2AProviderConfigManager.get_provider_config( + custom_llm_provider="vertex_ai", model=f"agent_engine/{ENGINE_PATH}" + ) + assert isinstance(cfg, VertexAgentEngineA2AConfig) + + def test_not_registered_for_other_vertex_models(self): + assert A2AProviderConfigManager.get_provider_config(custom_llm_provider="vertex_ai", model="gemini-pro") is None + + @pytest.mark.asyncio + async def test_non_streaming_forwards_session_to_bridge(self): + litellm_params = {**LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: "hash-a"} + with ( + patch( + "litellm.a2a_protocol.providers.vertex_agent_engine.config.merge_a2a_session_into_litellm_params", + AsyncMock(return_value={**litellm_params, "user_id": "u", "session_id": "s"}), + ) as merge, + patch( + "litellm.a2a_protocol.providers.vertex_agent_engine.config.A2ACompletionBridgeHandler.handle_non_streaming", + AsyncMock(return_value={"ok": True}), + ) as bridge, + ): + result = await VertexAgentEngineA2AConfig().handle_non_streaming( + request_id="r1", params=_params(), litellm_params=litellm_params + ) + assert result == {"ok": True} + merge.assert_awaited_once() + assert merge.call_args.args[2] == "hash-a" + kwargs = bridge.call_args.kwargs + assert kwargs["litellm_params"]["session_id"] == "s" + assert kwargs["_skip_a2a_provider_routing"] is True + + @pytest.mark.asyncio + async def test_streaming_forwards_session_to_bridge(self): + litellm_params = {**LITELLM_PARAMS, A2A_USER_API_KEY_HASH_PARAM: "hash-a"} + + async def fake_stream(**kwargs): + yield {"chunk": 1} + yield {"chunk": 2} + + bridge = MagicMock(side_effect=fake_stream) + with ( + patch( + "litellm.a2a_protocol.providers.vertex_agent_engine.config.merge_a2a_session_into_litellm_params", + AsyncMock(return_value={**litellm_params, "user_id": "u", "session_id": "s"}), + ), + patch( + "litellm.a2a_protocol.providers.vertex_agent_engine.config.A2ACompletionBridgeHandler.handle_streaming", + bridge, + ), + ): + chunks = [ + c + async for c in VertexAgentEngineA2AConfig().handle_streaming( + request_id="r1", params=_params(), litellm_params=litellm_params + ) + ] + assert chunks == [{"chunk": 1}, {"chunk": 2}] + kwargs = bridge.call_args.kwargs + assert kwargs["litellm_params"]["session_id"] == "s" + assert kwargs["_skip_a2a_provider_routing"] is True + + @pytest.mark.asyncio + async def test_missing_litellm_params_raises(self): + with pytest.raises(ValueError, match="litellm_params is required"): + await VertexAgentEngineA2AConfig().handle_non_streaming(request_id="r1", params=_params()) From d2f249dc9ab36ab3029f27547d83d635dafc2e9d Mon Sep 17 00:00:00 2001 From: GordonWei Date: Sun, 1 Nov 2026 23:57:39 +0800 Subject: [PATCH 2/2] fix(a2a): bypass response cache for Agent Engine session-backed calls The reply depends on server-side session history, but the response cache key ignores user_id/session_id, so a cached reply could be served to another key or to a later turn of the same conversation. --- litellm/llms/vertex_ai/agent_engine/a2a.py | 9 ++++++++- .../agent_engine/test_agent_engine_a2a.py | 17 +++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/litellm/llms/vertex_ai/agent_engine/a2a.py b/litellm/llms/vertex_ai/agent_engine/a2a.py index 4a51735648d..51e8eacc113 100644 --- a/litellm/llms/vertex_ai/agent_engine/a2a.py +++ b/litellm/llms/vertex_ai/agent_engine/a2a.py @@ -159,4 +159,11 @@ async def merge_a2a_session_into_litellm_params( existing_session: Final = _str_param(litellm_params, "session_id") session_id: Final = existing_session or await lookup_latest_session_id(litellm_params, user_id) session_update: Final[Mapping[str, object]] = {"session_id": session_id} if session_id else {} - return {**litellm_params, "user_id": user_id, **session_update} + # The reply depends on server-side session history, but the response cache key ignores + # user_id/session_id, so a cached reply could be served to another key or a later turn. + return { + **litellm_params, + "user_id": user_id, + **session_update, + "cache": {"no-cache": True, "no-store": True}, + } diff --git a/tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py b/tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py index 7414c55be45..da47923e478 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py @@ -192,6 +192,23 @@ class TestMergeSessionIntoLitellmParams: await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(), "hash-a") assert LITELLM_PARAMS == original + @pytest.mark.asyncio + @pytest.mark.parametrize("sessions", [[], None]) + async def test_session_backed_call_bypasses_response_cache(self, sessions): + user_id = derive_agent_engine_user_id("ctx-1", "hash-a") + existing = [{"name": f"{ENGINE_PATH}/sessions/777", "userId": user_id, "updateTime": "2026-01-01T00:00:00Z"}] + token_patch, client_patch, _ = _mock_sessions_api(existing if sessions is None else sessions) + with token_patch, client_patch: + merged = await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(), "hash-a") + assert merged["cache"] == {"no-cache": True, "no-store": True} + + @pytest.mark.asyncio + async def test_noop_path_leaves_cache_settings_alone(self): + token_patch, client_patch, _ = _mock_sessions_api([]) + with token_patch, client_patch: + merged = await merge_a2a_session_into_litellm_params(LITELLM_PARAMS, _params(None), "hash-a") + assert "cache" not in merged + class TestVertexAgentEngineA2AConfig: def test_registered_for_agent_engine_models(self):