From a143191c85d97344c6c487b54ccc2dd7f8f455d2 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 6 May 2026 15:53:11 -0700
Subject: [PATCH] feat(warm-pool): HydratePayload Pydantic schema (LIT-2890)

---
 .../warm_pool/types.py                        | 72 +++++++++++++++++++
 1 file changed, 72 insertions(+)
 create mode 100644 litellm/proxy/agent_session_endpoints/warm_pool/types.py

diff --git a/litellm/proxy/agent_session_endpoints/warm_pool/types.py b/litellm/proxy/agent_session_endpoints/warm_pool/types.py
new file mode 100644
index 00000000000..11ec7ac2690
--- /dev/null
+++ b/litellm/proxy/agent_session_endpoints/warm_pool/types.py
@@ -0,0 +1,72 @@
+"""Wire format for the hydrate protocol (LIT-2890).
+
+The proxy assembles a ``HydratePayload`` per session-attach and pushes it
+into the warm VM via SSM RunCommand (or long-poll, depending on the
+configured transport). The daemon writes:
+
+  ``env_vars``        -> ``/etc/litellm-agent/env``               (mode 0644)
+  ``secrets``         -> ``/etc/litellm-agent/secrets.env``       (mode 0600)
+  ``network_access``  -> ``iptables`` rules applied BEFORE user code runs
+  ``agent_config``    -> daemon's runtime config
+
+Schema is stable across daemon and proxy. Bumping a field requires a
+coordinated daemon release — keep additions backward-compatible (use
+``Optional`` + default ``None`` for new fields).
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Literal, Optional
+
+from pydantic import BaseModel, Field
+
+
+class RepoSpec(BaseModel):
+    """Single git repo the daemon clones at hydrate time."""
+
+    url: str
+    ref: Optional[str] = Field(
+        default=None, description="branch / tag / sha; daemon defaults to main"
+    )
+
+
+class NetworkAccess(BaseModel):
+    """Egress firewall posture applied to the VM before user code runs."""
+
+    mode: Literal["allow_all", "allowlist"] = "allow_all"
+    allowlist: Optional[List[str]] = Field(
+        default=None,
+        description=(
+            "When mode=='allowlist', daemon installs iptables rules permitting "
+            "only these hostnames. Ignored when mode=='allow_all'."
+        ),
+    )
+
+
+class AgentConfig(BaseModel):
+    """Per-session daemon runtime config (model, prompt, PR auto-create flag)."""
+
+    model: str
+    system_prompt: str = ""
+    auto_create_pr: bool = False
+
+
+class HydratePayload(BaseModel):
+    """Full hydrate payload pushed to the warm VM at session-attach.
+
+    Locked schema per LIT-2890. The daemon reads this and transitions from
+    ``LITELLM_AGENT_MODE=warm`` to ``LITELLM_AGENT_MODE=session``.
+    """
+
+    session_id: str
+    agent_id: str
+    jwt: str = Field(description="Session-scoped JWT, replaces the per-VM pool JWT.")
+    jwt_expires_at: str = Field(description="ISO-8601 UTC expiry timestamp.")
+    repos: List[RepoSpec] = Field(default_factory=list)
+    env_vars: Dict[str, str] = Field(default_factory=dict)
+    secrets: Dict[str, str] = Field(
+        default_factory=dict,
+        description="Sensitive values written to disk mode 0600. Never logged.",
+    )
+    network_access: NetworkAccess = Field(default_factory=NetworkAccess)
+    agent_config: AgentConfig