Merge d2f249dc9a into 527d882508

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
 

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.
+"""

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

litellm/llms/vertex_ai/agent_engine/a2a.py

@@ -0,0 +1,169 @@
+"""
+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/<p>/locations/<l>/reasoningEngines/<id>`` or
+    ``agent_engine/<id>`` 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 {}
+    # 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},
+    }

tests/unit/llms/vertex_ai/agent_engine/test_agent_engine_a2a.py

@@ -0,0 +1,279 @@
+"""
+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
+
+    @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):
+        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())