From 40ade8f35f6995779564bc1a1a370738e0b26586 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 7 May 2026 10:06:55 -0700
Subject: [PATCH] feat(v2 managed agents): add Pydantic types (Agent, Session,
 Message, SandboxSpec)

---
 litellm/managed_agents/types.py | 168 ++++++++++++++++++++++++++++++++
 1 file changed, 168 insertions(+)
 create mode 100644 litellm/managed_agents/types.py

diff --git a/litellm/managed_agents/types.py b/litellm/managed_agents/types.py
new file mode 100644
index 00000000000..4c0c83290bf
--- /dev/null
+++ b/litellm/managed_agents/types.py
@@ -0,0 +1,168 @@
+"""Pydantic v2 types for managed agents v2.
+
+Wire format: snake_case JSON. Timestamps ISO 8601 UTC.
+IDs: prefixed UUIDs — `agt_<uuid4>`, `ses_<uuid4>`, `msg_<uuid4>`.
+
+Request models forbid extra fields; response models allow extras for
+forward-compatibility with proxy-side enrichment.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+# ---------------------------------------------------------------------------
+# Status / sandbox enums
+# ---------------------------------------------------------------------------
+
+SessionStatus = Literal["provisioning", "ready", "terminated", "error"]
+MessageStatus = Literal["in_progress", "completed", "failed"]
+MessageRole = Literal["user", "assistant"]
+SandboxType = Literal["opencode"]
+SandboxSize = Literal["small", "medium", "large"]
+
+
+# ---------------------------------------------------------------------------
+# Agent
+# ---------------------------------------------------------------------------
+
+
+class AgentConfig(BaseModel):
+    """Inner config object on an agent.
+
+    The `litellm_api_key` is masked on read by `masking.mask_litellm_api_key`.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    model: str
+    system_prompt: str
+    tools: List[str]
+    litellm_api_key: str
+    litellm_base_url: str
+
+
+class CreateAgentRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    name: str
+    config: AgentConfig
+
+
+class AgentRow(BaseModel):
+    """Response shape for `POST /v2/agents`. `config.litellm_api_key` is masked."""
+
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    name: str
+    config: AgentConfig
+    created_by: Optional[str] = None
+    created_at: datetime
+    updated_at: datetime
+
+
+# ---------------------------------------------------------------------------
+# Sandbox / session
+# ---------------------------------------------------------------------------
+
+
+class SandboxSpec(BaseModel):
+    """Per-session sandbox configuration.
+
+    `image` is optional and falls back to a pinned built-in.
+    `timeout_minutes` is bounded to [1, 1440].
+    `idle_timeout_minutes` must be in [1, timeout_minutes] (enforced by handler).
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    type: SandboxType
+    size: SandboxSize = "small"
+    timeout_minutes: int = Field(default=60, ge=1, le=1440)
+    idle_timeout_minutes: int = Field(default=10, ge=1)
+    image: Optional[str] = None
+
+
+class Repo(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    url: str
+    starting_ref: str
+    checked_out_sha: Optional[str] = None
+
+
+class CreateSessionRequest(BaseModel):
+    """For `POST /v2/sessions` (Krrish, LIT-2919).
+
+    Wave 1 doesn't implement the endpoint — the type is here for Wave 2/3
+    handlers and adapter integration tests.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    agent_id: str
+    sandbox: SandboxSpec
+    repos: List[Repo] = Field(default_factory=list)
+    env_vars: Dict[str, str] = Field(default_factory=dict)
+
+
+class SessionRow(BaseModel):
+    """Response shape for session reads.
+
+    `sandbox_url` and `sandbox_metadata` are stored on the row internally but
+    NEVER returned in the API response — see contract §6.2.
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    agent_id: str
+    sandbox: SandboxSpec
+    status: SessionStatus
+    repos: List[Repo] = Field(default_factory=list)
+    created_by: Optional[str] = None
+    created_at: datetime
+    terminated_at: Optional[datetime] = None
+
+
+# ---------------------------------------------------------------------------
+# Message
+# ---------------------------------------------------------------------------
+
+
+class CreateMessageRequest(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    content: str
+    model: Optional[str] = None  # falls back to agent.config.model
+
+
+class MessageRow(BaseModel):
+    """Normalized message shape, mapped from the underlying sandbox provider.
+
+    `tools` is omitted when there are no tool calls (assistant-only field).
+    """
+
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    session_id: str
+    role: MessageRole
+    content: str
+    model: Optional[str] = None
+    tools: Optional[List[Dict[str, Any]]] = None
+    status: MessageStatus
+    created_at: datetime
+    completed_at: Optional[datetime] = None
+
+
+class MessageList(BaseModel):
+    """Paginated list response for `GET /v2/sessions/:id/messages`."""
+
+    model_config = ConfigDict(extra="allow")
+
+    data: List[MessageRow]
+    next_cursor: Optional[str] = None
+    has_more: bool = False