mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
fix(a2a): preserve Vertex Agent Engine sessions across messages
This commit is contained in:
Senthurapandi
parent
3129e44848
commit
59ddffe83e
5 changed files with 858 additions and 0 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
175
litellm/a2a_protocol/providers/vertex_agent_engine/config.py
Normal file
175
litellm/a2a_protocol/providers/vertex_agent_engine/config.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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)
|
||||
|
|
@ -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",
|
||||
]
|
||||
Loading…
Add table
Reference in a new issue