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.
This commit is contained in:
GordonWei 2026-11-01 23:57:39 +08:00
parent 71da0f70b4
commit d2f249dc9a
2 changed files with 25 additions and 1 deletions

View file

@ -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},
}

View file

@ -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):