scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 227

fix(proxy): bind JWT agent lookup at startup regardless of agent source
Some checks failed
ai-gateway image / ai-gateway release image (push) Has been cancelled
Details
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details

Browse source 
Move jwt_handler.bind_agent_lookup out of the YAML and DB agent loading paths and into
ProxyStartupEvent._initialize_jwt_auth so agents created via the API or UI after startup,
with no agents in config and no DB agent reload, still resolve for agent_id_jwt_field

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-14 21:38:00 +00:00
parent 4d4d3fb18a
commit f8e26deb54
2 changed files with 24 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/proxy_server.py
View file

@ -6217,7 +6217,6 @@ class ProxyConfig:
)
global_agent_registry.load_agents_from_config(agent_config)
jwt_handler.bind_agent_lookup(global_agent_registry)
mcp_servers_config: Final = config.get("mcp_servers", None)
if mcp_servers_config:
@ -8183,7 +8182,6 @@ class ProxyConfig:
global_agent_registry as AGENT_REGISTRY,
)
jwt_handler.bind_agent_lookup(AGENT_REGISTRY)
try:
async with AGENT_RECONCILE_LOCK:
db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client)
@ -9515,6 +9513,9 @@ class ProxyStartupEvent:
user_api_key_cache=user_api_key_cache,
litellm_jwtauth=litellm_jwtauth,
)
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
jwt_handler.bind_agent_lookup(global_agent_registry)
@classmethod
def _add_proxy_budget_to_db(cls):

34
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -3741,42 +3741,50 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_
@pytest.mark.asyncio
@pytest.mark.parametrize("agents_source", ["config", "db"])
async def test_ProxyConfig_agent_loading_binds_registry_to_jwt_agent_claims(clean_agent_registry, agents_source):
"""A JWT agent claim must resolve against the agents the proxy loaded, whichever source registered them."""
@pytest.mark.parametrize("agents_source", ["config", "db", "api"])
async def test_ProxyStartupEvent_jwt_auth_resolves_agent_claims_against_live_registry(
clean_agent_registry, agents_source
):
"""A JWT agent claim must resolve against every agent the proxy knows, including ones created after startup."""
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.auth.handle_jwt import JWTAuthManager
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.types.agents import AgentResponse
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="appid"),
)
original_lookup = proxy_server.jwt_handler.agent_lookup
proxy_server.jwt_handler.bind_agent_lookup(jwt_handler.agent_lookup)
try:
proxy_server.ProxyStartupEvent._initialize_jwt_auth(
general_settings={"litellm_jwtauth": {"agent_id_jwt_field": "appid"}},
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
)
if agents_source == "config":
await ProxyConfig()._init_non_llm_configs(
config={"agents": [_config_agent("loaded-agent")]},
config_file_path=None,
)
else:
elif agents_source == "db":
prisma_client = MagicMock()
prisma_client.db.litellm_agentstable.find_many = AsyncMock(
return_value=[_FakeAgentRow("db-id", "loaded-agent")]
)
await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
else:
clean_agent_registry.register_agent(
agent_config=AgentResponse(agent_id="api-id", **_config_agent("loaded-agent"))
)
resolved = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_handler=proxy_server.jwt_handler,
jwt_valid_token={"appid": "loaded-agent"},
agent_registry=proxy_server.jwt_handler.agent_lookup,
)
finally:
proxy_server.jwt_handler.bind_agent_lookup(original_lookup)
proxy_server.jwt_handler.update_environment(
prisma_client=None, user_api_key_cache=UserApiKeyCache(), litellm_jwtauth=LiteLLM_JWTAuth()
)
assert resolved == clean_agent_registry.get_agent_by_name(agent_name="loaded-agent").agent_id