From 51339531212bffe2d9d173155f7a791d27b7fccf Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 6 May 2026 18:04:07 -0700
Subject: [PATCH 1/2] test(managed_agents): unwrap Prisma relation syntax in
 fake DB so create payloads land on FK columns

Real Prisma translates {agent: {connect: {id: X}}} into the underlying
agent_id FK column at write time. The in-memory FakePrismaClient we
share with the proxy tests stores the dict verbatim, so the agent_id
column was None on every row our managed_agents code wrote. Tests
asserting session.agent_id == 'agent_1' (and the same for session_id /
run_id) all failed at the same boundary.

Fix: wrap each table's create() in a small adapter that unwraps the
relation-syntax keys into bare FK columns before delegating. Local to
this test conftest so the change does not leak into the proxy tests
that re-use the same fake.
---
 tests/test_litellm/managed_agents/conftest.py | 57 ++++++++++++++++++-
 1 file changed, 56 insertions(+), 1 deletion(-)

diff --git a/tests/test_litellm/managed_agents/conftest.py b/tests/test_litellm/managed_agents/conftest.py
index 233bfe6db3e..e68f74abfa1 100644
--- a/tests/test_litellm/managed_agents/conftest.py
+++ b/tests/test_litellm/managed_agents/conftest.py
@@ -40,6 +40,59 @@ if str(_ENDPOINTS_TEST_DIR) not in sys.path:
     sys.path.insert(0, str(_ENDPOINTS_TEST_DIR))
 
 
+# Map of relation-name -> FK column name used by our payloads. The real
+# Prisma client translates {"agent": {"connect": {"id": X}}} into the
+# underlying agent_id FK column; the in-memory FakePrismaClient we share
+# with the proxy tests doesn't, so we do the unwrap at write time before
+# handing the data dict to ``_Table.create``.
+_RELATION_FK_MAP = {
+    "agent": "agent_id",
+    "session": "session_id",
+    "run": "run_id",
+}
+
+
+def _unwrap_relations(data):
+    """Translate Prisma relation-syntax keys into bare FK columns.
+
+    Operates on a copy so callers' dicts are not mutated. Only handles the
+    ``{"connect": {"id": <value>}}`` shape, which is what every payload in
+    ``litellm/managed_agents`` uses.
+    """
+    out = dict(data)
+    for relation, fk in _RELATION_FK_MAP.items():
+        rel_value = out.pop(relation, None)
+        if isinstance(rel_value, dict):
+            connect = rel_value.get("connect")
+            if isinstance(connect, dict) and "id" in connect:
+                out[fk] = connect["id"]
+    return out
+
+
+def _patch_table_creates(db) -> None:
+    """Wrap ``_Table.create`` on the shared fake to handle Prisma relations.
+
+    We monkey the wrapper into each per-instance table so the change does
+    not leak into the proxy tests that import the same FakePrismaClient.
+    """
+    table_names = [
+        "litellm_agent",
+        "litellm_agentsession",
+        "litellm_agentrun",
+        "litellm_agentrunevent",
+    ]
+    for name in table_names:
+        table = getattr(db, name, None)
+        if table is None:
+            continue
+        original_create = table.create
+
+        async def create(data, _orig=original_create):
+            return await _orig(data=_unwrap_relations(data))
+
+        table.create = create  # type: ignore[assignment]
+
+
 @pytest.fixture
 def fake_db():
     """Bare in-memory Prisma client; use this directly via ``client.db``.
@@ -51,4 +104,6 @@ def fake_db():
     # Local import: see _ENDPOINTS_TEST_DIR sys.path manipulation above.
     from conftest import FakePrismaClient  # type: ignore  # noqa: E402
 
-    return FakePrismaClient()
+    client = FakePrismaClient()
+    _patch_table_creates(client.db)
+    return client

From 47415550ee65e35e87496b3001af927bc7d05ac5 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 6 May 2026 18:04:35 -0700
Subject: [PATCH 2/2] feat(managed_agents): expose AgentRuntime + ClaudeSDK +
 LiteLLM runtimes from agent_runtime package

The package __init__ now re-exports AgentRuntime/AgentConfig/SessionState
plus both runtime implementations (ClaudeSDKAgentRuntime and the new
LiteLLMAgentRuntime) so callers can do
'from litellm.managed_agents.agent_runtime import LiteLLMAgentRuntime'
without poking into submodules.

Required for the package-level __init__.py re-exports to resolve.
---
 .../managed_agents/agent_runtime/__init__.py    | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 litellm/managed_agents/agent_runtime/__init__.py

diff --git a/litellm/managed_agents/agent_runtime/__init__.py b/litellm/managed_agents/agent_runtime/__init__.py
new file mode 100644
index 00000000000..de794904e42
--- /dev/null
+++ b/litellm/managed_agents/agent_runtime/__init__.py
@@ -0,0 +1,17 @@
+"""Pluggable AgentRuntime implementations."""
+
+from litellm.managed_agents.agent_runtime.base import (
+    AgentConfig,
+    AgentRuntime,
+    SessionState,
+)
+from litellm.managed_agents.agent_runtime.claude_sdk import ClaudeSDKAgentRuntime
+from litellm.managed_agents.agent_runtime.litellm_native import LiteLLMAgentRuntime
+
+__all__ = [
+    "AgentConfig",
+    "AgentRuntime",
+    "SessionState",
+    "ClaudeSDKAgentRuntime",
+    "LiteLLMAgentRuntime",
+]