From 9f5dbc81d8538f7b52ade43b85221d8b9f4fe728 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 6 May 2026 16:05:04 -0700 Subject: [PATCH] test(warm-pool): unit tests for hydrate payload builder + scope filtering --- .../warm_pool/test_hydrate.py | 284 ++++++++++++++++++ 1 file changed, 284 insertions(+) create mode 100644 tests/test_litellm/proxy/agent_session_endpoints/warm_pool/test_hydrate.py diff --git a/tests/test_litellm/proxy/agent_session_endpoints/warm_pool/test_hydrate.py b/tests/test_litellm/proxy/agent_session_endpoints/warm_pool/test_hydrate.py new file mode 100644 index 00000000000..cb7f65a5de7 --- /dev/null +++ b/tests/test_litellm/proxy/agent_session_endpoints/warm_pool/test_hydrate.py @@ -0,0 +1,284 @@ +"""Unit tests for `litellm/proxy/agent_session_endpoints/warm_pool/hydrate.py`. + +Verify the payload builder: +* respects scope filtering — out-of-scope secrets never land in the payload +* tolerates a missing `LiteLLM_AgentVMConfig` row (defaults to allow_all) +* falls back to a sensible model when the agent has none +* stringifies env_vars correctly so the daemon's env file is well-formed +""" + +from __future__ import annotations + +import os +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional +from unittest.mock import patch + +import pytest + + +# Tests need a stable salt key BEFORE the encryption module is imported by +# the production code we're calling. +os.environ.setdefault("LITELLM_SALT_KEY", "sk-test-salt-key-do-not-use-in-prod") + + +from litellm.proxy.agent_session_endpoints.warm_pool.hydrate import ( + build_hydrate_payload, +) +from litellm.proxy.agent_settings_endpoints.encryption import encrypt_optional + + +class _Row: + def __init__(self, **fields: Any) -> None: + for k, v in fields.items(): + setattr(self, k, v) + + +class _Table: + def __init__(self, rows: Optional[List[_Row]] = None) -> None: + self.rows = rows or [] + + async def find_many(self, where=None) -> List[_Row]: + if where is None: + return list(self.rows) + return [ + r + for r in self.rows + if all(getattr(r, k, None) == v for k, v in where.items()) + ] + + async def find_unique(self, where: Dict[str, Any]) -> Optional[_Row]: + for r in self.rows: + if all(getattr(r, k, None) == v for k, v in where.items()): + return r + return None + + +class _DB: + def __init__(self) -> None: + self.litellm_agentsecret = _Table() + self.litellm_agentvmconfig = _Table() + + +class _Prisma: + def __init__(self) -> None: + self.db = _DB() + + +@pytest.fixture +def prisma(): + return _Prisma() + + +@pytest.fixture +def jwt_expires_at(): + return datetime(2026, 5, 7, 0, 0, 0, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +async def test_payload_with_no_secrets_no_config(prisma, jwt_expires_at): + """No secrets, no VMConfig row -> defaults all the way through.""" + payload = await build_hydrate_payload( + prisma_client=prisma, + session_id="sess-1", + agent_id="agent-1", + team_id="team-A", + jwt="jwt-x", + jwt_expires_at=jwt_expires_at, + repos=[{"url": "https://github.com/a/b", "ref": "main"}], + env_vars={"NODE_ENV": "test"}, + agent_row=_Row(model="gpt-4o", system_prompt="hi"), + ) + + assert payload.session_id == "sess-1" + assert payload.agent_id == "agent-1" + assert payload.jwt == "jwt-x" + assert payload.jwt_expires_at == jwt_expires_at.isoformat() + assert payload.secrets == {} + assert payload.network_access.mode == "allow_all" + assert payload.agent_config.model == "gpt-4o" + assert payload.agent_config.system_prompt == "hi" + assert len(payload.repos) == 1 + assert payload.repos[0].url == "https://github.com/a/b" + assert payload.env_vars == {"NODE_ENV": "test"} + + +@pytest.mark.asyncio +async def test_in_scope_secrets_decrypted(prisma, jwt_expires_at): + """`scope='all'` and matching repo scope land in the payload.""" + prisma.db.litellm_agentsecret.rows = [ + _Row( + team_id="team-A", + name="GLOBAL_TOKEN", + value_enc=encrypt_optional("global-secret-value"), + scope="all", + ), + _Row( + team_id="team-A", + name="REPO_TOKEN", + value_enc=encrypt_optional("repo-specific-value"), + scope=["foo/bar"], + ), + ] + + payload = await build_hydrate_payload( + prisma_client=prisma, + session_id="sess-1", + agent_id="agent-1", + team_id="team-A", + jwt="jwt", + jwt_expires_at=jwt_expires_at, + repos=[{"url": "https://github.com/foo/bar"}], + env_vars=None, + agent_row=None, + ) + + assert payload.secrets == { + "GLOBAL_TOKEN": "global-secret-value", + "REPO_TOKEN": "repo-specific-value", + } + + +@pytest.mark.asyncio +async def test_out_of_scope_secrets_excluded(prisma, jwt_expires_at): + """A secret scoped to repo `x/y` is dropped when session repos don't match.""" + prisma.db.litellm_agentsecret.rows = [ + _Row( + team_id="team-A", + name="SECRET_FOR_X", + value_enc=encrypt_optional("must-not-leak"), + scope=["x/y"], + ), + ] + + payload = await build_hydrate_payload( + prisma_client=prisma, + session_id="sess-1", + agent_id="agent-1", + team_id="team-A", + jwt="jwt", + jwt_expires_at=jwt_expires_at, + repos=[{"url": "https://github.com/foo/bar"}], + env_vars=None, + agent_row=None, + ) + + assert "SECRET_FOR_X" not in payload.secrets + + +@pytest.mark.asyncio +async def test_other_team_secrets_excluded(prisma, jwt_expires_at): + """A secret on a different team_id is invisible to this team.""" + prisma.db.litellm_agentsecret.rows = [ + _Row( + team_id="other-team", + name="OTHER_TEAM_SECRET", + value_enc=encrypt_optional("must-not-leak"), + scope="all", + ), + ] + + payload = await build_hydrate_payload( + prisma_client=prisma, + session_id="sess-1", + agent_id="agent-1", + team_id="team-A", + jwt="jwt", + jwt_expires_at=jwt_expires_at, + repos=[], + env_vars=None, + agent_row=None, + ) + + assert payload.secrets == {} + + +@pytest.mark.asyncio +async def test_network_access_loaded_from_vm_config(prisma, jwt_expires_at): + """`LiteLLM_AgentVMConfig.network_access` is reflected in the payload.""" + prisma.db.litellm_agentvmconfig.rows = [ + _Row( + team_id="team-A", + network_access={ + "mode": "allowlist", + "allowlist": ["github.com", "api.openai.com"], + }, + ), + ] + + payload = await build_hydrate_payload( + prisma_client=prisma, + session_id="sess-1", + agent_id="agent-1", + team_id="team-A", + jwt="jwt", + jwt_expires_at=jwt_expires_at, + repos=[], + env_vars=None, + agent_row=None, + ) + + assert payload.network_access.mode == "allowlist" + assert payload.network_access.allowlist == ["github.com", "api.openai.com"] + + +@pytest.mark.asyncio +async def test_decrypt_failure_skips_secret_does_not_crash(prisma, jwt_expires_at): + """A bad ciphertext is skipped — payload still builds, secret is absent.""" + good_ciphertext = encrypt_optional("good-value") + bad_ciphertext = "not-valid-ciphertext" + prisma.db.litellm_agentsecret.rows = [ + _Row( + team_id="team-A", + name="GOOD_SECRET", + value_enc=good_ciphertext, + scope="all", + ), + _Row( + team_id="team-A", + name="BAD_SECRET", + value_enc=bad_ciphertext, + scope="all", + ), + ] + + def fake_decrypt(value, key): + if value == good_ciphertext: + return "good-value" + raise ValueError("simulated decrypt fail") + + with patch( + "litellm.proxy.agent_session_endpoints.warm_pool.hydrate.decrypt_optional", + new=fake_decrypt, + ): + payload = await build_hydrate_payload( + prisma_client=prisma, + session_id="sess-1", + agent_id="agent-1", + team_id="team-A", + jwt="jwt", + jwt_expires_at=jwt_expires_at, + repos=[], + env_vars=None, + agent_row=None, + ) + + assert "GOOD_SECRET" in payload.secrets + assert "BAD_SECRET" not in payload.secrets + + +@pytest.mark.asyncio +async def test_default_model_when_agent_row_missing(prisma, jwt_expires_at): + """Builder uses 'gpt-4o-mini' when agent_row is None and no model passed.""" + payload = await build_hydrate_payload( + prisma_client=prisma, + session_id="sess-1", + agent_id="agent-1", + team_id="team-A", + jwt="jwt", + jwt_expires_at=jwt_expires_at, + repos=[], + env_vars=None, + agent_row=None, + ) + assert payload.agent_config.model == "gpt-4o-mini"