Merge upstream main into fix/bedrock-scheduled-batch-status

.circleci/scripts/unit_selection.sh

@@ -77,6 +77,7 @@ legacy_paths() {
       echo tests/unit/embeddings
       echo tests/unit/endpoints
       echo tests/unit/files
+      echo tests/unit/harness
       echo tests/unit/images
       echo tests/unit/interactions
       echo tests/unit/messages
@@ -145,6 +146,7 @@ legacy_paths() {
       echo tests/unit/proxy/test_proxy_token_counter.py
       echo tests/unit/proxy/test_server_root_path.py ;;
     proxy-db-proxy-server-core)
+      echo tests/unit/proxy/test__lazy_features.py
       echo tests/unit/proxy/test_aproxy_startup.py
       echo tests/unit/proxy/test_proxy_server.py ;;
     proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;

.github/ci-coverage-allowlist.yml

@@ -4,6 +4,14 @@ description: >-
   by a job nor listed here, so every entry below is a decision on the record.
 
 test_paths:
+  - reason: >-
+      litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
+      deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
+      LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
+      before changing litellm/harness; the mocked coverage runs in tests/unit/harness and
+      tests/unit/llms/*/harness
+    paths:
+      - tests/harness_e2e
   - reason: >-
       The Rust/Python parity harness is run manually through its local CLI. Recorded replay,
       fixture generation, and harness checks are intentionally outside pull request CI

README.md

@@ -268,6 +268,31 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
 
 </details>
 
+<details>
+<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
+
+### Python SDK - Agents
+
+```python
+import litellm
+from litellm import Harness, sandbox
+
+result = litellm.agent(
+    Harness.CLAUDE_CODE,  # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
+    "Find why tests/test_router.py is flaky and fix it.",
+    sandbox=sandbox.local("./repo"),
+    model="litellm_proxy/claude-sonnet-4-5",  # a model group on your AI Gateway
+)
+
+print(result.text, result.cost, [f.path for f in result.files])
+```
+
+Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
+
+[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
+
+</details>
+
 ### Supported Providers ([Website Supported Models](https://models.litellm.ai/) | [Docs](https://docs.litellm.ai/docs/providers))
 
 | Provider                                                                            | `/chat/completions` | `/messages` | `/responses` | `/embeddings` | `/image/generations` | `/audio/transcriptions` | `/audio/speech` | `/moderations` | `/batches` | `/rerank` |

litellm-proxy-extras/litellm_proxy_extras/migration_lock.py

@@ -87,3 +87,21 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
         f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
         f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
     )
+
+
+@contextmanager
+def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
+    """A session-level, non-blocking hold of the migration coordinator lock on an autocommit
+    connection, for DDL that cannot run inside a transaction (`CREATE INDEX CONCURRENTLY`).
+    Yields whether the lock was acquired; a v2 resolver or another migration job's index build
+    holding it yields False. Released on exit."""
+    from psycopg.rows import class_row
+
+    with connection.cursor(row_factory=class_row(_LockResult)) as cursor:
+        row: Final = cursor.execute("SELECT pg_try_advisory_lock(%s) AS acquired", (MIGRATION_LOCK_KEY,)).fetchone()
+    acquired: Final = row is not None and row.acquired
+    try:
+        yield acquired
+    finally:
+        if acquired:
+            connection.execute("SELECT pg_advisory_unlock(%s)", (MIGRATION_LOCK_KEY,))

litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py

@@ -1,4 +1,5 @@
 import hashlib
+import re
 import subprocess
 from collections.abc import Mapping
 from dataclasses import dataclass
@@ -156,3 +157,48 @@ def baseline_current_schema(
         "review any feature-specific backfill requirements.",
         len(migrations),
     )
+
+
+_LINE_COMMENT_RE: Final = re.compile(r"--[^\n]*")
+_BLOCK_COMMENT_RE: Final = re.compile(r"/\*.*?\*/", re.DOTALL)
+_NO_OP_STATEMENT_RE: Final = re.compile(r"^\s*SELECT\s+1\s*$", re.IGNORECASE)
+
+
+def is_inert_migration(script: str) -> bool:
+    """Whether a migration file changes nothing: only comments and `SELECT 1`, so
+    applying it can neither repeat nor skip a database change."""
+    stripped: Final = _LINE_COMMENT_RE.sub("", _BLOCK_COMMENT_RE.sub("", script))
+    return all(not part.strip() or _NO_OP_STATEMENT_RE.match(part) for part in stripped.split(";"))
+
+
+def roll_back_failed_inert_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool:
+    """Roll back the failed ledger row of a migration whose file in this build is inert,
+    so `migrate deploy` applies the inert file on its next pass. The row records an
+    earlier build's attempt at SQL this build no longer ships (an index now built by the
+    migration job), so no database change can be repeated or skipped by replaying
+    the empty file. The caller commits this checkpoint before the next Prisma command.
+    """
+    from psycopg import sql
+
+    if not is_inert_migration(migration.read_text(encoding="utf-8")):
+        return False
+    coordinator.acquire_prisma_lock()
+    records: Final = _migration_records(coordinator.connection, schema, migration)
+    unfinished: Final = tuple(record for record in records if not record.finished)
+    if len(unfinished) != 1:
+        return False
+    result: Final = coordinator.connection.execute(
+        sql.SQL(
+            "UPDATE {} SET rolled_back_at = current_timestamp "
+            "WHERE id = %s AND finished_at IS NULL AND rolled_back_at IS NULL"
+        ).format(sql.Identifier(schema, "_prisma_migrations")),
+        (unfinished[0].id,),
+    )
+    if result.rowcount != 1:
+        raise RuntimeError("Could not roll back the failed inert migration history row; rerun the database setup.")
+    logger.info(
+        "Rolled back the failed history row of %s: this build ships it as an inert migration, "
+        "its index is built by the migration job",
+        migration.parent.name,
+    )
+    return True

litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql

@@ -1,2 +1,6 @@
--- CreateIndex
-CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");
+-- The (api_key, startTime) index on LiteLLM_SpendLogs is built after migrate deploy,
+-- through litellm_proxy_extras/request_log_indexes.py: concurrently on a plain table and
+-- per partition on a partitioned one. The migration job builds it; a serving proxy that
+-- ran the migrations itself builds it in the background once it serves. A migration
+-- cannot do either without blocking spend-log writes or failing on a partitioned table.
+SELECT 1;

litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql

@@ -1,12 +1,6 @@
--- CreateIndex (CONCURRENTLY)
---
--- Disclaimer:
--- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
---   single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
--- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
---   interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
--- - Do not edit this file after it has been applied to any database: Prisma checksums
---   migrations; add a new migration instead.
--- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
---   without IF NOT EXISTS if you must support older versions).
-CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
+-- The litellm_call_id index on LiteLLM_SpendLogs is built after migrate deploy, through
+-- litellm_proxy_extras/request_log_indexes.py: concurrently on a plain table and per
+-- partition on a partitioned one. The migration job builds it; a serving proxy that ran
+-- the migrations itself builds it in the background once it serves. Postgres refuses
+-- CREATE INDEX CONCURRENTLY on a partitioned parent, so this migration no longer runs it.
+SELECT 1;

litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py

@@ -0,0 +1,448 @@
+"""The request-log indexes built after `prisma migrate deploy` instead of by a migration:
+by the migration job, or by a serving proxy that ran the migrations itself (in the
+background, once it serves).
+
+A migration cannot build them: a plain `CREATE INDEX` blocks spend-log inserts for the
+whole build, and `CREATE INDEX CONCURRENTLY` is refused on a partitioned parent
+(db_scripts/partition_spend_logs.sql). `REQUEST_LOG_INDEXES` is the one list to extend;
+names match what Prisma derives from the `@@index` declarations in schema.prisma, so an
+index a database already has is recognized and never rebuilt.
+"""
+
+import hashlib
+import random
+import re
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final
+
+from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras.migration_lock import held_migration_lock
+
+if TYPE_CHECKING:
+    import psycopg
+    from psycopg import sql
+
+
+@dataclass(frozen=True, slots=True)
+class RequestLogIndex:
+    """One index the migration job owns: the table, the exact Prisma index name and the
+    column list as it would be written after `ON <table>`."""
+
+    table: str
+    name: str
+    definition: str
+
+    @property
+    def columns(self) -> tuple[str, ...]:
+        return tuple(re.findall(r'"([^"]+)"', self.definition))
+
+    def partition_index_name(self, partition: str) -> str:
+        """The child index name for one partition, built the way Postgres names the
+        children of a partitioned index, and kept within the 63 byte identifier limit."""
+        name: Final = f"{partition}_{self.name.removeprefix(f'{self.table}_')}"
+        if len(name.encode()) <= _IDENTIFIER_MAX_BYTES:
+            return name
+        digest: Final = hashlib.sha256(name.encode()).hexdigest()[:_DIGEST_LENGTH]
+        budget: Final = _IDENTIFIER_MAX_BYTES - _DIGEST_LENGTH - 1
+        kept: Final = next(name[:length] for length in range(len(name), 0, -1) if len(name[:length].encode()) <= budget)
+        return f"{kept}_{digest}"
+
+
+REQUEST_LOG_INDEXES: Final = (
+    RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'),
+    RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'),
+)
+
+_IDENTIFIER_MAX_BYTES: Final = 63
+_PARENT_LOCK_TIMEOUT: Final = "2s"
+_PARENT_LOCK_ATTEMPTS: Final = 30
+_LOCK_HANDOVER_SECONDS: Final = 2.0
+_DIGEST_LENGTH: Final = 8
+_CREATE_INDEX_STATEMENT: Final = re.compile(
+    r'^\s*CREATE\s+(?:UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?"(?P<index>[^"]+)"\s+ON\b',
+    re.IGNORECASE,
+)
+_TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)"
+_CHILDREN_WITHOUT_THE_INDEX_SQL: Final = (
+    "SELECT child.relname AS name, n.nspname AS schema, child.relkind = 'p' AS partitioned "
+    "FROM pg_inherits i JOIN pg_class child ON child.oid = i.inhrelid "
+    "JOIN pg_namespace n ON n.oid = child.relnamespace "
+    "WHERE i.inhparent = to_regclass(%s) AND NOT EXISTS ("
+    "SELECT 1 FROM pg_inherits attached JOIN pg_index x ON x.indexrelid = attached.inhrelid "
+    "WHERE attached.inhparent = to_regclass(%s) AND x.indrelid = child.oid) "
+    "ORDER BY child.relname"
+)
+_EQUIVALENT_INDEXES_SQL: Final = (
+    "SELECT i.relname AS name, x.indisvalid AS valid "
+    "FROM pg_index x JOIN pg_class i ON i.oid = x.indexrelid JOIN pg_am am ON am.oid = i.relam "
+    "WHERE x.indrelid = to_regclass(%s) AND i.relname <> %s AND am.amname = 'btree' AND NOT x.indisunique "
+    "AND x.indexprs IS NULL AND x.indpred IS NULL AND x.indnkeyatts = x.indnatts "
+    "AND NOT EXISTS (SELECT 1 FROM unnest(x.indoption::int2[]) o WHERE o <> 0) "
+    "AND NOT EXISTS (SELECT 1 FROM unnest(x.indclass::oid[]) c JOIN pg_opclass oc ON oc.oid = c WHERE NOT oc.opcdefault) "
+    "AND NOT EXISTS (SELECT 1 FROM unnest(x.indcollation::oid[]) WITH ORDINALITY c(coll, ord) "
+    "JOIN unnest(x.indkey::int2[]) WITH ORDINALITY k(attnum, ord) ON k.ord = c.ord "
+    "JOIN pg_attribute a ON a.attrelid = x.indrelid AND a.attnum = k.attnum "
+    "WHERE c.coll <> 0 AND c.coll <> a.attcollation) "
+    "AND (SELECT array_agg(a.attname::text ORDER BY k.ord) FROM unnest(x.indkey::int2[]) WITH ORDINALITY k(attnum, ord) "
+    "JOIN pg_attribute a ON a.attrelid = x.indrelid AND a.attnum = k.attnum) = %s::text[] "
+    "AND NOT EXISTS (SELECT 1 FROM pg_inherits WHERE inhrelid = x.indexrelid) "
+    "ORDER BY x.indisvalid DESC, i.relname"
+)
+_INDEX_STATE_SQL: Final = (
+    'SELECT x.indisvalid AS valid, t.relname AS "table" '
+    "FROM pg_index x JOIN pg_class t ON t.oid = x.indrelid WHERE x.indexrelid = to_regclass(%s)"
+)
+
+
+@dataclass(frozen=True, slots=True)
+class _Relation:
+    name: str
+    schema: str
+    partitioned: bool
+
+
+@dataclass(frozen=True, slots=True)
+class _IndexState:
+    valid: bool
+    table: str
+
+
+@dataclass(frozen=True, slots=True)
+class _EquivalentIndex:
+    name: str
+    valid: bool
+
+
+@dataclass(frozen=True, slots=True)
+class _TableKind:
+    partitioned: bool
+
+
+def filter_request_log_index_diff(diff_sql: str, indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES) -> str:
+    """The `prisma migrate diff` script without the statements that create a migration-job-owned
+    index, which the schema declares and the migrations deliberately do not build."""
+    names: Final = frozenset(index.name for index in indexes)
+    statements: Final = diff_sql.split(";")
+    kept: Final = tuple(statement for statement in statements if not _creates_one_of(statement, names))
+    return ";".join(kept) if any(part.strip() for part in kept) else ""
+
+
+def _creates_one_of(statement: str, names: frozenset[str]) -> bool:
+    match: Final = _CREATE_INDEX_STATEMENT.match(_without_comments(statement))
+    return match is not None and match["index"] in names
+
+
+def _without_comments(statement: str) -> str:
+    return "\n".join(line for line in statement.splitlines() if not line.lstrip().startswith("--"))
+
+
+def _connect(database_url: str) -> "psycopg.Connection[tuple[object, ...]]":
+    import psycopg
+
+    return psycopg.connect(database_url, connect_timeout=10, autocommit=True)
+
+
+def ensure_request_log_indexes(
+    database_url: str,
+    schema: str,
+    indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES,
+    connect: "Callable[[str], psycopg.Connection[tuple[object, ...]]]" = _connect,
+) -> bool:
+    """Build every listed index that is missing or invalid. Each build step runs under
+    the migration coordinator lock, held per statement so a resolver booting on another
+    replica gets in between partitions rather than waiting for the whole table. Any
+    failure is logged and left for the next index build; the result says whether
+    every index ended up valid. Never raises."""
+    import psycopg
+
+    try:
+        with connect(database_url) as connection:
+            connection.execute("SET statement_timeout = 0")
+            results: Final = tuple(_ensure_index(connection, schema, index) for index in indexes)
+    except psycopg.Error as exc:
+        logger.warning("Could not build the request-log indexes, leaving them for the next index build: %s", exc)
+        return False
+    if not all(results):
+        logger.warning("Some request-log indexes are not in place yet, leaving them for the next index build")
+        return False
+    logger.info("Request-log indexes are all in place")
+    return True
+
+
+def _under_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]", step: Callable[[], bool]) -> bool:
+    with held_migration_lock(connection) as held:
+        if not held:
+            logger.info(
+                "Another process holds the migration lock, leaving the request-log indexes to the next index build"
+            )
+            return False
+        return step()
+
+
+def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex) -> bool:
+    from psycopg.rows import class_row
+
+    with connection.cursor(row_factory=class_row(_TableKind)) as cursor:
+        table: Final = cursor.execute(_TABLE_KIND_SQL, (_regclass_name(connection, schema, index.table),)).fetchone()
+    if table is None:
+        logger.info("Table %s does not exist yet, skipping index %s", index.table, index.name)
+        return True
+    if table.partitioned:
+        return build_index_on_partitioned_table(connection, schema, index)
+    return _build_leaf_index(connection, schema, index.table, index.name, index)
+
+
+def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str:
+    from psycopg import sql
+
+    return sql.Identifier(schema, name).as_string(connection)
+
+
+def _create_index_statement(
+    connection: "psycopg.Connection[tuple[object, ...]]", prefix: "sql.Composed", definition: str
+) -> bytes:
+    return (prefix.as_string(connection) + definition).encode()
+
+
+def _index_state(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: str) -> "_IndexState | None":
+    from psycopg.rows import class_row
+
+    with connection.cursor(row_factory=class_row(_IndexState)) as cursor:
+        return cursor.execute(_INDEX_STATE_SQL, (_regclass_name(connection, schema, index),)).fetchone()
+
+
+def _equivalent_indexes(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    table: str,
+    name: str,
+    index: RequestLogIndex,
+) -> tuple[_EquivalentIndex, ...]:
+    """The indexes on `table` other than `name` with the same definition: default btree
+    over the same columns in the same order, no expression, predicate, DESC or custom
+    opclass or collation, and not attached under a partitioned index. Valid ones first."""
+    from psycopg.rows import class_row
+
+    with connection.cursor(row_factory=class_row(_EquivalentIndex)) as cursor:
+        return tuple(
+            cursor.execute(
+                _EQUIVALENT_INDEXES_SQL, (_regclass_name(connection, schema, table), name, list(index.columns))
+            ).fetchall()
+        )
+
+
+def _adopt_equivalent_index(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    table: str,
+    name: str,
+    index: RequestLogIndex,
+) -> bool:
+    """Rename a valid index of the same definition under another name (an operator's
+    hand-built copy, say) to the name this code expects, instead of building a second
+    one. RENAME on an index is a catalog change that lets writes through."""
+    from psycopg import sql
+
+    equivalent: Final = next(
+        (found for found in _equivalent_indexes(connection, schema, table, name, index) if found.valid), None
+    )
+    if equivalent is None:
+        return False
+    logger.info(
+        "Renaming the equivalent index %s on %s to %s instead of building a second one", equivalent.name, table, name
+    )
+    connection.execute(
+        sql.SQL("ALTER INDEX {} RENAME TO {}").format(sql.Identifier(schema, equivalent.name), sql.Identifier(name))
+    )
+    return True
+
+
+def _report_second_copies(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    table: str,
+    name: str,
+    index: RequestLogIndex,
+    concurrently: bool,
+) -> None:
+    """Log every other index of the same definition with the statement that removes it.
+    Dropping is the operator's call: a second copy costs writes and disk, never results."""
+    from psycopg import sql
+
+    drop: Final = "DROP INDEX CONCURRENTLY" if concurrently else "DROP INDEX"
+    for copy in _equivalent_indexes(connection, schema, table, name, index):
+        logger.warning(
+            "Index %s on %s is a second copy of %s and only costs writes and disk; remove it with: %s %s",
+            copy.name,
+            table,
+            name,
+            drop,
+            sql.Identifier(schema, copy.name).as_string(connection),
+        )
+
+
+def _children_without_the_index(
+    connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str, index: str
+) -> tuple[_Relation, ...]:
+    from psycopg.rows import class_row
+
+    with connection.cursor(row_factory=class_row(_Relation)) as cursor:
+        return tuple(
+            cursor.execute(
+                _CHILDREN_WITHOUT_THE_INDEX_SQL,
+                (_regclass_name(connection, schema, table), _regclass_name(connection, schema, index)),
+            ).fetchall()
+        )
+
+
+def _build_leaf_index(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    table: str,
+    name: str,
+    index: RequestLogIndex,
+) -> bool:
+    """Build one plain table's or partition's index with CONCURRENTLY so writes keep
+    flowing. The catalog is read under the migration lock, so a replica that saw an
+    invalid index before the lock finds the valid one another replica just built and
+    leaves it. An invalid index left by an interrupted build is dropped and rebuilt; a
+    valid index of the same definition under another name is renamed rather than
+    duplicated; an index of that name on another table is a collision this code will
+    not touch."""
+    from psycopg import sql
+
+    def build() -> bool:
+        existing: Final = _index_state(connection, schema, name)
+        if existing is not None and existing.table != table:
+            logger.warning(
+                "Index %s already exists on %s rather than %s, leaving it alone", name, existing.table, table
+            )
+            return False
+        if existing is not None and existing.valid:
+            return True
+        if existing is not None:
+            logger.info("Dropping the invalid index %s left by an interrupted build on %s", name, table)
+            connection.execute(sql.SQL("DROP INDEX CONCURRENTLY {}").format(sql.Identifier(schema, name)))
+        elif _adopt_equivalent_index(connection, schema, table, name, index):
+            return True
+        logger.info("Building index %s on %s concurrently", name, table)
+        prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format(
+            sql.Identifier(name), sql.Identifier(schema, table)
+        )
+        connection.execute(_create_index_statement(connection, prefix, index.definition))
+        built: Final = _index_state(connection, schema, name)
+        return built is not None and built.valid
+
+    current: Final = _index_state(connection, schema, name)
+    if current is None or not current.valid or current.table != table:
+        if not _under_migration_lock(connection, build):
+            return False
+        time.sleep(_LOCK_HANDOVER_SECONDS)
+    _report_second_copies(connection, schema, table, name, index, concurrently=True)
+    return True
+
+
+def build_index_on_partitioned_table(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    index: RequestLogIndex,
+    table: "str | None" = None,
+    name: "str | None" = None,
+) -> bool:
+    """Build the index the way Postgres allows on a partitioned parent: a metadata-only
+    parent index ON ONLY the parent, one CONCURRENTLY build per partition, and ATTACH
+    PARTITION for each child. Partitions that are themselves partitioned get the same
+    treatment one level down. Every step checks the catalog before acting, so an
+    interrupted run resumes where it stopped and a second run finds nothing to do; a
+    parent or child index of the same definition under another name is renamed and
+    used rather than duplicated. The connection must be in autocommit mode. True when
+    the parent index ends up valid."""
+
+    parent_table: Final = index.table if table is None else table
+    parent_index: Final = index.name if name is None else name
+    existing: Final = _index_state(connection, schema, parent_index)
+    if existing is not None and existing.table != parent_table:
+        logger.warning(
+            "Index %s already exists on %s rather than %s, leaving it alone", parent_index, existing.table, parent_table
+        )
+        return False
+    if existing is None and not _under_migration_lock(
+        connection,
+        lambda: (
+            _adopt_equivalent_index(connection, schema, parent_table, parent_index, index)
+            or _create_parent_index(connection, schema, parent_index, parent_table, index)
+        ),
+    ):
+        return False
+    children: Final = _children_without_the_index(connection, schema, parent_table, parent_index)
+    if not all(_attach_child_index(connection, schema, parent_index, child, index) for child in children):
+        return False
+    final: Final = _index_state(connection, schema, parent_index)
+    if final is None or not final.valid:
+        return False
+    _report_second_copies(connection, schema, parent_table, parent_index, index, concurrently=False)
+    return True
+
+
+def _create_parent_index(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    name: str,
+    table: str,
+    index: RequestLogIndex,
+) -> bool:
+    """Create the metadata-only parent index. Postgres takes a SHARE lock on the
+    parent for that statement, so it waits for in-flight writes and queues new ones
+    behind it; a short lock_timeout with retries keeps every such pause bounded."""
+    import psycopg
+    from psycopg import sql
+
+    prefix: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format(
+        sql.Identifier(name), sql.Identifier(schema, table)
+    )
+    statement: Final = _create_index_statement(connection, prefix, index.definition)
+    connection.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(_PARENT_LOCK_TIMEOUT)))
+    try:
+        for _ in range(_PARENT_LOCK_ATTEMPTS):
+            try:
+                connection.execute(statement)
+                return True
+            except psycopg.errors.LockNotAvailable:
+                logger.info("Waiting for in-flight writes to %s before creating the parent index %s", table, name)
+                time.sleep(random.uniform(0.1, 0.5))
+    finally:
+        connection.execute("SET lock_timeout = 0")
+    logger.warning("Could not get the parent lock on %s to create %s, leaving it for the next index build", table, name)
+    return False
+
+
+def _attach_child_index(
+    connection: "psycopg.Connection[tuple[object, ...]]",
+    schema: str,
+    parent_index: str,
+    child: _Relation,
+    index: RequestLogIndex,
+) -> bool:
+    from psycopg import sql
+
+    child_index: Final = index.partition_index_name(child.name)
+    built: Final = (
+        build_index_on_partitioned_table(connection, child.schema, index, child.name, child_index)
+        if child.partitioned
+        else _build_leaf_index(connection, child.schema, child.name, child_index, index)
+    )
+    if not built:
+        return False
+
+    def attach() -> bool:
+        connection.execute(
+            sql.SQL("ALTER INDEX {} ATTACH PARTITION {}").format(
+                sql.Identifier(schema, parent_index), sql.Identifier(child.schema, child_index)
+            )
+        )
+        logger.info("Attached index %s on partition %s to %s", child_index, child.name, parent_index)
+        return True
+
+    return _under_migration_lock(connection, attach)

litellm-proxy-extras/litellm_proxy_extras/utils.py

@@ -5,6 +5,7 @@ import re
 import shutil
 import subprocess
 import tempfile
+import threading
 import time
 from collections.abc import Callable
 from dataclasses import dataclass, replace
@@ -13,6 +14,7 @@ from typing import TYPE_CHECKING, Final, Optional
 
 from litellm_proxy_extras import prisma_toolchain
 from litellm_proxy_extras._logging import logger
+from litellm_proxy_extras.migration_lock import held_migration_lock
 from litellm_proxy_extras.prisma_toolchain import (
     PRISMA_COMMAND_TIMEOUT_ENV_VAR,
     PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@@ -24,6 +26,7 @@ from litellm_proxy_extras.replica_identity import (
     REPLICA_IDENTITY_FULL_ENV_VAR,
     apply_replica_identity_full,
 )
+from litellm_proxy_extras.request_log_indexes import ensure_request_log_indexes, filter_request_log_index_diff
 
 if TYPE_CHECKING:
     import psycopg
@@ -433,6 +436,21 @@ class ProxyExtrasDBManager:
                 return True
         return False
 
+    @staticmethod
+    def _filter_migration_job_owned_drift(diff_sql: str, partitioned: bool | None = None) -> str:
+        """The drift script without the indexes the migration job builds (the schema
+        declares them, the migrations deliberately do not) and, when LiteLLM_SpendLogs
+        is partitioned, without its primary-key rewrite and partitioning artifacts."""
+        without_indexes: Final = filter_request_log_index_diff(diff_sql)
+        is_partitioned: Final = ProxyExtrasDBManager.spend_logs_is_partitioned() if partitioned is None else partitioned
+        if not is_partitioned:
+            return without_indexes
+        logger.info(
+            "LiteLLM_SpendLogs is partitioned; removed its primary-key "
+            "rewrite and partitioning artifacts from the drift script"
+        )
+        return filter_partitioned_spend_logs_diff(without_indexes)
+
     @staticmethod
     def _resolve_all_migrations(
         migrations_dir: str, schema_path: str, mark_all_applied: bool = True
@@ -513,21 +531,14 @@ class ProxyExtrasDBManager:
             return
         logger.info(f"Migration diff created at {diff_sql_path}")
 
-        if ProxyExtrasDBManager.spend_logs_is_partitioned():
-            filtered_sql = filter_partitioned_spend_logs_diff(
-                diff_sql_path.read_text()
-            )
-            diff_sql_path.write_text(filtered_sql)
-            logger.info(
-                "LiteLLM_SpendLogs is partitioned; removed its primary-key "
-                "rewrite and partitioning artifacts from the drift script"
-            )
-            if not filtered_sql.strip():
-                logger.info("Drift script is empty after filtering; nothing to apply")
-                if not mark_all_applied:
-                    return
-                ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
+        filtered_sql: Final = ProxyExtrasDBManager._filter_migration_job_owned_drift(diff_sql_path.read_text())
+        diff_sql_path.write_text(filtered_sql)
+        if not filtered_sql.strip():
+            logger.info("Drift script is empty after filtering; nothing to apply")
+            if not mark_all_applied:
                 return
+            ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
+            return
 
         # 2. Run prisma db execute to apply the migration
         applied_ok = False
@@ -800,7 +811,7 @@ class ProxyExtrasDBManager:
             conn.execute(statement)
         except psycopg.Error as e:
             logger.warning(
-                "Could not repair invalid index %s.%s, will retry on the next startup. "
+                "Could not repair invalid index %s.%s, will retry on the next database setup run. "
                 "If this keeps happening, run `%s` by hand as the index owner. Error: %s",
                 index.schema,
                 index.name,
@@ -811,16 +822,21 @@ class ProxyExtrasDBManager:
         logger.info("%s invalid index %s.%s", action, index.schema, index.name)
 
     @staticmethod
-    def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
+    def repair_invalid_indexes(
+        lock_timeout: str = "30s",
+        repair: "Callable[[psycopg.Connection[tuple[str, str, str]], _InvalidIndex], None] | None" = None,
+    ) -> bool:
         """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
         INVALID (a migration deadlock between replicas is the usual cause; the
         retried migration skips them because of IF NOT EXISTS). Never raises:
         returns True when no invalid index remains, False when the repair was
-        skipped or failed and will be retried on the next startup. Looks in the
+        skipped or failed and will be retried on the next database setup run. Looks in the
         schema DATABASE_URL names, the only URL Prisma migrates through, but
         connects over DIRECT_URL when set: the session settings, the advisory
         lock and REINDEX CONCURRENTLY all need one server session, which a
-        transaction pooler does not give."""
+        transaction pooler does not give. Each rebuild holds the migration
+        coordinator lock on its own, like the migration job's index build, so a resolver
+        booting on another replica waits for one index at most."""
         prisma_url: Final = os.getenv("DATABASE_URL")
         if not prisma_url:
             return False
@@ -856,20 +872,53 @@ class ProxyExtrasDBManager:
                 if lock_row is None or not lock_row[0]:
                     logger.info("Another replica is already rebuilding the invalid indexes, skipping")
                     return False
-                for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
-                    ProxyExtrasDBManager._repair_index(conn, index)
+                repair_one: Final = repair or ProxyExtrasDBManager._repair_index
+                repaired: Final = all(
+                    ProxyExtrasDBManager._repair_under_migration_lock(conn, schema, index, repair_one)
+                    for index in found
+                )
+                if not repaired:
+                    return False
                 remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
         except psycopg.Error as e:
-            logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
+            logger.warning(
+                "Could not check for invalid indexes, will retry on the next database setup run. Error: %s", e
+            )
             return False
         return not remaining
 
+    @staticmethod
+    def _repair_under_migration_lock(
+        conn: "psycopg.Connection[tuple[str, str, str]]",
+        schema: str,
+        index: _InvalidIndex,
+        repair: "Callable[[psycopg.Connection[tuple[str, str, str]], _InvalidIndex], None]",
+    ) -> bool:
+        """Rebuild one index under the migration coordinator lock, skipping it when a
+        migration job finished or dropped it in the meantime. False when another process
+        holds the lock, so the check waits for the next database setup run."""
+        with held_migration_lock(conn) as held:
+            if not held:
+                logger.info(
+                    "Another process is building indexes under the migration lock, leaving the "
+                    "invalid index check to the next database setup run"
+                )
+                return False
+            still_invalid: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
+            if any(found.schema == index.schema and found.name == index.name for found in still_invalid):
+                repair(conn, index)
+            return True
+
     @staticmethod
     def _setup_database_v2(use_migrate: bool) -> bool:
         if not use_migrate:
             return ProxyExtrasDBManager._run_database_v2(False)
         from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
-        from litellm_proxy_extras.migration_recovery import baseline_current_schema, recover_completed_migration
+        from litellm_proxy_extras.migration_recovery import (
+            baseline_current_schema,
+            recover_completed_migration,
+            roll_back_failed_inert_migration,
+        )
 
         database_url: Final = os.environ.get("DATABASE_URL")
         if not database_url:
@@ -884,7 +933,9 @@ class ProxyExtrasDBManager:
             if not migration.is_file():
                 return False
             with migration_lock(lock_url) as coordinator:
-                return recover_completed_migration(coordinator, schema, migration)
+                return recover_completed_migration(coordinator, schema, migration) or roll_back_failed_inert_migration(
+                    coordinator, schema, migration
+                )
 
         def baseline_existing(migrations_dir: str) -> None:
             with migration_lock(lock_url) as coordinator:
@@ -1177,13 +1228,16 @@ class ProxyExtrasDBManager:
         )
 
     @staticmethod
-    def setup_database(
-        use_migrate: bool = False, use_v2_resolver: bool = False
-    ) -> bool:
+    def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool:
         """
         Set up the database using either prisma migrate or prisma db push
         Uses migrations from litellm-proxy-extras package
 
+        The request-log indexes in `REQUEST_LOG_INDEXES` are not built here: the
+        migration job builds them through `run_migration_job`, and a serving proxy that
+        ran the migrations itself starts them through `start_request_log_index_build`
+        once it is ready to serve.
+
         Args:
             use_migrate: Whether to use prisma migrate instead of db push
             use_v2_resolver: Opt into the v2 migration resolver (safer during
@@ -1200,10 +1254,48 @@ class ProxyExtrasDBManager:
         migrated = ProxyExtrasDBManager._run_migrations(
             use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
         )
-        if migrated:
-            ProxyExtrasDBManager.repair_invalid_indexes()
-            ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
-        return migrated
+        if not migrated:
+            return False
+        ProxyExtrasDBManager.repair_invalid_indexes()
+        ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
+        return True
+
+    @staticmethod
+    def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
+        """Build the indexes in `REQUEST_LOG_INDEXES` on the writer, in the schema the
+        migrations target. Idempotent and never raises; False when an index is still
+        missing or invalid, so the migration job reports it and gets rerun instead of
+        leaving the table unindexed until the next deploy."""
+        database_url: Final = os.environ.get("DATABASE_URL")
+        if not database_url:
+            return True
+        direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
+            os.environ.get("DIRECT_URL") or database_url
+        )
+        schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
+        return build(direct_url, schema)
+
+    @staticmethod
+    def run_migration_job(
+        use_migrate: bool = False,
+        use_v2_resolver: bool = False,
+        setup: Callable[[bool, bool], bool] = setup_database,
+        build: Callable[[], bool] = build_request_log_indexes,
+    ) -> bool:
+        """The migration job's whole run: `setup_database`, then the request-log indexes,
+        built synchronously so the job exits only once they are in place. False when the
+        migrations failed or an index could not be built, so the Job is rerun."""
+        return setup(use_migrate, use_v2_resolver) and build()
+
+    @staticmethod
+    def start_request_log_index_build(build: Callable[[], bool] = build_request_log_indexes) -> threading.Thread:
+        """A serving proxy that ran the migrations itself (schema updates not disabled)
+        builds the request-log indexes on a daemon thread, so a long build never delays
+        readiness. A build that could not finish is logged and picked up by the next boot
+        or the migration job."""
+        thread: Final = threading.Thread(target=build, name="litellm-request-log-indexes", daemon=True)
+        thread.start()
+        return thread
 
     @staticmethod
     def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
@@ -1247,15 +1339,16 @@ class ProxyExtrasDBManager:
                         logger.info("✅ Post-migration sanity check completed")
                         return True
                     except subprocess.CalledProcessError as e:
-                        logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}")
-                        if "P3009" in e.stderr:
+                        stderr: Final = str(e.stderr or "")
+                        logger.info(f"prisma db error: {stderr}, e: {e.stdout}")
+                        if "P3009" in stderr:
                             # Extract the failed migration name from the error message
                             migration_match = re.search(
-                                r"`(\d+_.*)` migration", e.stderr
+                                r"`(\d+_.*)` migration", stderr
                             )
                             if migration_match:
                                 failed_migration = migration_match.group(1)
-                                if ProxyExtrasDBManager._is_idempotent_error(e.stderr):
+                                if ProxyExtrasDBManager._is_idempotent_error(stderr):
                                     logger.info(
                                         f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied"
                                     )
@@ -1311,8 +1404,8 @@ class ProxyExtrasDBManager:
                                         f"✅ Migration {failed_migration} marked as rolled back... retrying"
                                     )
                         elif (
-                            "P3005" in e.stderr
-                            and "database schema is not empty" in e.stderr
+                            "P3005" in stderr
+                            and "database schema is not empty" in stderr
                         ):
                             logger.info(
                                 "Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system"
@@ -1326,13 +1419,13 @@ class ProxyExtrasDBManager:
                             )
                             logger.info("✅ All migrations resolved.")
                             return True
-                        elif "P3018" in e.stderr:
+                        elif "P3018" in stderr:
                             # Check if this is a permission error or idempotent error
-                            if ProxyExtrasDBManager._is_permission_error(e.stderr):
+                            if ProxyExtrasDBManager._is_permission_error(stderr):
                                 # Permission errors should NOT be marked as applied
                                 # Extract migration name for logging
                                 migration_match = re.search(
-                                    r"Migration name: (\d+_.*)", e.stderr
+                                    r"Migration name: (\d+_.*)", stderr
                                 )
                                 migration_name = (
                                     migration_match.group(1)
@@ -1342,7 +1435,7 @@ class ProxyExtrasDBManager:
 
                                 logger.error(
                                     f"❌ Migration {migration_name} failed due to insufficient permissions. "
-                                    f"Please check database user privileges. Error: {e.stderr}"
+                                    f"Please check database user privileges. Error: {stderr}"
                                 )
 
                                 # Mark as rolled back and exit with error
@@ -1365,7 +1458,7 @@ class ProxyExtrasDBManager:
                                     f"was NOT applied. Please grant necessary database permissions and retry."
                                 ) from e
 
-                            elif ProxyExtrasDBManager._is_idempotent_error(e.stderr):
+                            elif ProxyExtrasDBManager._is_idempotent_error(stderr):
                                 # Idempotent errors mean the migration has effectively been applied
                                 logger.info(
                                     "Migration failed due to idempotent error (e.g., column already exists), "
@@ -1373,7 +1466,7 @@ class ProxyExtrasDBManager:
                                 )
                                 # Extract the migration name from the error message
                                 migration_match = re.search(
-                                    r"Migration name: (\d+_.*)", e.stderr
+                                    r"Migration name: (\d+_.*)", stderr
                                 )
                                 if migration_match:
                                     migration_name = migration_match.group(1)
@@ -1422,7 +1515,7 @@ class ProxyExtrasDBManager:
                                 logger.warning(
                                     f"P3018 error encountered but could not classify "
                                     f"as permission or idempotent error. "
-                                    f"Error: {e.stderr}"
+                                    f"Error: {stderr}"
                                 )
                                 raise
                 else:

litellm/__init__.py

@@ -2282,6 +2282,24 @@ if TYPE_CHECKING:
 # Track if async client cleanup has been registered (for lazy loading)
 _async_client_cleanup_registered = False
 
+# litellm.agent() entrypoints, resolved lazily from litellm.harness by __getattr__.
+_AGENT_EXPORTS: Final = frozenset(
+    {
+        "agent",
+        "aagent",
+        "agent_session",
+        "aagent_session",
+        "agent_resume",
+        "aagent_resume",
+        "agent_capabilities",
+        "Harness",
+        "ClaudeCodeOptions",
+        "CodexOptions",
+        "OpenCodeOptions",
+        "DeepAgentsOptions",
+    }
+)
+
 # Eager loading for backwards compatibility with VCR and other HTTP recording tools
 # When LITELLM_DISABLE_LAZY_LOADING is set, lazy-loaded attributes are loaded at import time
 # For now, this only affects encoding (tiktoken) as it was the only reported issue
@@ -2315,6 +2333,13 @@ def __getattr__(name: str) -> Any:
         handler_func: Final = registry[name]
         return handler_func(name)
 
+    # litellm.agent() and friends: imported on first access (not needed for completion calls)
+    if name == "harness" or name in _AGENT_EXPORTS:
+        import importlib
+
+        harness_module = importlib.import_module("litellm.harness")
+        return harness_module if name == "harness" else getattr(harness_module, name)
+
     # Lazy load encoding from main.py to avoid heavy tiktoken import
     if name == "encoding":
         from ._lazy_imports import get_litellm_globals

litellm/constants.py

@@ -2199,3 +2199,26 @@ EMPTY_MAPPING: Final = MappingProxyType({})
 
 # API endpoint for breached password k-anonymity search
 HIBP_RANGE_API_BASE: Final = "https://api.pwnedpasswords.com/range"
+
+# litellm.harness defaults
+HARNESS_ENDPOINT_HOST: Final = "127.0.0.1"
+HARNESS_ENDPOINT_STARTUP_TIMEOUT_SECONDS: Final = 10.0
+HARNESS_ENDPOINT_REQUEST_TIMEOUT_SECONDS: Final = 600.0
+HARNESS_SESSION_TOKEN_BYTES: Final = 32
+HARNESS_MAX_DIFF_BYTES: Final = 256 * 1024
+HARNESS_STDERR_TAIL_LINES: Final = 40
+HARNESS_STREAM_READ_CHUNK_BYTES: Final = 64 * 1024
+HARNESS_EVENT_QUEUE_MAX_SIZE: Final = 1024
+HARNESS_PROCESS_KILL_GRACE_SECONDS: Final = 5.0
+HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset(
+    {
+        ".git",
+        "node_modules",
+        ".venv",
+        "venv",
+        "__pycache__",
+        ".mypy_cache",
+        ".pytest_cache",
+        ".ruff_cache",
+    }
+)

litellm/harness/__init__.py

@@ -0,0 +1,98 @@
+"""Agent harnesses: run Claude Code, Codex, OpenCode or Deep Agents on any LiteLLM model.
+
+The entrypoints live on the top-level package:
+
+    import litellm
+    from litellm import Harness, sandbox
+
+    result = litellm.agent(
+        Harness.CLAUDE_CODE,
+        "fix the failing test",
+        sandbox=sandbox.local("."),
+        model="litellm_proxy/claude-sonnet-4-5",  # a model group on your AI Gateway
+    )
+
+This module holds the types you get back: events, Result, State, errors.
+"""
+
+from litellm.harness.errors import (
+    CapabilityUnsupported,
+    HarnessError,
+    HarnessInstallFailed,
+    OptionsMismatch,
+    OutputInvalid,
+    SandboxError,
+    SessionClosed,
+    StateIncompatible,
+)
+from litellm.harness.options import (
+    ClaudeCodeOptions,
+    CodexOptions,
+    DeepAgentsOptions,
+    OpenCodeOptions,
+)
+from litellm.harness.runtime import (
+    AsyncEventStream,
+    AsyncSession,
+    aagent,
+    aagent_resume,
+    aagent_session,
+    agent_capabilities,
+)
+from litellm.harness.sync import EventStream, Session, agent, agent_resume, agent_session
+from litellm.harness.types import (
+    Approval,
+    Capabilities,
+    Compaction,
+    Done,
+    Event,
+    FileChange,
+    Harness,
+    Reasoning,
+    Result,
+    State,
+    Text,
+    ToolCall,
+    ToolResult,
+    Usage,
+)
+
+__all__ = (
+    "Approval",
+    "AsyncEventStream",
+    "AsyncSession",
+    "Capabilities",
+    "CapabilityUnsupported",
+    "ClaudeCodeOptions",
+    "CodexOptions",
+    "Compaction",
+    "DeepAgentsOptions",
+    "Done",
+    "Event",
+    "EventStream",
+    "FileChange",
+    "Harness",
+    "HarnessError",
+    "HarnessInstallFailed",
+    "OpenCodeOptions",
+    "OptionsMismatch",
+    "OutputInvalid",
+    "Reasoning",
+    "Result",
+    "SandboxError",
+    "Session",
+    "SessionClosed",
+    "State",
+    "StateIncompatible",
+    "Text",
+    "ToolCall",
+    "ToolResult",
+    "Usage",
+    "aagent",
+    "aagent_resume",
+    "aagent_session",
+    "agent",
+    "agent_capabilities",
+    "agent_resume",
+    "agent_session",
+)

litellm/harness/context.py

@@ -0,0 +1,62 @@
+"""Per-session state shared by the runtime, handlers and harness configs."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+from pydantic import BaseModel
+
+from litellm.harness.options import HarnessOptions
+from litellm.harness.sandbox.base import Sandbox
+from litellm.harness.types import Approval, Harness, PermissionMode
+
+if TYPE_CHECKING:
+    from litellm.harness.endpoint import ModelEndpoint
+
+ApprovalHandler: TypeAlias = Callable[
+    [Approval], bool | Awaitable[bool]  # mutable-ok: Callable parameter list in a type alias, not a runtime collection
+]
+
+
+@dataclass(frozen=True)
+class GatewayTarget:
+    """Resolved LiteLLM AI Gateway for `litellm_proxy/` models. Internal, not exported."""
+
+    api_base: str
+    api_key: str
+
+
+@dataclass
+class SessionContext:
+    """Everything a handler and config need for a session. Owned by the runtime."""
+
+    harness: Harness
+    sandbox: Sandbox
+    session_id: str
+    # Model name as sent to the runtime (litellm_proxy/ prefix already stripped).
+    model: str | None = None
+    gateway: GatewayTarget | None = None
+    api_key: str | None = None
+    api_base: str | None = None
+    endpoint: ModelEndpoint | None = None
+    instructions: str | None = None
+    tools: Sequence[Callable[..., Any]] = ()
+    skills: Sequence[str] = ()
+    disable_tools: Sequence[str] = ()
+    permissions: PermissionMode = "full"
+    on_approval: ApprovalHandler | None = None
+    output: type[BaseModel] | None = None
+    max_turns: int | None = None
+    timeout: float | None = None
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+    options: HarnessOptions | None = None
+    # Set by the handler after each turn.
+    final_text: str = ""
+    output_json: str | None = None
+    # Usage for in-process harnesses that call LiteLLM directly (no model endpoint).
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cost: float = 0.0
+    calls: int = 0

litellm/harness/endpoint.py

@@ -0,0 +1,689 @@
+"""Per-session local model endpoint every CLI harness talks to.
+
+The runtime inside the sandbox points its Anthropic / OpenAI base URL at this endpoint and
+authenticates with a random per-session token. The endpoint either reverse-proxies to a LiteLLM
+AI Gateway (gateway mode) or calls the LiteLLM SDK directly (SDK mode), and counts usage + cost.
+
+starlette and uvicorn are optional: they are imported only when an endpoint starts.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import itertools
+import json
+import logging
+import secrets
+from collections.abc import AsyncIterable, AsyncIterator, Mapping
+from dataclasses import dataclass
+from types import MappingProxyType, ModuleType
+from typing import TYPE_CHECKING, Any, Final
+
+import httpx
+import openai
+
+import litellm
+from litellm.constants import (
+    DEFAULT_POLLING_INTERVAL,
+    HARNESS_ENDPOINT_HOST,
+    HARNESS_ENDPOINT_REQUEST_TIMEOUT_SECONDS,
+    HARNESS_ENDPOINT_STARTUP_TIMEOUT_SECONDS,
+    HARNESS_PROCESS_KILL_GRACE_SECONDS,
+    HARNESS_SESSION_TOKEN_BYTES,
+)
+from litellm.harness.context import GatewayTarget
+from litellm.harness.errors import HarnessError, HarnessInstallFailed
+from litellm.harness.types import Harness, Usage
+from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
+from litellm.types.llms.custom_http import httpxSpecialProvider
+
+if TYPE_CHECKING:
+    from starlette.applications import Starlette
+    from starlette.requests import Request
+    from starlette.responses import Response
+    from uvicorn import Server
+
+verbose_logger: Final = logging.getLogger("LiteLLM")
+
+MISSING_DEPS_MESSAGE = "litellm.harness needs starlette and uvicorn: pip install starlette uvicorn"
+
+ROUTE_MESSAGES = "messages"
+ROUTE_CHAT = "chat/completions"
+ROUTE_RESPONSES = "responses"
+POST_ROUTES = (ROUTE_MESSAGES, ROUTE_CHAT, ROUTE_RESPONSES)
+ROUTE_PREFIXES: Final = ("", "/v1")
+
+# What an SDK call or its stream can raise: LiteLLM maps provider failures onto openai's
+# exception hierarchy; transport errors, bad request kwargs and unserializable chunks remain.
+SDK_ERRORS: Final = (openai.OpenAIError, httpx.HTTPError, HarnessError, ValueError, TypeError)
+
+HOP_BY_HOP_HEADERS = frozenset(
+    {
+        "connection",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "te",
+        "trailer",
+        "trailers",
+        "transfer-encoding",
+        "upgrade",
+        "host",
+        "content-length",
+    }
+)
+DROPPED_REQUEST_HEADERS = HOP_BY_HOP_HEADERS | frozenset(
+    (
+        "authorization",
+        "x-api-key",
+        "accept-encoding",
+    )
+)
+DROPPED_RESPONSE_HEADERS = HOP_BY_HOP_HEADERS | frozenset(("content-encoding",))
+COST_HEADER = "x-litellm-response-cost"
+SSE_MEDIA_TYPE = "text/event-stream"
+
+
+@dataclass(frozen=True)
+class _ServerDeps:
+    uvicorn: ModuleType
+    applications: ModuleType
+    routing: ModuleType
+    responses: ModuleType
+
+
+def _load_server_deps() -> _ServerDeps:
+    """Import starlette + uvicorn on demand; they are not litellm dependencies."""
+    try:
+        import uvicorn
+        from starlette import applications, responses, routing
+    except ImportError as e:
+        raise HarnessInstallFailed(MISSING_DEPS_MESSAGE) from e
+    return _ServerDeps(
+        uvicorn=uvicorn,
+        applications=applications,
+        routing=routing,
+        responses=responses,
+    )
+
+
+@dataclass
+class UsageTracker:
+    """Running token + cost totals for one session."""
+
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cost: float = 0.0
+    calls: int = 0
+
+    def add(self, input_tokens: int = 0, output_tokens: int = 0, cost: float = 0.0) -> None:
+        self.input_tokens += input_tokens
+        self.output_tokens += output_tokens
+        self.cost += cost
+        self.calls += 1
+
+    def snapshot(self) -> Usage:
+        return Usage(
+            input_tokens=self.input_tokens,
+            output_tokens=self.output_tokens,
+            calls=self.calls,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Usage parsing
+# ---------------------------------------------------------------------------
+
+
+def _as_int(value: object) -> int:
+    if isinstance(value, bool):
+        return 0
+    if isinstance(value, (int, float)):
+        return int(value)
+    return 0
+
+
+def usage_from_mapping(usage: object) -> tuple[int, int]:
+    """(input, output) from a usage dict using OpenAI or Anthropic/Responses field names."""
+    if not isinstance(usage, Mapping):
+        return 0, 0
+    input_tokens = usage.get("input_tokens", usage.get("prompt_tokens"))
+    output_tokens = usage.get("output_tokens", usage.get("completion_tokens"))
+    return _as_int(input_tokens), _as_int(output_tokens)
+
+
+def usage_from_body(body: object) -> tuple[int, int]:
+    """Usage from a non-streaming JSON response body."""
+    if not isinstance(body, Mapping):
+        return 0, 0
+    if isinstance(body.get("usage"), Mapping):
+        return usage_from_mapping(body["usage"])
+    response = body.get("response")
+    if isinstance(response, Mapping):
+        return usage_from_mapping(response.get("usage"))
+    return 0, 0
+
+
+class SSEUsageParser:
+    """Collects token usage from an SSE byte stream as it passes through."""
+
+    def __init__(self) -> None:
+        self.input_tokens = 0
+        self.output_tokens = 0
+        self._buffer = b""
+
+    def feed(self, chunk: bytes) -> None:
+        self._buffer += chunk
+        *lines, self._buffer = self._buffer.split(b"\n")
+        for line in lines:
+            self._feed_line(line)
+
+    def close(self) -> None:
+        if self._buffer:
+            self._feed_line(self._buffer)
+            self._buffer = b""
+
+    def _feed_line(self, line: bytes) -> None:
+        text = line.strip()
+        if not text.startswith(b"data:"):
+            return
+        payload = text[len(b"data:") :].strip()
+        if not payload or payload == b"[DONE]":
+            return
+        try:
+            event = json.loads(payload)
+        except ValueError:
+            return
+        if isinstance(event, Mapping):
+            self.absorb(event)
+
+    def absorb(self, event: Mapping[str, Any]) -> None:
+        event_type = event.get("type")
+        if event_type == "message_start":
+            self._absorb_message_start(event)
+        elif event_type == "message_delta":
+            self._absorb_message_delta(event)
+        elif event_type == "response.completed":
+            self._absorb_response_completed(event)
+        elif isinstance(event.get("usage"), Mapping):
+            self._set(*usage_from_mapping(event["usage"]))
+
+    def _absorb_message_start(self, event: Mapping[str, Any]) -> None:
+        message = event.get("message")
+        if isinstance(message, Mapping):
+            self._set(*usage_from_mapping(message.get("usage")))
+
+    def _absorb_message_delta(self, event: Mapping[str, Any]) -> None:
+        # message_delta output_tokens is cumulative for the whole message.
+        self._set(*usage_from_mapping(event.get("usage")))
+
+    def _absorb_response_completed(self, event: Mapping[str, Any]) -> None:
+        response = event.get("response")
+        if isinstance(response, Mapping):
+            self._set(*usage_from_mapping(response.get("usage")))
+
+    def _set(self, input_tokens: int, output_tokens: int) -> None:
+        if input_tokens:
+            self.input_tokens = input_tokens
+        if output_tokens:
+            self.output_tokens = output_tokens
+
+
+# ---------------------------------------------------------------------------
+# Cost + helpers
+# ---------------------------------------------------------------------------
+
+
+def compute_cost(model: str | None, input_tokens: int, output_tokens: int) -> float:
+    """Cost from LiteLLM's price map. Never raises; unknown models cost 0.0."""
+    if not model or not (input_tokens or output_tokens):
+        return 0.0
+    try:
+        prompt_cost, completion_cost = litellm.cost_per_token(
+            model=model, prompt_tokens=input_tokens, completion_tokens=output_tokens
+        )
+        return float(prompt_cost) + float(completion_cost)
+    except Exception:  # accounting must never break a call; the price-map lookup raises bare Exception
+        verbose_logger.debug("harness endpoint: cost lookup failed for %s", model, exc_info=True)
+        return 0.0
+
+
+def header_cost(headers: Mapping[str, str]) -> float | None:
+    raw = headers.get(COST_HEADER)
+    if raw is None:
+        return None
+    try:
+        return float(raw)
+    except (TypeError, ValueError):
+        return None
+
+
+def hidden_cost(response: object) -> float | None:
+    hidden = getattr(response, "_hidden_params", None)
+    if not isinstance(hidden, Mapping):
+        return None
+    try:
+        cost = hidden.get("response_cost")
+        return None if cost is None else float(cost)
+    except (TypeError, ValueError):
+        return None
+
+
+def extract_token(headers: Mapping[str, str]) -> str | None:
+    auth = headers.get("authorization") or ""
+    if auth.lower().startswith("bearer "):
+        return auth[len("bearer ") :].strip()
+    return headers.get("x-api-key")
+
+
+def gateway_headers(
+    incoming: Mapping[str, str],
+    gateway: GatewayTarget,
+    harness: Harness,
+    metadata: Mapping[str, Any] | None,
+) -> Mapping[str, str]:
+    """Incoming headers minus hop-by-hop/auth/x-litellm-*, plus gateway auth, tags, metadata."""
+    kept = (
+        (name, value)
+        for name, value in incoming.items()
+        if name.lower() not in DROPPED_REQUEST_HEADERS and not name.lower().startswith("x-litellm-")
+    )
+    metadata_json = json.dumps(dict(metadata), default=str) if metadata else None  # mutable-ok: for json.dumps
+    metadata_header = (("x-litellm-spend-logs-metadata", metadata_json),) if metadata_json is not None else ()
+    added = (
+        ("authorization", f"Bearer {gateway.api_key}"),
+        ("x-litellm-tags", f"harness,{harness.value}"),
+        *metadata_header,
+    )
+    return MappingProxyType(dict(itertools.chain(kept, added)))
+
+
+def response_headers(upstream: Mapping[str, str]) -> Mapping[str, str]:
+    return MappingProxyType(
+        {name: value for name, value in upstream.items() if name.lower() not in DROPPED_RESPONSE_HEADERS}
+    )
+
+
+def sanitize(message: str, secret_values: tuple[str | None, ...]) -> str:
+    for value in secret_values:
+        if value:
+            message = message.replace(value, "***")
+    return message
+
+
+def error_status(exc: BaseException) -> int:
+    status = getattr(exc, "status_code", None)
+    if isinstance(status, int) and 400 <= status <= 599:
+        return status
+    return 500
+
+
+def error_body(exc: BaseException, message: str) -> dict[str, Any]:  # mutable-ok: JSONResponse body
+    return {"error": {"type": type(exc).__name__, "message": message}}  # mutable-ok: JSONResponse body
+
+
+def to_jsonable(obj: object) -> object:
+    if hasattr(obj, "model_dump"):
+        return obj.model_dump(mode="json", exclude_none=True)
+    if isinstance(obj, Mapping):
+        return dict(obj)  # mutable-ok: plain-dict copy so json.dumps can serialize any Mapping
+    return obj
+
+
+def encode_anthropic_chunk(chunk: object) -> bytes:
+    if isinstance(chunk, bytes):
+        return chunk
+    if isinstance(chunk, str):
+        return chunk.encode()
+    data = to_jsonable(chunk)
+    event_type = data.get("type", "message") if isinstance(data, Mapping) else "message"
+    return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
+
+
+def encode_chat_chunk(chunk: object) -> bytes:
+    if hasattr(chunk, "model_dump_json"):
+        return f"data: {chunk.model_dump_json()}\n\n".encode()
+    return f"data: {json.dumps(to_jsonable(chunk))}\n\n".encode()
+
+
+def encode_responses_chunk(chunk: object) -> bytes:
+    data = to_jsonable(chunk)
+    event_type = data.get("type", "message") if isinstance(data, Mapping) else "message"
+    return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
+
+
+STREAM_ENCODERS = MappingProxyType(
+    {
+        ROUTE_MESSAGES: encode_anthropic_chunk,
+        ROUTE_CHAT: encode_chat_chunk,
+        ROUTE_RESPONSES: encode_responses_chunk,
+    }
+)
+STREAM_TRAILERS = MappingProxyType({ROUTE_CHAT: b"data: [DONE]\n\n"})
+
+
+def route_of(path: str) -> str:
+    stripped = path.strip("/")
+    stripped = stripped.removeprefix("v1/")
+    return stripped
+
+
+def _noop() -> None:
+    return None
+
+
+# ---------------------------------------------------------------------------
+# ModelEndpoint
+# ---------------------------------------------------------------------------
+
+
+class ModelEndpoint:
+    """Local HTTP endpoint for one harness session. Use as an async context manager."""
+
+    def __init__(
+        self,
+        harness: Harness,
+        model: str | None,
+        gateway: GatewayTarget | None,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        metadata: Mapping[str, Any] | None = None,
+        *,
+        client: httpx.AsyncClient | None = None,
+    ) -> None:
+        self.harness = harness
+        self.model = model
+        self.gateway = gateway
+        self.api_key = api_key
+        self.api_base = api_base
+        self.metadata: Mapping[str, Any] = MappingProxyType(dict(metadata or ()))
+        self.token = secrets.token_urlsafe(HARNESS_SESSION_TOKEN_BYTES)
+        self.usage = UsageTracker()
+        self.port = 0
+        # Injected client (tests); production uses LiteLLM's shared cached client.
+        self._injected_client = client
+        self._deps: _ServerDeps | None = None
+        self._client: httpx.AsyncClient | None = None
+        self._server: Any = None
+        self._task: asyncio.Task[None] | None = None
+
+    @property
+    def url(self) -> str:
+        return f"http://{HARNESS_ENDPOINT_HOST}:{self.port}"
+
+    # -- lifecycle ----------------------------------------------------------
+
+    async def __aenter__(self) -> ModelEndpoint:
+        await self.start()
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.stop()
+
+    async def start(self) -> None:
+        self._deps = _load_server_deps()
+        if self.gateway is not None:
+            self._client = self._gateway_client()
+        self._server = self._build_server(self._deps)
+        self._task = asyncio.create_task(self._server.serve())
+        try:
+            await asyncio.wait_for(self._wait_started(), HARNESS_ENDPOINT_STARTUP_TIMEOUT_SECONDS)
+        except BaseException:
+            await self.stop()
+            raise
+        self.port = self._server.servers[0].sockets[0].getsockname()[1]
+
+    async def stop(self) -> None:
+        if self._server is not None:
+            self._server.should_exit = True
+        if self._task is not None:
+            with contextlib.suppress(BaseException):
+                await self._task
+            self._task = None
+        # Never close the client: the shared cached one may still serve other requests,
+        # and an injected one belongs to its caller.
+        self._client = None
+
+    def _gateway_client(self) -> httpx.AsyncClient:
+        """LiteLLM's shared cached async client, unless one was injected."""
+        if self._injected_client is not None:
+            return self._injected_client
+        handler = get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.AgentHarness,
+            params={  # mutable-ok: get_async_httpx_client takes a dict params argument
+                "timeout": HARNESS_ENDPOINT_REQUEST_TIMEOUT_SECONDS
+            },
+        )
+        return handler.client
+
+    async def _wait_started(self) -> None:
+        while not self._server.started:
+            if self._task is not None and self._task.done():
+                raise HarnessError("harness model endpoint failed to start")
+            await asyncio.sleep(DEFAULT_POLLING_INTERVAL)
+
+    def _build_server(self, deps: _ServerDeps) -> Server:
+        config = deps.uvicorn.Config(
+            self._build_app(deps),
+            host=HARNESS_ENDPOINT_HOST,
+            port=0,
+            log_config=None,
+            log_level="warning",
+            access_log=False,
+            lifespan="off",
+            timeout_graceful_shutdown=HARNESS_PROCESS_KILL_GRACE_SECONDS,
+        )
+        server = deps.uvicorn.Server(config)
+        # Never touch the host process's signal handlers.
+        if hasattr(server, "capture_signals"):
+            server.capture_signals = contextlib.nullcontext
+        if hasattr(server, "install_signal_handlers"):
+            server.install_signal_handlers = _noop
+        return server
+
+    def _build_app(self, deps: _ServerDeps) -> Starlette:
+        Route = deps.routing.Route
+        post_routes = tuple(
+            Route(
+                f"{prefix}/{route}",
+                self._handle,
+                methods=["POST"],  # mutable-ok: Starlette Route takes a methods list
+            )
+            for prefix, route in itertools.product(ROUTE_PREFIXES, POST_ROUTES)
+        )
+        get_routes = tuple(
+            Route(f"{prefix}/models", self._models, methods=["GET"])  # mutable-ok: Starlette Route takes a methods list
+            for prefix in ROUTE_PREFIXES
+        )
+        return deps.applications.Starlette(
+            routes=[*post_routes, *get_routes]  # mutable-ok: Starlette takes a routes list
+        )
+
+    # -- request handling ---------------------------------------------------
+
+    @property
+    def _responses(self) -> ModuleType:
+        if self._deps is None:
+            raise HarnessError("harness model endpoint is not started")
+        return self._deps.responses
+
+    def _authorized(self, request: Request) -> bool:
+        token = extract_token(request.headers)
+        return token is not None and secrets.compare_digest(token.encode(), self.token.encode())
+
+    def _json(self, body: object, status_code: int = 200) -> Response:
+        return self._responses.JSONResponse(body, status_code=status_code)
+
+    def _unauthorized(self) -> Response:
+        return self._json(
+            {"error": {"type": "authentication_error", "message": "invalid token"}},  # mutable-ok: JSONResponse body
+            401,
+        )
+
+    def _error(self, exc: BaseException, status_code: int | None = None) -> Response:
+        message = sanitize(str(exc), self._secrets())
+        return self._json(error_body(exc, message), status_code or error_status(exc))
+
+    def _secrets(self) -> tuple[str | None, ...]:
+        gateway_key = self.gateway.api_key if self.gateway else None
+        return (gateway_key, self.api_key, self.token)
+
+    async def _models(self, request: Request) -> Response:
+        if not self._authorized(request):
+            return self._unauthorized()
+        entry = {"id": self.model, "object": "model", "created": 0, "owned_by": "litellm"}  # mutable-ok: JSON body
+        data = (entry,) if self.model else ()
+        return self._json({"object": "list", "data": data})  # mutable-ok: JSON response body for Starlette JSONResponse
+
+    async def _handle(self, request: Request) -> Response:
+        if not self._authorized(request):
+            return self._unauthorized()
+        try:
+            body = json.loads(await request.body())
+        except ValueError as e:
+            return self._error(e, 400)
+        if not isinstance(body, dict):
+            return self._error(ValueError("request body must be a JSON object"), 400)
+        route = route_of(request.url.path)
+        if self.gateway is not None:
+            return await self._forward(request, route, body)
+        return await self._call_sdk(route, body)
+
+    def _cost_model(self, body: Mapping[str, Any]) -> str | None:
+        model = self.model or body.get("model")
+        return model if isinstance(model, str) else None
+
+    def _record(
+        self,
+        model: str | None,
+        input_tokens: int,
+        output_tokens: int,
+        cost: float | None,
+    ) -> None:
+        if cost is None:
+            cost = compute_cost(model, input_tokens, output_tokens)
+        self.usage.add(input_tokens, output_tokens, cost)
+
+    # -- gateway mode -------------------------------------------------------
+
+    async def _forward(self, request: Request, route: str, body: Mapping[str, Any]) -> Response:
+        if self._client is None or self.gateway is None:
+            raise HarnessError("gateway client is not started")
+        if self.model:
+            body = {**body, "model": self.model}  # mutable-ok: JSON request body re-sent upstream via httpx json=
+        upstream_request = self._client.build_request(
+            "POST",
+            f"{self.gateway.api_base}/v1/{route}",
+            json=body,
+            headers=gateway_headers(request.headers, self.gateway, self.harness, self.metadata),
+        )
+        try:
+            upstream = await self._client.send(upstream_request, stream=True)
+        except httpx.HTTPError as e:
+            return self._error(e, 502)
+        return self._responses.StreamingResponse(
+            self._relay(upstream, self._cost_model(body)),
+            status_code=upstream.status_code,
+            headers=response_headers(upstream.headers),
+        )
+
+    async def _relay(self, upstream: httpx.Response, model: str | None) -> AsyncIterator[bytes]:
+        is_sse = SSE_MEDIA_TYPE in upstream.headers.get("content-type", "")
+        parser = SSEUsageParser()
+        collected = bytearray()
+        try:
+            async for chunk in upstream.aiter_bytes():
+                if is_sse:
+                    parser.feed(chunk)
+                else:
+                    collected.extend(chunk)
+                yield chunk
+        finally:
+            await upstream.aclose()
+            if upstream.status_code < 400:
+                self._record_relayed(upstream, model, parser, is_sse, bytes(collected))
+
+    def _record_relayed(
+        self,
+        upstream: httpx.Response,
+        model: str | None,
+        parser: SSEUsageParser,
+        is_sse: bool,
+        collected: bytes,
+    ) -> None:
+        if is_sse:
+            parser.close()
+            tokens = (parser.input_tokens, parser.output_tokens)
+        else:
+            try:
+                tokens = usage_from_body(json.loads(collected))
+            except ValueError:
+                tokens = (0, 0)
+        self._record(model, tokens[0], tokens[1], header_cost(upstream.headers))
+
+    # -- SDK mode -----------------------------------------------------------
+
+    def _sdk_kwargs(
+        self, body: Mapping[str, Any]
+    ) -> dict[str, Any]:  # mutable-ok: SDK call kwargs, mutated by _invoke_sdk then splatted
+        kwargs: dict[str, Any] = {**body}  # mutable-ok: SDK call kwargs built from the JSON body, then overridden
+        if self.model:
+            kwargs["model"] = self.model
+        if self.api_key:
+            kwargs["api_key"] = self.api_key
+        if self.api_base:
+            kwargs["api_base"] = self.api_base
+        return kwargs
+
+    async def _invoke_sdk(
+        self,
+        route: str,
+        kwargs: dict[str, Any],  # mutable-ok: injects stream_options into the SDK kwargs
+    ) -> object:
+        if route == ROUTE_MESSAGES:
+            return await litellm.anthropic.messages.acreate(**kwargs)
+        if route == ROUTE_CHAT:
+            if kwargs.get("stream"):
+                stream_options = kwargs.get("stream_options") or {}  # mutable-ok: empty default for a JSON field
+                kwargs["stream_options"] = {  # mutable-ok: JSON field sent to litellm.acompletion
+                    "include_usage": True,
+                    **stream_options,
+                }
+            return await litellm.acompletion(**kwargs)
+        return await litellm.aresponses(**kwargs)
+
+    async def _call_sdk(self, route: str, body: Mapping[str, Any]) -> Response:
+        kwargs = self._sdk_kwargs(body)
+        model = self._cost_model(kwargs)
+        try:
+            response = await self._invoke_sdk(route, kwargs)
+        except SDK_ERRORS as e:
+            verbose_logger.debug("harness endpoint: SDK call failed: %s", type(e).__name__)
+            return self._error(e)
+        if kwargs.get("stream") and isinstance(response, AsyncIterable):
+            return self._responses.StreamingResponse(
+                self._sdk_stream(route, response, model), media_type=SSE_MEDIA_TYPE
+            )
+        data = to_jsonable(response)
+        input_tokens, output_tokens = usage_from_body(data)
+        self._record(model, input_tokens, output_tokens, hidden_cost(response))
+        return self._json(data)
+
+    async def _sdk_stream(self, route: str, iterator: AsyncIterable[object], model: str | None) -> AsyncIterator[bytes]:
+        encode = STREAM_ENCODERS[route]
+        parser = SSEUsageParser()
+        try:
+            async for chunk in iterator:
+                encoded = encode(chunk)
+                parser.feed(encoded)
+                yield encoded
+            trailer = STREAM_TRAILERS.get(route)
+            if trailer:
+                yield trailer
+        except SDK_ERRORS as e:
+            message = sanitize(str(e), self._secrets())
+            yield f"event: error\ndata: {json.dumps(error_body(e, message))}\n\n".encode()
+        finally:
+            parser.close()
+            self._record(model, parser.input_tokens, parser.output_tokens, None)

litellm/harness/errors.py

@@ -0,0 +1,45 @@
+"""Exceptions raised by litellm.harness."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from litellm.harness.types import Result
+
+
+class HarnessError(Exception):
+    """Base class for every litellm.harness error."""
+
+
+class CapabilityUnsupported(HarnessError):
+    """The harness cannot do what was asked. Raised before the runtime starts."""
+
+
+class OptionsMismatch(HarnessError):
+    """Options for a different harness, or a native option LiteLLM manages itself."""
+
+
+class HarnessInstallFailed(HarnessError):
+    """The runtime is missing from the sandbox or failed to start."""
+
+
+class SandboxError(HarnessError):
+    """The sandbox failed to start, run a command, or reach the host."""
+
+
+class SessionClosed(HarnessError):
+    """A turn was started on a session that is closed or detached."""
+
+
+class StateIncompatible(HarnessError):
+    """resume() was given a State from another harness or an unreadable version."""
+
+
+class OutputInvalid(HarnessError):
+    """The final answer did not validate against output=."""
+
+    def __init__(self, message: str, raw: str, result: Result | None = None) -> None:
+        super().__init__(message)
+        self.raw = raw
+        self.result = result

litellm/harness/handlers/__init__.py

@@ -0,0 +1,35 @@
+"""Handlers run a harness config: CLI runtimes as subprocesses, Deep Agents in-process."""
+
+from __future__ import annotations
+
+from litellm.harness.errors import HarnessError
+from litellm.harness.handlers.base import BaseHarnessHandler
+from litellm.harness.types import Harness, require_harness
+from litellm.llms.base_llm.harness.transformation import (
+    BaseCLIHarnessConfig,
+    BaseHarnessConfig,
+)
+from litellm.utils import ProviderConfigManager
+
+
+def get_harness_config(harness: Harness) -> BaseHarnessConfig:
+    config = ProviderConfigManager.get_provider_harness_config(require_harness(harness))
+    if config is None:
+        raise HarnessError(f"No harness config registered for Harness.{harness.name}")
+    return config
+
+
+def get_harness_handler(config: BaseHarnessConfig) -> BaseHarnessHandler:
+    """The handler that knows how to run this kind of config."""
+    if isinstance(config, BaseCLIHarnessConfig):
+        from litellm.harness.handlers.cli_handler import CLIHarnessHandler
+
+        return CLIHarnessHandler(config)
+    if config.harness is Harness.DEEPAGENTS:
+        from litellm.harness.handlers.deepagents_handler import DeepAgentsHandler
+
+        return DeepAgentsHandler(config)
+    raise HarnessError(f"No handler for Harness.{config.harness.name}")
+
+
+__all__ = ("BaseHarnessHandler", "get_harness_config", "get_harness_handler")

litellm/harness/handlers/base.py

@@ -0,0 +1,42 @@
+"""The handler interface the runtime drives. A handler owns I/O for one session."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Any
+
+from litellm.harness.context import SessionContext
+from litellm.harness.errors import CapabilityUnsupported
+from litellm.harness.types import Event
+from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
+
+
+class BaseHarnessHandler(ABC):
+    """Runs one harness session. The config decides what to run; the handler runs it."""
+
+    def __init__(self, config: BaseHarnessConfig) -> None:
+        self.config = config
+
+    @abstractmethod
+    async def start(self, ctx: SessionContext) -> None:
+        """Prepare the runtime (config files, skills, agent build). Called again after an interrupt."""
+
+    @abstractmethod
+    def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
+        """Run one turn and yield events (never Done). Sets ctx.final_text / ctx.output_json."""
+
+    @abstractmethod
+    async def stop(self, ctx: SessionContext) -> None:
+        """Stop anything this handler started. Safe to call twice."""
+
+    @abstractmethod
+    def native_session_id(self) -> str | None:
+        """The runtime's own session id, for State / resume."""
+
+    @abstractmethod
+    async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
+        """Continue the runtime's own session on the next turn."""
+
+    async def history(self, ctx: SessionContext) -> list[dict[str, Any]]:  # mutable-ok: public history() API shape
+        raise CapabilityUnsupported(f"Harness.{self.config.harness.name} does not expose history")

litellm/harness/handlers/cli_handler.py

@@ -0,0 +1,161 @@
+"""
+Generic handler for CLI harnesses (Claude Code, Codex, OpenCode).
+
+The config (`litellm/llms/<harness>/harness/transformation.py`) says what to run and how to
+read it; this handler does every sandbox and process operation: binary check, private dir,
+config files, persisted dirs, skills, spawning the turn, streaming stdout lines into the
+config's parser, collecting stderr, and killing the process on early exit.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections import deque
+from collections.abc import AsyncIterator, Sequence
+from typing import Any, Final
+
+from litellm._logging import verbose_logger
+from litellm.constants import HARNESS_STDERR_TAIL_LINES, HARNESS_STREAM_READ_CHUNK_BYTES
+from litellm.harness.context import SessionContext
+from litellm.harness.errors import HarnessInstallFailed, SandboxError
+from litellm.harness.handlers.base import BaseHarnessHandler
+from litellm.harness.sandbox.base import Process, Sandbox
+from litellm.harness.types import Event
+from litellm.llms.base_llm.harness.transformation import BaseCLIHarnessConfig, HarnessSessionSetup
+from litellm.llms.base_llm.harness.utils import decode_json_line, read_skill_files
+
+# Link <private_dir>/<dir> to a LiteLLM-owned cache dir so a later session can resume.
+PERSIST_DIR_SCRIPT: Final = (
+    'd="${HOME:-/tmp}/.cache/litellm-harness/$2"; mkdir -p "$d" && mkdir -p "$(dirname "$1")" && ln -sfn "$d" "$1"'
+)
+
+
+async def iter_stream_lines(stream: asyncio.StreamReader) -> AsyncIterator[bytes]:
+    """Newline-delimited lines without StreamReader's 64KiB readline limit."""
+    buffer = b""
+    while True:
+        chunk = await stream.read(HARNESS_STREAM_READ_CHUNK_BYTES)
+        if not chunk:
+            break
+        buffer += chunk
+        *lines, buffer = buffer.split(b"\n")
+        for line in lines:
+            yield line
+    if buffer:
+        yield buffer
+
+
+async def drain_stderr(stream: asyncio.StreamReader, tail: deque[str]) -> None:  # mutable-ok: stderr ring
+    async for line in iter_stream_lines(stream):
+        tail.append(line.decode("utf-8", errors="replace"))
+
+
+async def send_stdin(proc: Process, data: str) -> None:
+    if proc.stdin is None:
+        raise SandboxError("harness process has no stdin")
+    proc.stdin.write(data.encode("utf-8"))
+    await proc.stdin.drain()
+    proc.stdin.close()
+
+
+async def private_dir_for(sandbox: Sandbox) -> str:
+    tempdir = getattr(sandbox, "tempdir", None)
+    if tempdir is None:
+        raise SandboxError(f"{type(sandbox).__name__} has no tempdir(); CLI harnesses need a private config dir")
+    path: str = await tempdir()
+    return path
+
+
+def sandbox_path(private_dir: str, path: str) -> str:
+    return path if path.startswith("/") else f"{private_dir}/{path}"
+
+
+async def persist_dir(sandbox: Sandbox, link_path: str, cache_subpath: str) -> None:
+    script_args: Final = ("-c", PERSIST_DIR_SCRIPT, "sh", link_path, cache_subpath)
+    cmd: Final = ["sh", *script_args]  # mutable-ok: Sandbox.run takes list[str]
+    run = await sandbox.run(cmd)
+    if run.exit_code != 0:
+        verbose_logger.debug(
+            "harness: could not persist %s, resume across sessions disabled: %s", cache_subpath, run.stderr.strip()
+        )
+
+
+async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str) -> None:
+    for skill in skills:
+        name = os.path.basename(os.path.realpath(os.fspath(skill)))
+        for rel, data in await asyncio.to_thread(read_skill_files, skill):
+            await sandbox.write(f"{skills_root}/{name}/{rel.replace(os.sep, '/')}", data)
+
+
+class CLIHarnessHandler(BaseHarnessHandler):
+    config: BaseCLIHarnessConfig
+
+    def __init__(self, config: BaseCLIHarnessConfig) -> None:
+        super().__init__(config)
+        self._private_dir: str | None = None
+        self._setup: HarnessSessionSetup | None = None
+        self._native_id: str | None = None
+        self._proc: Process | None = None
+
+    async def start(self, ctx: SessionContext) -> None:
+        self.config.validate_environment(ctx)
+        binary = self.config.get_binary()
+        if not await ctx.sandbox.which(binary):
+            raise HarnessInstallFailed(
+                f"`{binary}` was not found on PATH in the sandbox. Install it with: {self.config.get_install_hint()}"
+            )
+        private_dir = await private_dir_for(ctx.sandbox)
+        setup = self.config.transform_session_setup(ctx, private_dir)
+        for link, cache_subpath in setup.persisted_dirs:
+            await persist_dir(ctx.sandbox, sandbox_path(private_dir, link), cache_subpath)
+        for rel_path, data in setup.files.items():
+            await ctx.sandbox.write(sandbox_path(private_dir, rel_path), data)
+        if ctx.skills and setup.skills_dir:
+            await copy_skills(ctx.sandbox, tuple(ctx.skills), sandbox_path(private_dir, setup.skills_dir))
+        self._private_dir = private_dir
+        self._setup = setup
+
+    async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
+        if self._setup is None or self._private_dir is None:
+            raise RuntimeError("CLIHarnessHandler.turn() called before start()")
+        request = self.config.transform_turn_request(ctx, self._setup, self._private_dir, prompt, self._native_id)
+        argv: Final = list(request.argv)  # mutable-ok: Sandbox.exec takes list[str]
+        proc = await ctx.sandbox.exec(argv, env=request.env, cwd=request.cwd)
+        self._proc = proc
+        tail: Final[deque[str]] = deque(maxlen=HARNESS_STDERR_TAIL_LINES)  # mutable-ok: bounded stderr ring buffer
+        stderr_task = asyncio.ensure_future(drain_stderr(proc.stderr, tail))
+        state: Any = self.config.create_stream_state()
+        exit_code: int | None = None
+        try:
+            await send_stdin(proc, request.stdin)
+            async for raw in iter_stream_lines(proc.stdout):
+                line = decode_json_line(raw)
+                if line is None:
+                    continue
+                for event in self.config.transform_stream_line(line, state):
+                    yield event
+                self._native_id = self.config.get_native_session_id(state) or self._native_id
+            exit_code = await proc.wait()
+            await stderr_task
+        finally:
+            self._proc = None
+            if exit_code is None:
+                # Consumer stopped early, timed out or errored: don't leave the runtime running.
+                await proc.kill()
+            if not stderr_task.done():
+                stderr_task.cancel()
+        response = self.config.transform_turn_response(ctx, state, exit_code, tuple(tail))
+        ctx.final_text = response.final_text
+        ctx.output_json = response.output_json
+
+    async def stop(self, ctx: SessionContext) -> None:
+        proc, self._proc = self._proc, None
+        if proc is not None:
+            await proc.kill()
+
+    def native_session_id(self) -> str | None:
+        return self._native_id
+
+    async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
+        self._native_id = native_session_id

litellm/harness/handlers/deepagents_handler.py

@@ -0,0 +1,261 @@
+"""
+In-process handler for Deep Agents.
+
+Deep Agents is a Python library, so there is no process or model endpoint: the handler
+builds the agent with a LiteLLM chat model, streams the LangGraph run, turns interrupts into
+Approval events and counts usage. Translation lives in
+`litellm/llms/deepagents/harness/transformation.py`.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass
+from types import MappingProxyType, ModuleType
+from typing import TYPE_CHECKING, Any
+
+from litellm.harness.context import SessionContext
+from litellm.harness.errors import HarnessError, HarnessInstallFailed
+from litellm.harness.handlers.base import BaseHarnessHandler
+from litellm.harness.handlers.cli_handler import copy_skills
+from litellm.harness.options import DeepAgentsOptions
+from litellm.harness.types import Approval, Event
+from litellm.llms.deepagents.harness.transformation import (
+    EXECUTE_TOOLS,
+    INSTALL_HINT,
+    SKILLS_DIR,
+    WRITE_TOOLS,
+    TurnState,
+    approval_requests,
+    blocked_tools,
+    chat_model_kwargs,
+    decision,
+    final_ai_text,
+    interrupt_config,
+    interrupts_in,
+    normalized_tool_name,
+    recursion_limit,
+    stream_events,
+    structured_json,
+    update_events,
+)
+
+if TYPE_CHECKING:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.language_models import BaseChatModel
+    from langchain_core.runnables import RunnableConfig
+    from langgraph.checkpoint.base import BaseCheckpointSaver
+    from langgraph.graph.state import CompiledStateGraph
+    from langgraph.types import Command
+
+    from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
+
+_MODEL_NODE = "model"
+
+
+@dataclass(frozen=True)
+class DeepAgentsDeps:
+    """The optional-dependency entrypoints this handler uses."""
+
+    create_deep_agent: Any
+    chat_litellm: Any
+    checkpointer_cls: Any
+    command_cls: Any
+    subagent_defaults: Mapping[str, Any]
+    convert_to_openai_messages: Any
+    backend: ModuleType
+
+
+def load_deps() -> DeepAgentsDeps:
+    """Import deepagents + langchain-litellm, or raise HarnessInstallFailed."""
+    try:
+        deepagents = importlib.import_module("deepagents")
+        subagents = importlib.import_module("deepagents.middleware.subagents")
+        chat = importlib.import_module("langchain_litellm")
+        memory = importlib.import_module("langgraph.checkpoint.memory")
+        lg_types = importlib.import_module("langgraph.types")
+        messages = importlib.import_module("langchain_core.messages")
+        backend = importlib.import_module("litellm.llms.deepagents.harness.sandbox_backend")
+    except ImportError as e:
+        raise HarnessInstallFailed(f"{INSTALL_HINT} ({e})") from e
+    return DeepAgentsDeps(
+        create_deep_agent=deepagents.create_deep_agent,
+        chat_litellm=chat.ChatLiteLLM,
+        checkpointer_cls=memory.InMemorySaver,
+        command_cls=lg_types.Command,
+        subagent_defaults=subagents.GENERAL_PURPOSE_SUBAGENT,
+        convert_to_openai_messages=messages.convert_to_openai_messages,
+        backend=backend,
+    )
+
+
+_SHARED_CHECKPOINTER: dict[str, Any] = {}  # mutable-ok: process-wide lazy singleton slot for the in-memory checkpointer
+
+
+def shared_checkpointer(deps: DeepAgentsDeps) -> BaseCheckpointSaver:
+    """One in-memory checkpointer per process, so resume() works across sessions in-process."""
+    saver = _SHARED_CHECKPOINTER.get("saver")
+    if saver is None:
+        saver = deps.checkpointer_cls()
+        _SHARED_CHECKPOINTER["saver"] = saver
+    return saver
+
+
+def build_chat_model(ctx: SessionContext, deps: DeepAgentsDeps) -> BaseChatModel:
+    """The LangChain chat model for this session. Tests monkeypatch this."""
+    return deps.chat_litellm(**chat_model_kwargs(ctx))
+
+
+class DeepAgentsHandler(BaseHarnessHandler):
+    def __init__(self, config: BaseHarnessConfig) -> None:
+        super().__init__(config)
+        self._deps: DeepAgentsDeps | None = None
+        self._agent: Any = None
+        self._thread_id: str | None = None
+        self._skip_tools: frozenset[str] = frozenset()
+
+    async def start(self, ctx: SessionContext) -> None:
+        self.config.validate_environment(ctx)
+        deps = load_deps()
+        self._deps = deps
+        blocked = blocked_tools(ctx.permissions, ctx.disable_tools)
+        backend = deps.backend.SandboxBackend(
+            ctx.sandbox,
+            loop=asyncio.get_running_loop(),
+            writable=WRITE_TOOLS.isdisjoint(blocked),
+            allow_execute=EXECUTE_TOOLS.isdisjoint(blocked),
+        )
+        self._agent = deps.create_deep_agent(
+            model=build_chat_model(ctx, deps),
+            tools=list(ctx.tools),  # mutable-ok: deepagents create_deep_agent(tools=) takes a list
+            system_prompt=ctx.instructions,
+            middleware=self._middleware(deps, blocked),
+            subagents=self._subagents(ctx, deps, blocked),
+            skills=await self._install_skills(ctx),
+            backend=backend,
+            interrupt_on=interrupt_config(ctx.permissions, blocked),
+            response_format=ctx.output,
+            checkpointer=shared_checkpointer(deps),
+        )
+        self._skip_tools = frozenset({ctx.output.__name__}) if ctx.output is not None else frozenset()
+        if self._thread_id is None:
+            self._thread_id = ctx.session_id
+
+    async def stop(self, ctx: SessionContext) -> None:
+        self._agent = None
+
+    def native_session_id(self) -> str | None:
+        return self._thread_id
+
+    async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
+        self._thread_id = native_session_id
+
+    async def history(
+        self, ctx: SessionContext
+    ) -> list[dict[str, Any]]:  # mutable-ok: BaseHarnessHandler.history API returns OpenAI message dicts
+        agent, deps = self._require_agent()
+        snapshot = await agent.aget_state(self._run_config(ctx, None))
+        messages = (snapshot.values or MappingProxyType({})).get("messages") or ()
+        converted: list[dict[str, Any]] = deps.convert_to_openai_messages(  # mutable-ok: LangChain returns a list
+            messages
+        )
+        return converted
+
+    async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
+        agent, deps = self._require_agent()
+        run_config = self._run_config(ctx, deps.backend.UsageCallback(ctx, ctx.model))
+        user_message = {"role": "user", "content": prompt}  # mutable-ok: LangGraph input message dict
+        payload: dict[str, object] | Command = {"messages": [user_message]}  # mutable-ok: LangGraph input state
+        while True:
+            state = TurnState()
+            async for event in self._stream_pass(agent, payload, run_config, state):
+                yield event
+            if not state.interrupts:
+                break
+            resume: dict[str, Any] = {}  # mutable-ok: Command(resume=) payload, filled per answered approval
+            for interrupt in state.interrupts:
+                decisions: list[dict[str, Any]] = []  # mutable-ok: HITL decisions collected across awaited approvals
+                for request in approval_requests(getattr(interrupt, "value", None)):
+                    approval = Approval(
+                        tool=normalized_tool_name(str(request.get("name") or "")),
+                        input=dict(request.get("args") or ()),  # mutable-ok: Approval.input is a public dict field
+                    )
+                    yield approval
+                    decisions.append(decision(*await approval.wait()))
+                resume[interrupt.id] = {"decisions": decisions}  # mutable-ok: LangGraph HITL resume payload
+            payload = deps.command_cls(resume=resume)
+        await self._finish_turn(ctx, agent, run_config)
+
+    async def _stream_pass(
+        self,
+        agent: CompiledStateGraph,
+        payload: dict[str, object] | Command,  # mutable-ok: LangGraph astream input type
+        run_config: RunnableConfig,
+        state: TurnState,
+    ) -> AsyncIterator[Event]:
+        async for part in agent.astream(
+            payload,
+            run_config,
+            stream_mode=["messages", "updates"],  # mutable-ok: LangGraph stream_mode takes a list
+        ):
+            # A list stream_mode yields (mode, chunk) tuples; LangGraph's overloads don't say so.
+            if not isinstance(part, tuple) or len(part) != 2:
+                continue
+            mode, chunk = part
+            if mode == "messages":
+                message, meta = chunk
+                if isinstance(meta, Mapping) and meta.get("langgraph_node") == _MODEL_NODE:
+                    for event in stream_events(message):
+                        yield event
+            elif mode == "updates":
+                state.interrupts = (*state.interrupts, *interrupts_in(chunk))
+                for event in update_events(chunk, self._skip_tools):
+                    yield event
+
+    async def _finish_turn(self, ctx: SessionContext, agent: CompiledStateGraph, run_config: RunnableConfig) -> None:
+        snapshot = await agent.aget_state(run_config)
+        values = snapshot.values or MappingProxyType({})
+        ctx.final_text = final_ai_text(values.get("messages") or ())
+        if ctx.output is not None:
+            ctx.output_json = structured_json(values.get("structured_response"))
+
+    def _require_agent(self) -> tuple[Any, DeepAgentsDeps]:
+        if self._agent is None or self._deps is None:
+            raise HarnessError("Deep Agents session is not started")
+        return self._agent, self._deps
+
+    def _run_config(self, ctx: SessionContext, usage_callback: BaseCallbackHandler | None) -> RunnableConfig:
+        run_config: RunnableConfig = {
+            "configurable": {"thread_id": self._thread_id or ctx.session_id},
+            "recursion_limit": recursion_limit(ctx),
+        }
+        if usage_callback is not None:
+            run_config["callbacks"] = [usage_callback]  # mutable-ok: LangChain RunnableConfig.callbacks is a list
+        return run_config
+
+    @staticmethod
+    def _middleware(deps: DeepAgentsDeps, blocked: frozenset[str]) -> list[Any]:  # mutable-ok: deepagents API
+        filters = (deps.backend.ToolFilterMiddleware(blocked),) if blocked else ()
+        return list(filters)  # mutable-ok: deepagents create_deep_agent(middleware=) takes a list
+
+    def _subagents(
+        self, ctx: SessionContext, deps: DeepAgentsDeps, blocked: frozenset[str]
+    ) -> list[Any]:  # mutable-ok: deepagents create_deep_agent(subagents=) takes a list
+        """User subagents, plus a general-purpose one that honours disable_tools when set."""
+        options = ctx.options if isinstance(ctx.options, DeepAgentsOptions) else None
+        user_subagents = tuple(options.subagents) if options is not None else ()
+        has_general = any(
+            isinstance(s, Mapping) and s.get("name") == deps.subagent_defaults["name"] for s in user_subagents
+        )
+        spec = {**deps.subagent_defaults, "middleware": self._middleware(deps, blocked)}  # mutable-ok: SubAgent dict
+        general = (spec,) if blocked and not has_general else ()
+        return [*general, *user_subagents]  # mutable-ok: deepagents create_deep_agent(subagents=) takes a list
+
+    @staticmethod
+    async def _install_skills(ctx: SessionContext) -> list[str] | None:  # mutable-ok: deepagents skills= takes a list
+        if not ctx.skills:
+            return None
+        await copy_skills(ctx.sandbox, ctx.skills, f"{ctx.sandbox.workdir}/{SKILLS_DIR}")
+        return [f"/{SKILLS_DIR}/"]  # mutable-ok: deepagents create_deep_agent(skills=) takes a list

litellm/harness/options.py

@@ -0,0 +1,37 @@
+"""Typed per-harness options. Settings that only make sense for one runtime live here."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+@dataclass(frozen=True)
+class ClaudeCodeOptions:
+    config: Mapping[str, Any] = field(default_factory=dict)
+    env: Mapping[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class CodexOptions:
+    reasoning_effort: Literal["low", "medium", "high", "xhigh"] | None = None
+    web_search: bool = False
+    config: Mapping[str, Any] = field(default_factory=dict)
+    env: Mapping[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class OpenCodeOptions:
+    agent: str = "build"
+    config: Mapping[str, Any] = field(default_factory=dict)
+    env: Mapping[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class DeepAgentsOptions:
+    subagents: Sequence[Any] = ()
+    recursion_limit: int | None = None
+
+
+HarnessOptions = ClaudeCodeOptions | CodexOptions | OpenCodeOptions | DeepAgentsOptions

litellm/harness/sandbox/__init__.py

@@ -0,0 +1,25 @@
+"""Sandboxes for litellm.harness: where the runtime runs and which files it can touch."""
+
+from litellm.harness.sandbox.base import CompletedRun, Process, Sandbox
+from litellm.harness.sandbox.docker import DockerSandbox, docker
+from litellm.harness.sandbox.local import LocalSandbox, local
+from litellm.harness.sandbox.snapshot import (
+    build_file_changes,
+    capture_text_contents,
+    diff_snapshots,
+    snapshot_local,
+)
+
+__all__ = (
+    "CompletedRun",
+    "DockerSandbox",
+    "LocalSandbox",
+    "Process",
+    "Sandbox",
+    "build_file_changes",
+    "capture_text_contents",
+    "diff_snapshots",
+    "docker",
+    "local",
+    "snapshot_local",
+)

litellm/harness/sandbox/base.py

@@ -0,0 +1,62 @@
+"""The Sandbox protocol: where a harness runtime runs and which files it can touch."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Protocol, runtime_checkable
+
+
+@dataclass(frozen=True)
+class CompletedRun:
+    stdout: str
+    stderr: str
+    exit_code: int
+
+
+@runtime_checkable
+class Process(Protocol):
+    stdin: asyncio.StreamWriter | None
+    stdout: asyncio.StreamReader
+    stderr: asyncio.StreamReader
+
+    async def wait(self) -> int: ...
+
+    async def kill(self) -> None: ...
+
+
+@runtime_checkable
+class Sandbox(Protocol):
+    workdir: str
+
+    async def exec(
+        self,
+        cmd: Sequence[str],
+        *,
+        env: Mapping[str, str] | None = None,
+        cwd: str | None = None,
+    ) -> Process: ...
+
+    async def run(
+        self,
+        cmd: Sequence[str],
+        *,
+        env: Mapping[str, str] | None = None,
+        cwd: str | None = None,
+        timeout: float | None = None,
+    ) -> CompletedRun: ...
+
+    async def read(self, path: str) -> bytes: ...
+
+    async def write(self, path: str, data: bytes) -> None: ...
+
+    def host_url(self, port: int) -> str: ...
+
+    async def which(self, binary: str) -> str | None: ...
+
+    async def snapshot(self) -> Mapping[str, str]: ...
+
+    async def tempdir(self) -> str: ...
+
+    async def close(self) -> None: ...

litellm/harness/sandbox/docker.py

@@ -0,0 +1,298 @@
+"""DockerSandbox: run the harness runtime inside a container via the docker CLI."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import posixpath
+import shutil
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
+from typing import Final
+
+from litellm.constants import HARNESS_SNAPSHOT_SKIP_DIRS
+from litellm.harness.errors import SandboxError
+from litellm.harness.sandbox.base import CompletedRun
+from litellm.harness.sandbox.local import SubprocessHandle, collect_output
+from litellm.harness.sandbox.snapshot import HARNESS_SNAPSHOT_MAX_FILE_BYTES
+
+DOCKER_HOST_ALIAS: Final = "host.docker.internal"
+_SHA256_HEX_LEN: Final = 64
+_WRITE_SCRIPT: Final = 'mkdir -p "$(dirname "$1")" && cat > "$1"'
+_WHICH_SCRIPT: Final = 'command -v "$1"'
+
+
+def _snapshot_script() -> str:
+    prune = " -o ".join(f"-name '{name}'" for name in sorted(HARNESS_SNAPSHOT_SKIP_DIRS))
+    return (
+        'cd "$1" && find . -type d \\( '
+        + prune
+        + " \\) -prune -o -type f -size -"
+        + f"{HARNESS_SNAPSHOT_MAX_FILE_BYTES + 1}c"
+        + " -exec sha256sum {} +"
+    )
+
+
+def parse_sha256sum(output: str) -> Mapping[str, str]:
+    """Parse `sha256sum` lines ("<hex>  ./rel/path") into {rel/path: hex}."""
+    return MappingProxyType(
+        {
+            line[_SHA256_HEX_LEN + 2 :].removeprefix("./"): line[:_SHA256_HEX_LEN]
+            for line in output.splitlines()
+            if len(line) > _SHA256_HEX_LEN + 2
+        }
+    )
+
+
+class DockerSandbox:
+    """Sandbox backed by a long-lived `sleep infinity` container."""
+
+    # Harness configs read this to skip a runtime's own nested OS sandbox.
+    is_container = True
+
+    def __init__(
+        self,
+        image: str,
+        mounts: Mapping[str | os.PathLike[str], str] | None = None,
+        workdir: str = "/workspace",
+        env: Mapping[str, str] | None = None,
+        name: str | None = None,
+    ) -> None:
+        if not image:
+            raise SandboxError("docker sandbox needs an image")
+        if not posixpath.isabs(workdir):
+            raise SandboxError(f"docker workdir must be absolute: {workdir}")
+        self.image = image
+        self.workdir: str = posixpath.normpath(workdir)
+        self.mounts: Mapping[str, str] = MappingProxyType(
+            {os.path.abspath(os.fspath(host)): container for host, container in (mounts.items() if mounts else ())}
+        )
+        self.env: Mapping[str, str] = MappingProxyType(dict(env or ()))
+        self.name = name
+        self.container_id: str | None = None
+        self._start_lock = asyncio.Lock()
+        self._processes: set[SubprocessHandle] = set()  # mutable-ok: live-process registry (add/discard)
+        self._closed = False
+
+    def __repr__(self) -> str:
+        return f"DockerSandbox({self.image!r}, workdir={self.workdir!r})"
+
+    # -- docker CLI plumbing (tests monkeypatch these two) ---------------------
+
+    def _docker_binary(self) -> str:
+        binary = shutil.which("docker")
+        if binary is None:
+            raise SandboxError(
+                "docker sandbox requires the `docker` CLI on PATH; install Docker or use sandbox.local(path)"
+            )
+        return binary
+
+    async def _spawn(self, args: Sequence[str]) -> SubprocessHandle:
+        """Start `docker <args>` with stdin/stdout/stderr pipes."""
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                self._docker_binary(),
+                *args,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+        except (FileNotFoundError, PermissionError) as exc:
+            raise SandboxError(f"could not run docker: {exc}") from exc
+        return SubprocessHandle(proc)
+
+    async def _docker(
+        self,
+        args: Sequence[str],
+        *,
+        input: bytes | None = None,
+        timeout: float | None = None,
+    ) -> tuple[int, bytes, bytes]:
+        """Run `docker <args>` to completion; returns (exit_code, stdout, stderr)."""
+        handle = await self._spawn(args)
+        try:
+            return await asyncio.wait_for(_communicate(handle, input), timeout)
+        except asyncio.TimeoutError:
+            await handle.kill()
+            raise SandboxError(f"docker {args[0]} timed out after {timeout}s")
+
+    # -- command construction --------------------------------------------------
+
+    def run_args(
+        self,
+    ) -> list[str]:  # mutable-ok: argv is returned as a list, the shape callers and tests compare against
+        name_args = ("--name", self.name) if self.name else ()
+        mount_args = tuple(
+            arg for host, container in self.mounts.items() for arg in ("-v", f"{host}:{container}")
+        )  # comprehension-ok: flattens (flag, value) pairs into argv
+        env_args = tuple(
+            arg for key, value in self.env.items() for arg in ("-e", f"{key}={value}")
+        )  # comprehension-ok: flattens (flag, value) pairs into argv
+        return [  # mutable-ok: argv is returned as a list, the shape callers and tests compare against
+            "run",
+            "-d",
+            "--rm",
+            f"--add-host={DOCKER_HOST_ALIAS}:host-gateway",
+            *name_args,
+            *mount_args,
+            *env_args,
+            "-w",
+            self.workdir,
+            self.image,
+            "sleep",
+            "infinity",
+        ]
+
+    def exec_args(
+        self,
+        container_id: str,
+        cmd: Sequence[str],
+        *,
+        env: Mapping[str, str] | None = None,
+        cwd: str | None = None,
+    ) -> list[str]:  # mutable-ok: argv is returned as a list, the shape callers and tests compare against
+        env_args = tuple(
+            arg for key, value in (env.items() if env else ()) for arg in ("-e", f"{key}={value}")
+        )  # comprehension-ok: flattens (flag, value) pairs into argv
+        return [  # mutable-ok: argv is returned as a list, the shape callers and tests compare against
+            "exec",
+            "-i",
+            "-w",
+            self.container_path(cwd or self.workdir),
+            *env_args,
+            container_id,
+            *cmd,
+        ]
+
+    def container_path(self, path: str) -> str:
+        """Absolute container path; relative paths resolve against workdir."""
+        joined = path if posixpath.isabs(path) else posixpath.join(self.workdir, path)
+        return posixpath.normpath(joined)
+
+    # -- lifecycle ---------------------------------------------------------------
+
+    async def start(self) -> str:
+        """Start the container if needed and return its id."""
+        if self._closed:
+            raise SandboxError("sandbox is closed")
+        async with self._start_lock:
+            if self.container_id is not None:
+                return self.container_id
+            code, out, err = await self._docker(self.run_args())
+            if code != 0:
+                raise SandboxError(f"docker run {self.image} failed ({code}): {err.decode(errors='replace').strip()}")
+            container_id = out.decode().strip()
+            if not container_id:
+                raise SandboxError("docker run returned no container id")
+            self.container_id = container_id
+            return container_id
+
+    async def _exec_capture(self, cmd: Sequence[str], *, input: bytes | None = None) -> tuple[int, bytes, bytes]:
+        container_id = await self.start()
+        return await self._docker(self.exec_args(container_id, cmd), input=input)
+
+    # -- Sandbox protocol --------------------------------------------------------
+
+    async def exec(
+        self,
+        cmd: Sequence[str],
+        *,
+        env: Mapping[str, str] | None = None,
+        cwd: str | None = None,
+    ) -> SubprocessHandle:
+        if not cmd:
+            raise SandboxError("exec() needs a non-empty command")
+        container_id = await self.start()
+        handle = await self._spawn(self.exec_args(container_id, cmd, env=env, cwd=cwd))
+        self._processes.add(handle)
+        return handle
+
+    async def run(
+        self,
+        cmd: Sequence[str],
+        *,
+        env: Mapping[str, str] | None = None,
+        cwd: str | None = None,
+        timeout: float | None = None,
+    ) -> CompletedRun:
+        handle = await self.exec(cmd, env=env, cwd=cwd)
+        try:
+            return await collect_output(handle, cmd, timeout)
+        finally:
+            self._processes.discard(handle)
+
+    async def read(self, path: str) -> bytes:
+        target = self.container_path(path)
+        code, out, err = await self._exec_capture(("cat", target))
+        if code != 0:
+            raise SandboxError(f"could not read {target}: {err.decode(errors='replace').strip()}")
+        return out
+
+    async def write(self, path: str, data: bytes) -> None:
+        target = self.container_path(path)
+        code, _, err = await self._exec_capture(("sh", "-c", _WRITE_SCRIPT, "sh", target), input=data)
+        if code != 0:
+            raise SandboxError(f"could not write {target}: {err.decode(errors='replace').strip()}")
+
+    def host_url(self, port: int) -> str:
+        return f"http://{DOCKER_HOST_ALIAS}:{port}"
+
+    async def which(self, binary: str) -> str | None:
+        code, out, _ = await self._exec_capture(("sh", "-lc", _WHICH_SCRIPT, "sh", binary))
+        found = out.decode(errors="replace").strip()
+        return found if code == 0 and found else None
+
+    async def tempdir(self) -> str:
+        """A fresh `mktemp -d` directory inside the container."""
+        code, out, err = await self._exec_capture(("mktemp", "-d"))
+        path = out.decode(errors="replace").strip()
+        if code != 0 or not path:
+            raise SandboxError(f"mktemp -d failed: {err.decode(errors='replace').strip()}")
+        return path
+
+    async def snapshot(self) -> Mapping[str, str]:
+        code, out, err = await self._exec_capture(("sh", "-c", _snapshot_script(), "sh", self.workdir))
+        if code != 0:
+            raise SandboxError(f"snapshot failed: {err.decode(errors='replace').strip()}")
+        return parse_sha256sum(out.decode("utf-8", errors="replace"))
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        live = tuple(h for h in self._processes if h.returncode is None)
+        await asyncio.gather(*(h.kill() for h in live), return_exceptions=True)
+        self._processes.clear()
+        if self.container_id is not None:
+            container_id, self.container_id = self.container_id, None
+            await self._docker(
+                ["rm", "-f", container_id]  # mutable-ok: argv list, the shape _spawn records and tests assert on
+            )
+
+    async def __aenter__(self) -> DockerSandbox:
+        await self.start()
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.close()
+
+
+async def _communicate(handle: SubprocessHandle, data: bytes | None) -> tuple[int, bytes, bytes]:
+    if handle.stdin is not None:
+        if data:
+            handle.stdin.write(data)
+            await handle.stdin.drain()
+        handle.stdin.close()
+    stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read())
+    return await handle.wait(), stdout, stderr
+
+
+def docker(
+    image: str,
+    mounts: Mapping[str | os.PathLike[str], str] | None = None,
+    workdir: str = "/workspace",
+    env: Mapping[str, str] | None = None,
+    name: str | None = None,
+) -> DockerSandbox:
+    """Sandbox in a new container of `image`, started lazily on first use."""
+    return DockerSandbox(image, mounts=mounts, workdir=workdir, env=env, name=name)

litellm/harness/sandbox/local.py

@@ -0,0 +1,277 @@
+"""LocalSandbox: run the harness runtime as a subprocess on this machine."""
+
+from __future__ import annotations
+
+import asyncio
+import itertools
+import os
+import shutil
+import signal
+import tempfile
+from collections.abc import Mapping, Sequence
+from types import MappingProxyType
+from typing import Final
+
+from litellm.constants import HARNESS_PROCESS_KILL_GRACE_SECONDS
+from litellm.harness.errors import SandboxError
+from litellm.harness.sandbox.base import CompletedRun
+from litellm.harness.sandbox.snapshot import snapshot_local
+
+_SECRET_PREFIXES: Final = (
+    "ANTHROPIC_",
+    "OPENAI_",
+    "LITELLM_",
+    "AZURE_",
+    "AWS_",
+    "GEMINI_",
+    "CODEX_",
+    "CURSOR_",
+    "VERTEX",
+    # A parent Claude Code session's socket/session vars make a child `claude` attach to
+    # the parent's login instead of the harness token.
+    "CLAUDE_CODE_",
+    "CLAUDE_PID",
+    "CLAUDECODE",
+)
+_SECRET_NAMES: Final = frozenset({"GOOGLE_API_KEY", "GOOGLE_APPLICATION_CREDENTIALS"})
+_SECRET_SUBSTRINGS: Final = ("API_KEY", "TOKEN", "SECRET")
+_TEMPDIR_PREFIX: Final = "litellm-harness-"
+
+
+def is_secret_env_name(name: str) -> bool:
+    """True if an env var name looks like a provider credential."""
+    upper = name.upper()
+    if upper in _SECRET_NAMES or upper.startswith(_SECRET_PREFIXES):
+        return True
+    return any(part in upper for part in _SECRET_SUBSTRINGS)
+
+
+def filtered_environ(
+    base: Mapping[str, str] | None = None,
+    extra: Mapping[str, str] | None = None,
+) -> Mapping[str, str]:
+    """base (default os.environ) without provider secrets, then extra on top."""
+    source = os.environ if base is None else base
+    kept = ((k, v) for k, v in source.items() if not is_secret_env_name(k))
+    overlay = extra.items() if extra else ()
+    return MappingProxyType(dict(itertools.chain(kept, overlay)))
+
+
+def _signal_process(proc: asyncio.subprocess.Process, sig: int) -> None:
+    try:
+        os.killpg(proc.pid, sig)
+    except (ProcessLookupError, PermissionError, OSError):
+        try:
+            proc.send_signal(sig)
+        except ProcessLookupError:
+            pass
+
+
+class SubprocessHandle:
+    """Process-protocol wrapper around an asyncio subprocess."""
+
+    def __init__(self, proc: asyncio.subprocess.Process) -> None:
+        if proc.stdout is None or proc.stderr is None:
+            raise SandboxError("subprocess was started without stdout/stderr pipes")
+        self._proc = proc
+        self.stdin: asyncio.StreamWriter | None = proc.stdin
+        self.stdout: asyncio.StreamReader = proc.stdout
+        self.stderr: asyncio.StreamReader = proc.stderr
+
+    @property
+    def pid(self) -> int:
+        return self._proc.pid
+
+    @property
+    def returncode(self) -> int | None:
+        return self._proc.returncode
+
+    async def wait(self) -> int:
+        return await self._proc.wait()
+
+    async def kill(self) -> None:
+        """SIGTERM, wait HARNESS_PROCESS_KILL_GRACE_SECONDS, then SIGKILL."""
+        if self._proc.returncode is not None:
+            return
+        _signal_process(self._proc, signal.SIGTERM)
+        try:
+            await asyncio.wait_for(self._proc.wait(), timeout=HARNESS_PROCESS_KILL_GRACE_SECONDS)
+            return
+        except asyncio.TimeoutError:
+            pass
+        _signal_process(self._proc, signal.SIGKILL)
+        await self._proc.wait()
+
+
+async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
+    if handle.stdin is not None:
+        handle.stdin.close()
+    stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read())
+    exit_code = await handle.wait()
+    return stdout, stderr, exit_code
+
+
+async def collect_output(handle: SubprocessHandle, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
+    """Close stdin, read stdout/stderr to EOF; kill and raise SandboxError on timeout."""
+    try:
+        stdout, stderr, code = await asyncio.wait_for(_read_all(handle), timeout)
+    except asyncio.TimeoutError:
+        await handle.kill()
+        raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}")
+    return CompletedRun(
+        stdout=stdout.decode("utf-8", errors="replace"),
+        stderr=stderr.decode("utf-8", errors="replace"),
+        exit_code=code,
+    )
+
+
+def _is_within(path: str, root: str) -> bool:
+    return path == root or path.startswith(root.rstrip(os.sep) + os.sep)
+
+
+class LocalSandbox:
+    """Sandbox backed by the local filesystem and asyncio subprocesses."""
+
+    def __init__(self, path: str | os.PathLike[str]) -> None:
+        resolved = os.path.realpath(os.path.abspath(os.fspath(path)))
+        if not os.path.isdir(resolved):
+            raise SandboxError(f"sandbox path does not exist or is not a directory: {resolved}")
+        self.workdir: str = resolved
+        self._processes: set[SubprocessHandle] = set()  # mutable-ok: live-process registry (add/discard)
+        self._tempdirs: list[str] = []  # mutable-ok: tempdirs created on demand by tempdir(), removed on close()
+        self._closed = False
+
+    def __repr__(self) -> str:
+        return f"LocalSandbox({self.workdir!r})"
+
+    def _check_open(self) -> None:
+        if self._closed:
+            raise SandboxError("sandbox is closed")
+
+    def _allowed_roots(self) -> tuple[str, ...]:
+        return (self.workdir, *self._tempdirs)
+
+    def resolve_path(self, path: str) -> str:
+        """Absolute real path for path; SandboxError if it escapes the sandbox."""
+        joined = path if os.path.isabs(path) else os.path.join(self.workdir, path)
+        real = os.path.realpath(joined)
+        if not any(_is_within(real, root) for root in self._allowed_roots()):
+            raise SandboxError(f"path escapes the sandbox: {path}")
+        return real
+
+    def _resolve_cwd(self, cwd: str | None) -> str:
+        if cwd is None:
+            return self.workdir
+        resolved = self.resolve_path(cwd)
+        if not os.path.isdir(resolved):
+            raise SandboxError(f"cwd is not a directory: {cwd}")
+        return resolved
+
+    def child_env(self, env: Mapping[str, str] | None = None) -> Mapping[str, str]:
+        return filtered_environ(extra=env)
+
+    async def exec(
+        self,
+        cmd: Sequence[str],
+        *,
+        env: Mapping[str, str] | None = None,
+        cwd: str | None = None,
+    ) -> SubprocessHandle:
+        self._check_open()
+        if not cmd:
+            raise SandboxError("exec() needs a non-empty command")
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                *cmd,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                cwd=self._resolve_cwd(cwd),
+                env=self.child_env(env),
+                start_new_session=True,
+            )
+        except (FileNotFoundError, PermissionError) as exc:
+            raise SandboxError(f"could not start {cmd[0]}: {exc}") from exc
+        handle = SubprocessHandle(proc)
+        self._processes.add(handle)
+        return handle
+
+    async def run(
+        self,
+        cmd: Sequence[str],
+        *,
+        env: Mapping[str, str] | None = None,
+        cwd: str | None = None,
+        timeout: float | None = None,
+    ) -> CompletedRun:
+        handle = await self.exec(cmd, env=env, cwd=cwd)
+        try:
+            return await collect_output(handle, cmd, timeout)
+        finally:
+            self._processes.discard(handle)
+
+    async def read(self, path: str) -> bytes:
+        self._check_open()
+        resolved = self.resolve_path(path)
+        try:
+            return await asyncio.to_thread(_read_bytes, resolved)
+        except OSError as exc:
+            raise SandboxError(f"could not read {path}: {exc}") from exc
+
+    async def write(self, path: str, data: bytes) -> None:
+        self._check_open()
+        resolved = self.resolve_path(path)
+        try:
+            await asyncio.to_thread(_write_bytes, resolved, data)
+        except OSError as exc:
+            raise SandboxError(f"could not write {path}: {exc}") from exc
+
+    def host_url(self, port: int) -> str:
+        return f"http://127.0.0.1:{port}"
+
+    async def which(self, binary: str) -> str | None:
+        return shutil.which(binary, path=self.child_env().get("PATH"))
+
+    async def tempdir(self) -> str:
+        """A private temp dir (e.g. for CODEX_HOME), removed on close()."""
+        self._check_open()
+        path = os.path.realpath(tempfile.mkdtemp(prefix=_TEMPDIR_PREFIX))
+        self._tempdirs.append(path)
+        return path
+
+    async def snapshot(self) -> Mapping[str, str]:
+        self._check_open()
+        return await snapshot_local(self.workdir)
+
+    async def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        live = tuple(h for h in self._processes if h.returncode is None)
+        await asyncio.gather(*(h.kill() for h in live), return_exceptions=True)
+        self._processes.clear()
+        for path in self._tempdirs:
+            shutil.rmtree(path, ignore_errors=True)
+        self._tempdirs.clear()
+
+    async def __aenter__(self) -> LocalSandbox:
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        await self.close()
+
+
+def _read_bytes(path: str) -> bytes:
+    with open(path, "rb") as fh:
+        return fh.read()
+
+
+def _write_bytes(path: str, data: bytes) -> None:
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "wb") as fh:
+        fh.write(data)
+
+
+def local(path: str | os.PathLike[str]) -> LocalSandbox:
+    """Sandbox rooted at an existing local directory."""
+    return LocalSandbox(path)

litellm/harness/sandbox/snapshot.py

@@ -0,0 +1,183 @@
+"""Workspace snapshots and FileChange construction.
+
+A snapshot maps a workspace-relative POSIX path to the sha256 of its contents.
+Diffing two snapshots tells us which files a turn created, modified or deleted;
+`build_file_changes` turns that into `FileChange` events with unified diffs for
+small text files.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import difflib
+import functools
+import hashlib
+import os
+from collections.abc import Iterator, Mapping
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Final
+
+from litellm.constants import HARNESS_MAX_DIFF_BYTES, HARNESS_SNAPSHOT_SKIP_DIRS
+from litellm.harness.errors import HarnessError
+from litellm.harness.types import FileChange, FileChangeKind
+
+if TYPE_CHECKING:
+    from litellm.harness.sandbox.base import Sandbox
+
+# Files larger than this are left out of snapshots entirely.
+HARNESS_SNAPSHOT_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
+# Upper bound on bytes read by capture_text_contents() for one turn.
+HARNESS_SNAPSHOT_MAX_TOTAL_BYTES: Final = 16 * 1024 * 1024
+_HASH_CHUNK_BYTES: Final = 1024 * 1024
+_NO_NEWLINE_MARKER: Final = "\\ No newline at end of file\n"
+
+
+def _hash_file(path: str) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as fh:
+        for chunk in iter(functools.partial(fh.read, _HASH_CHUNK_BYTES), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _hash_entry(root: str, dirpath: str, filename: str) -> tuple[str, str] | None:
+    full = os.path.join(dirpath, filename)
+    try:
+        info = os.lstat(full)
+    except OSError:
+        return None
+    if not os.path.isfile(full) or os.path.islink(full):
+        return None
+    if info.st_size > HARNESS_SNAPSHOT_MAX_FILE_BYTES:
+        return None
+    try:
+        digest = _hash_file(full)
+    except OSError:
+        return None
+    rel = os.path.relpath(full, root).replace(os.sep, "/")
+    return rel, digest
+
+
+def _walk_entries(root: str) -> Iterator[tuple[str, str]]:
+    for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
+        dirnames[:] = [  # mutable-ok: os.walk prunes only via in-place mutation of its dirnames list
+            d for d in dirnames if d not in HARNESS_SNAPSHOT_SKIP_DIRS
+        ]
+        for filename in filenames:
+            entry = _hash_entry(root, dirpath, filename)
+            if entry is not None:
+                yield entry
+
+
+def snapshot_local_sync(root: str) -> Mapping[str, str]:
+    """Hash every regular file under root. Symlinks are never followed."""
+    return MappingProxyType(dict(_walk_entries(root)))
+
+
+async def snapshot_local(root: str) -> Mapping[str, str]:
+    """Async wrapper around snapshot_local_sync (runs in a worker thread)."""
+    return await asyncio.to_thread(snapshot_local_sync, root)
+
+
+def _change_kind(path: str, before: Mapping[str, str], after: Mapping[str, str]) -> FileChangeKind | None:
+    if path not in before:
+        return "created"
+    if path not in after:
+        return "deleted"
+    if before[path] != after[path]:
+        return "modified"
+    return None
+
+
+def diff_snapshots(
+    before: Mapping[str, str], after: Mapping[str, str]
+) -> list[tuple[str, FileChangeKind]]:  # mutable-ok: public sandbox helper; callers compare against a list
+    """Return (path, kind) for every changed file, sorted by path."""
+    kinds = ((path, _change_kind(path, before, after)) for path in sorted(frozenset(before) | frozenset(after)))
+    return [  # mutable-ok: public sandbox helper returns a list
+        (path, kind) for path, kind in kinds if kind is not None
+    ]
+
+
+def _as_text(data: bytes) -> str | None:
+    if len(data) > HARNESS_MAX_DIFF_BYTES or b"\0" in data:
+        return None
+    try:
+        return data.decode("utf-8")
+    except UnicodeDecodeError:
+        return None
+
+
+def unified_diff(path: str, old: str | None, new: str | None) -> str:
+    """Unified diff between two versions of path; None means the file is absent."""
+    from_file = "/dev/null" if old is None else f"a/{path}"
+    to_file = "/dev/null" if new is None else f"b/{path}"
+    lines = difflib.unified_diff(
+        (old or "").splitlines(keepends=True),
+        (new or "").splitlines(keepends=True),
+        fromfile=from_file,
+        tofile=to_file,
+    )
+    return "".join(line if line.endswith("\n") else line + "\n" + _NO_NEWLINE_MARKER for line in lines)
+
+
+async def _read_or_none(sandbox: Sandbox, path: str) -> bytes | None:
+    try:
+        return await sandbox.read(path)
+    except (HarnessError, OSError):
+        return None
+
+
+async def capture_text_contents(sandbox: Sandbox, paths_hashes: Mapping[str, str]) -> Mapping[str, bytes]:
+    """Read small text files before a turn so "modified"/"deleted" diffs can be built.
+
+    Each kept file is <= HARNESS_MAX_DIFF_BYTES; every byte read (kept or not) counts
+    toward HARNESS_SNAPSHOT_MAX_TOTAL_BYTES, after which capture stops.
+    """
+    captured: dict[str, bytes] = {}  # mutable-ok: async accumulator (awaits per read), frozen on return
+    total = 0
+    for path in sorted(paths_hashes):
+        if total >= HARNESS_SNAPSHOT_MAX_TOTAL_BYTES:
+            break
+        data = await _read_or_none(sandbox, path)
+        if data is None:
+            continue
+        total += len(data)
+        if _as_text(data) is not None:
+            captured[path] = data
+    return MappingProxyType(captured)
+
+
+async def _change_for(
+    sandbox: Sandbox,
+    path: str,
+    kind: FileChangeKind,
+    before_contents: Mapping[str, bytes],
+) -> FileChange:
+    old_bytes = before_contents.get(path)
+    old = _as_text(old_bytes) if old_bytes is not None else None
+    if kind == "deleted":
+        diff = unified_diff(path, old, None) if old is not None else None
+        return FileChange(path=path, kind=kind, diff=diff)
+    new_bytes = await _read_or_none(sandbox, path)
+    new = _as_text(new_bytes) if new_bytes is not None else None
+    if new is None or (kind == "modified" and old is None):
+        return FileChange(path=path, kind=kind, diff=None)
+    return FileChange(
+        path=path,
+        kind=kind,
+        diff=unified_diff(path, old if kind == "modified" else None, new),
+    )
+
+
+async def build_file_changes(
+    sandbox: Sandbox,
+    before: Mapping[str, str],
+    after: Mapping[str, str],
+    before_contents: Mapping[str, bytes] | None = None,
+) -> list[FileChange]:  # mutable-ok: feeds the public Result.files list
+    """FileChange per changed path. diff is None when it cannot be built as text."""
+    contents: Mapping[str, bytes] = before_contents or MappingProxyType({})
+    return [  # mutable-ok: feeds the public Result.files list
+        await _change_for(sandbox, path, kind, contents) for path, kind in diff_snapshots(before, after)
+    ]

litellm/harness/sync.py

@@ -0,0 +1,443 @@
+"""Sync API for litellm.harness: one daemon event-loop thread runs every async call."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import threading
+from collections.abc import AsyncIterator, Callable, Coroutine, Mapping, Sequence
+from concurrent.futures import Future
+from typing import (
+    Any,
+    TypeVar,
+)
+
+from pydantic import BaseModel
+
+from litellm.harness.context import ApprovalHandler
+from litellm.harness.options import HarnessOptions
+from litellm.harness.runtime import (
+    AsyncEventStream,
+    AsyncSession,
+    aagent_resume,
+    aagent_session,
+    arun_agent,
+    astream_agent,
+)
+from litellm.harness.sandbox.base import Sandbox
+from litellm.harness.types import (
+    Done,
+    Event,
+    Harness,
+    PermissionMode,
+    Result,
+    State,
+    Usage,
+)
+
+T = TypeVar("T")
+
+IN_LOOP_MESSAGE = (
+    "litellm.{name}() cannot be called from a running event loop; use `await litellm.a{name}(...)` instead"
+)
+
+
+class _LoopThread:
+    """A single background event loop shared by every sync call in the process."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._thread: threading.Thread | None = None
+
+    def loop(self) -> asyncio.AbstractEventLoop:
+        with self._lock:
+            if self._loop is None or self._thread is None or not self._thread.is_alive():
+                self._loop = asyncio.new_event_loop()
+                self._thread = threading.Thread(
+                    target=self._loop.run_forever,
+                    name="litellm-harness-loop",
+                    daemon=True,
+                )
+                self._thread.start()
+            return self._loop
+
+    def in_loop_thread(self) -> bool:
+        return self._thread is not None and threading.current_thread() is self._thread
+
+    def submit(self, coro: Coroutine[Any, Any, T]) -> Future[T]:
+        return asyncio.run_coroutine_threadsafe(coro, self.loop())
+
+
+_LOOP = _LoopThread()
+
+
+def _ensure_sync_context(name: str) -> None:
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return
+    raise RuntimeError(IN_LOOP_MESSAGE.format(name=name))
+
+
+def run_sync(coro: Coroutine[Any, Any, T], name: str) -> T:
+    """Run coro on the harness loop thread and block for its result."""
+    try:
+        _ensure_sync_context(name)
+    except RuntimeError:
+        coro.close()
+        raise
+    future = _LOOP.submit(coro)
+    try:
+        return future.result()
+    except KeyboardInterrupt:
+        future.cancel()
+        raise
+
+
+async def _anext(iterator: AsyncIterator[Event]) -> Event | None:
+    try:
+        return await iterator.__anext__()
+    except StopAsyncIteration:
+        return None
+
+
+async def _aclose_stream(stream: AsyncEventStream) -> None:
+    await stream.aclose()
+
+
+class EventStream:
+    """Sync iterator of events for one turn. `.result` is set once Done is seen."""
+
+    def __init__(self, stream: AsyncEventStream, name: str = "stream") -> None:
+        self._stream = stream
+        self._name = name
+        self._result: Result | None = None
+        self._finished = False
+
+    def __iter__(self) -> EventStream:
+        return self
+
+    def __next__(self) -> Event:
+        if self._finished:
+            raise StopIteration
+        event = run_sync(_anext(self._stream), self._name)
+        if event is None:
+            self._finished = True
+            raise StopIteration
+        if isinstance(event, Done):
+            self._result = event.result
+        return event
+
+    def __enter__(self) -> EventStream:
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.close()
+
+    @property
+    def result(self) -> Result | None:
+        return self._result
+
+    def cancel(self) -> None:
+        """Stop the turn. Iteration still ends with Done(stop_reason='cancelled')."""
+        _LOOP.loop().call_soon_threadsafe(self._stream.cancel)
+
+    def close(self) -> None:
+        """Abandon the stream and release the session behind it."""
+        if self._finished:
+            return
+        self._finished = True
+        run_sync(_aclose_stream(self._stream), self._name)
+
+
+class Session:
+    """Sync multi-turn session. Use as a context manager."""
+
+    def __init__(self, inner: AsyncSession) -> None:
+        self._inner = inner
+
+    @property
+    def aio(self) -> AsyncSession:
+        """The underlying AsyncSession (runs on the harness loop thread)."""
+        return self._inner
+
+    def start(self) -> Session:
+        run_sync(self._inner.start(), "session")
+        return self
+
+    def __enter__(self) -> Session:
+        return self.start()
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.close()
+
+    def run(self, prompt: str) -> Result:
+        return run_sync(self._inner.arun(prompt), "run")
+
+    def stream(self, prompt: str) -> EventStream:
+        return EventStream(self._inner.astream(prompt))
+
+    def close(self) -> None:
+        run_sync(self._inner.aclose(), "close")
+
+    def detach(self) -> State:
+        return run_sync(self._inner.adetach(), "detach")
+
+    def stop(self) -> State:
+        return run_sync(self._inner.astop(), "stop")
+
+    def history(
+        self,
+    ) -> list[dict[str, Any]]:  # mutable-ok: public API returns OpenAI-format message dicts from the handler
+        return run_sync(self._inner.history(), "history")
+
+    @property
+    def cost(self) -> float:
+        return self._inner.cost
+
+    @property
+    def usage(self) -> Usage:
+        return self._inner.usage
+
+    @property
+    def results(self) -> list[Result]:  # mutable-ok: public property; returns a detached copy of the session's results
+        return list(self._inner.results)  # mutable-ok: detached copy so callers cannot mutate the session's accumulator
+
+    @property
+    def session_id(self) -> str:
+        return self._inner.session_id
+
+
+def _run(
+    harness: Harness,
+    prompt: str,
+    *,
+    sandbox: Sandbox,
+    model: str | None = None,
+    api_key: str | None = None,
+    api_base: str | None = None,
+    instructions: str | None = None,
+    tools: Sequence[Callable[..., Any]] = (),
+    skills: Sequence[str | os.PathLike[str]] = (),
+    disable_tools: Sequence[str] = (),
+    permissions: PermissionMode = "full",
+    on_approval: ApprovalHandler | None = None,
+    output: type[BaseModel] | None = None,
+    max_turns: int | None = None,
+    timeout: float | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    options: HarnessOptions | None = None,
+    install: bool = False,
+) -> Result:
+    """Run one prompt to completion (blocking) and return the Result."""
+    return run_sync(
+        arun_agent(
+            harness,
+            prompt,
+            sandbox=sandbox,
+            model=model,
+            api_key=api_key,
+            api_base=api_base,
+            instructions=instructions,
+            tools=tools,
+            skills=skills,
+            disable_tools=disable_tools,
+            permissions=permissions,
+            on_approval=on_approval,
+            output=output,
+            max_turns=max_turns,
+            timeout=timeout,
+            metadata=metadata,
+            options=options,
+            install=install,
+        ),
+        "agent",
+    )
+
+
+def _stream(
+    harness: Harness,
+    prompt: str,
+    *,
+    sandbox: Sandbox,
+    model: str | None = None,
+    api_key: str | None = None,
+    api_base: str | None = None,
+    instructions: str | None = None,
+    tools: Sequence[Callable[..., Any]] = (),
+    skills: Sequence[str | os.PathLike[str]] = (),
+    disable_tools: Sequence[str] = (),
+    permissions: PermissionMode = "full",
+    on_approval: ApprovalHandler | None = None,
+    output: type[BaseModel] | None = None,
+    max_turns: int | None = None,
+    timeout: float | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    options: HarnessOptions | None = None,
+    install: bool = False,
+) -> EventStream:
+    """Stream events for one prompt (sync iterator). Validation errors raise here."""
+    _ensure_sync_context("agent")
+    inner = astream_agent(
+        harness,
+        prompt,
+        sandbox=sandbox,
+        model=model,
+        api_key=api_key,
+        api_base=api_base,
+        instructions=instructions,
+        tools=tools,
+        skills=skills,
+        disable_tools=disable_tools,
+        permissions=permissions,
+        on_approval=on_approval,
+        output=output,
+        max_turns=max_turns,
+        timeout=timeout,
+        metadata=metadata,
+        options=options,
+        install=install,
+    )
+    return EventStream(inner)
+
+
+def agent_session(
+    harness: Harness,
+    *,
+    sandbox: Sandbox,
+    model: str | None = None,
+    api_key: str | None = None,
+    api_base: str | None = None,
+    instructions: str | None = None,
+    tools: Sequence[Callable[..., Any]] = (),
+    skills: Sequence[str | os.PathLike[str]] = (),
+    disable_tools: Sequence[str] = (),
+    permissions: PermissionMode = "full",
+    on_approval: ApprovalHandler | None = None,
+    output: type[BaseModel] | None = None,
+    max_turns: int | None = None,
+    timeout: float | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    options: HarnessOptions | None = None,
+    install: bool = False,
+) -> Session:
+    """A multi-turn agent session: `with litellm.agent_session(...) as s: s.run(...)`."""
+    _ensure_sync_context("agent_session")
+    return Session(
+        aagent_session(
+            harness,
+            sandbox=sandbox,
+            model=model,
+            api_key=api_key,
+            api_base=api_base,
+            instructions=instructions,
+            tools=tools,
+            skills=skills,
+            disable_tools=disable_tools,
+            permissions=permissions,
+            on_approval=on_approval,
+            output=output,
+            max_turns=max_turns,
+            timeout=timeout,
+            metadata=metadata,
+            options=options,
+            install=install,
+        )
+    )
+
+
+def agent_resume(
+    state: State | bytes,
+    *,
+    sandbox: Sandbox,
+    model: str | None = None,
+    api_key: str | None = None,
+    api_base: str | None = None,
+    instructions: str | None = None,
+    tools: Sequence[Callable[..., Any]] = (),
+    skills: Sequence[str | os.PathLike[str]] = (),
+    disable_tools: Sequence[str] = (),
+    permissions: PermissionMode = "full",
+    on_approval: ApprovalHandler | None = None,
+    output: type[BaseModel] | None = None,
+    max_turns: int | None = None,
+    timeout: float | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    options: HarnessOptions | None = None,
+    install: bool = False,
+) -> Session:
+    """Continue a detached or stopped agent session from its State."""
+    _ensure_sync_context("agent_resume")
+    return Session(
+        aagent_resume(
+            state,
+            sandbox=sandbox,
+            model=model,
+            api_key=api_key,
+            api_base=api_base,
+            instructions=instructions,
+            tools=tools,
+            skills=skills,
+            disable_tools=disable_tools,
+            permissions=permissions,
+            on_approval=on_approval,
+            output=output,
+            max_turns=max_turns,
+            timeout=timeout,
+            metadata=metadata,
+            options=options,
+            install=install,
+        )
+    )
+
+
+def agent(
+    harness: Harness,
+    prompt: str,
+    *,
+    sandbox: Sandbox,
+    stream: bool = False,
+    model: str | None = None,
+    api_key: str | None = None,
+    api_base: str | None = None,
+    instructions: str | None = None,
+    tools: Sequence[Callable[..., Any]] = (),
+    skills: Sequence[str | os.PathLike[str]] = (),
+    disable_tools: Sequence[str] = (),
+    permissions: PermissionMode = "full",
+    on_approval: ApprovalHandler | None = None,
+    output: type[BaseModel] | None = None,
+    max_turns: int | None = None,
+    timeout: float | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    options: HarnessOptions | None = None,
+    install: bool = False,
+) -> Result | EventStream:
+    """Run an agent harness (Claude Code, Codex, OpenCode, Deep Agents) on one prompt.
+
+    Returns a Result. With stream=True it returns an iterator of events instead.
+    Prefix the model with `litellm_proxy/` to route every model call through your
+    LiteLLM AI Gateway.
+    """
+    kwargs: dict[str, Any] = {  # mutable-ok: forwarded as **kwargs to _run/_stream
+        "sandbox": sandbox,
+        "model": model,
+        "api_key": api_key,
+        "api_base": api_base,
+        "instructions": instructions,
+        "tools": tools,
+        "skills": skills,
+        "disable_tools": disable_tools,
+        "permissions": permissions,
+        "on_approval": on_approval,
+        "output": output,
+        "max_turns": max_turns,
+        "timeout": timeout,
+        "metadata": metadata,
+        "options": options,
+        "install": install,
+    }
+    if stream:
+        return _stream(harness, prompt, **kwargs)
+    return _run(harness, prompt, **kwargs)

litellm/harness/types.py

@@ -0,0 +1,214 @@
+"""Public types for litellm.harness: the Harness enum, events, results and state."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+from litellm.harness.errors import StateIncompatible
+
+StopReason = Literal["done", "max_turns", "timeout", "cancelled", "runtime_error"]
+PermissionMode = Literal["read-only", "ask", "edit", "full"]
+FileChangeKind = Literal["created", "modified", "deleted"]
+
+STATE_VERSION = 1
+
+
+class Harness(Enum):
+    """Supported agent runtimes. A plain Enum on purpose: strings are rejected."""
+
+    CLAUDE_CODE = "claude_code"
+    CODEX = "codex"
+    OPENCODE = "opencode"
+    DEEPAGENTS = "deepagents"
+
+
+def require_harness(harness: object) -> Harness:
+    """Return harness if it is a Harness member, else raise TypeError with a hint."""
+    if isinstance(harness, Harness):
+        return harness
+    hint = ""
+    if isinstance(harness, str):
+        normalized = harness.strip().lower().replace("-", "_")
+        for member in Harness:
+            if normalized in (member.value, member.name.lower()):
+                hint = f" Did you mean Harness.{member.name}?"
+    raise TypeError(
+        f"harness must be a litellm.harness.Harness member, got {type(harness).__name__} {harness!r}.{hint}"
+    )
+
+
+@dataclass(frozen=True)
+class Usage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    calls: int = 0
+
+    @property
+    def total_tokens(self) -> int:
+        return self.input_tokens + self.output_tokens
+
+
+@dataclass(frozen=True)
+class Text:
+    delta: str
+
+
+@dataclass(frozen=True)
+class Reasoning:
+    delta: str
+
+
+@dataclass(frozen=True)
+class ToolCall:
+    id: str
+    name: str
+    native_name: str
+    input: Mapping[str, Any]
+    builtin: bool = True
+
+
+@dataclass(frozen=True)
+class ToolResult:
+    id: str
+    output: str
+    is_error: bool = False
+
+
+@dataclass(frozen=True)
+class FileChange:
+    path: str
+    kind: FileChangeKind
+    diff: str | None = None
+
+
+@dataclass(frozen=True)
+class Compaction:
+    tokens_before: int | None = None
+    tokens_after: int | None = None
+
+
+@dataclass(frozen=True)
+class Approval:
+    """A request to run a tool. The turn waits until allow() or deny() is called."""
+
+    tool: str
+    input: Mapping[str, Any]
+    _decision: asyncio.Future[tuple[bool, str]] = field(
+        default_factory=lambda: asyncio.get_event_loop().create_future(),
+        compare=False,
+        repr=False,
+    )
+
+    def allow(self) -> None:
+        self._resolve(True, "")
+
+    def deny(self, reason: str = "") -> None:
+        self._resolve(False, reason)
+
+    @property
+    def answered(self) -> bool:
+        return self._decision.done()
+
+    async def wait(self) -> tuple[bool, str]:
+        return await self._decision
+
+    def _resolve(self, allowed: bool, reason: str) -> None:
+        if self._decision.done():
+            return
+        loop = self._decision.get_loop()
+        loop.call_soon_threadsafe(self._set_result, allowed, reason)
+
+    def _set_result(self, allowed: bool, reason: str) -> None:
+        if not self._decision.done():
+            self._decision.set_result((allowed, reason))
+
+
+@dataclass(frozen=True)
+class Result:
+    text: str
+    output: BaseModel | None
+    files: list[FileChange]  # mutable-ok: public Result field; users index/iterate it as a list
+    events: list[Event]  # mutable-ok: public Result field; users index/iterate it as a list
+    usage: Usage
+    cost: float
+    stop_reason: StopReason
+    session_id: str
+
+
+@dataclass(frozen=True)
+class Done:
+    result: Result
+
+    @property
+    def usage(self) -> Usage:
+        return self.result.usage
+
+    @property
+    def cost(self) -> float:
+        return self.result.cost
+
+    @property
+    def stop_reason(self) -> StopReason:
+        return self.result.stop_reason
+
+
+Event = Text | Reasoning | ToolCall | ToolResult | FileChange | Compaction | Approval | Done
+
+
+@dataclass(frozen=True)
+class Capabilities:
+    structured_output: bool
+    tool_approval: bool
+    tool_filtering: bool
+    history: bool
+    custom_tools: bool
+    skills: bool
+    resume: bool
+    permission_modes: frozenset[str]
+
+
+@dataclass(frozen=True)
+class State:
+    """Resume state for a detached or stopped session. Contains no credentials."""
+
+    harness: Harness
+    native_session_id: str | None
+    workdir: str
+    model: str | None = None
+    version: int = STATE_VERSION
+
+    def dumps(self) -> bytes:
+        return json.dumps(
+            {  # mutable-ok: JSON payload serialized immediately by json.dumps
+                "harness": self.harness.value,
+                "native_session_id": self.native_session_id,
+                "workdir": self.workdir,
+                "model": self.model,
+                "version": self.version,
+            }
+        ).encode("utf-8")
+
+    @classmethod
+    def loads(cls, data: bytes) -> State:
+        try:
+            raw = json.loads(data.decode("utf-8"))
+            harness = Harness(raw["harness"])
+            version = int(raw["version"])
+        except (ValueError, KeyError, TypeError, UnicodeDecodeError) as e:
+            raise StateIncompatible(f"Unreadable harness state: {e}") from e
+        if version != STATE_VERSION:
+            raise StateIncompatible(f"State version {version} is not supported (expected {STATE_VERSION})")
+        return cls(
+            harness=harness,
+            native_session_id=raw.get("native_session_id"),
+            workdir=raw["workdir"],
+            model=raw.get("model"),
+            version=version,
+        )

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -3826,7 +3826,7 @@ def _build_bedrock_tool_result_content_blocks(
         if tool_result_content_blocks:
             return tool_result_content_blocks, True
 
-    message_content: Final = message["content"]
+    message_content: Final = message.get("content")
     if isinstance(message_content, str):
         return [BedrockToolResultContentBlock(text=message_content)], False
     if isinstance(message_content, list):
@@ -4095,23 +4095,20 @@ def get_user_message_block_or_continue_message(
 ) -> ChatCompletionUserMessage:
     """
     Returns the user content block
-    if content block is an empty string, then return the default continue message
+    if content block is missing or an empty string, then return the default continue message
 
     Relevant Issue: https://github.com/BerriAI/litellm/issues/7169
     """
     content_block: Final = message.get("content", None)
 
-    # Handle None case
-    if content_block is None or (user_continue_message is None and litellm.modify_params is False):
+    if user_continue_message is None and litellm.modify_params is False:
         return skip_empty_text_blocks(message=message)
 
-    # Handle string case
+    if content_block is None or (isinstance(content_block, str) and not content_block.strip()):
+        return ChatCompletionUserMessage(**(user_continue_message or DEFAULT_USER_CONTINUE_MESSAGE))
+
     if isinstance(content_block, str):
-        # check if content is empty
-        if content_block.strip():
-            return message
-        else:
-            return ChatCompletionUserMessage(**(user_continue_message or DEFAULT_USER_CONTINUE_MESSAGE))
+        return message
 
     # Handle list case
     if isinstance(content_block, list):
@@ -4374,9 +4371,10 @@ class BedrockConverseMessagesProcessor:
                     message=messages[msg_i],
                     user_continue_message=user_continue_message,
                 )
-                if isinstance(message_block["content"], list):
+                message_content = message_block.get("content")
+                if isinstance(message_content, list):
                     _parts: list[BedrockContentBlock] = []
-                    for element in message_block["content"]:
+                    for element in message_content:
                         if isinstance(element, dict):
                             if element["type"] == "text":
                                 _part = BedrockContentBlock(text=element["text"])
@@ -4418,8 +4416,8 @@ class BedrockConverseMessagesProcessor:
                             if _cache_point_block is not None:
                                 _parts.append(_cache_point_block)
                     user_content.extend(_parts)
-                elif message_block["content"] and isinstance(message_block["content"], str):
-                    _part = BedrockContentBlock(text=messages[msg_i]["content"])
+                elif message_content and isinstance(message_content, str):
+                    _part = BedrockContentBlock(text=message_content)
                     _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block(
                         message_block, block_type="content_block", model=model
                     )
@@ -4746,9 +4744,10 @@ def _bedrock_converse_messages_pt(
                 message=messages[msg_i],
                 user_continue_message=user_continue_message,
             )
-            if isinstance(message_block["content"], list):
+            message_content = message_block.get("content")
+            if isinstance(message_content, list):
                 _parts: list[BedrockContentBlock] = []
-                for element in message_block["content"]:
+                for element in message_content:
                     if isinstance(element, dict):
                         if element["type"] == "text":
                             _part = BedrockContentBlock(text=element["text"])
@@ -4791,8 +4790,8 @@ def _bedrock_converse_messages_pt(
                         if _cache_point_block is not None:
                             _parts.append(_cache_point_block)
                 user_content.extend(_parts)
-            elif message_block["content"] and isinstance(message_block["content"], str):
-                _part = BedrockContentBlock(text=messages[msg_i]["content"])
+            elif message_content and isinstance(message_content, str):
+                _part = BedrockContentBlock(text=message_content)
                 _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block(
                     message_block, block_type="content_block", model=model
                 )

litellm/llms/base_llm/harness/transformation.py

@@ -0,0 +1,151 @@
+"""
+Base agent-harness transformation configuration.
+
+A harness is a complete agent runtime (Claude Code, Codex, OpenCode, Deep Agents).
+Like the LLM provider configs in `litellm/llms/base_llm/chat/transformation.py`, a
+harness config only translates: LiteLLM's session parameters in, the runtime's native
+command / config / event stream out. It never does I/O. A handler in
+`litellm/harness/handlers/` owns the sandbox, the process and the per-session model
+endpoint, and calls these transforms.
+
+Adding a CLI harness is one subclass of `BaseCLIHarnessConfig` in
+`litellm/llms/<harness>/harness/transformation.py`, plus one line in
+`ProviderConfigManager.get_provider_harness_config`.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
+
+from litellm.harness.errors import HarnessError, OptionsMismatch
+from litellm.harness.types import Capabilities, Event, Harness
+
+if TYPE_CHECKING:
+    from litellm.harness.context import SessionContext
+
+# A config's typed options (ClaudeCodeOptions, CodexOptions, ...) and its per-turn parser state.
+OptionsT = TypeVar("OptionsT")
+StreamStateT = TypeVar("StreamStateT")
+
+
+def event_list(*events: Event) -> Sequence[Event]:
+    """A transform_stream_line result. One place builds it so every parser returns the same shape."""
+    return list(events)  # mutable-ok: stream-line results are list-shaped; callers and tests compare with list literals
+
+
+class HarnessTurnError(HarnessError):
+    """The runtime reported a failed turn. The runtime maps this to stop_reason='runtime_error'."""
+
+
+@dataclass(frozen=True)
+class HarnessSessionSetup:
+    """What the handler must prepare in the sandbox before the first turn.
+
+    Paths are relative to `private_dir` (a per-session temp dir inside the sandbox)
+    unless they are absolute.
+    """
+
+    files: Mapping[str, bytes] = field(default_factory=dict)
+    # (dir inside private_dir, cache subpath under ~/.cache/litellm-harness) linked so a
+    # later session can resume the runtime's own conversation.
+    persisted_dirs: Sequence[tuple[str, str]] = ()
+    # Where skill folders are copied, relative to private_dir, or absolute.
+    skills_dir: str | None = None
+    # Env passed on every turn. Values may contain `{private_dir}`.
+    env: Mapping[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class HarnessTurnRequest:
+    """One turn of a CLI runtime: the process to run and what to send on stdin."""
+
+    argv: Sequence[str]
+    env: Mapping[str, str]
+    stdin: str
+    cwd: str | None = None
+
+
+@dataclass(frozen=True)
+class HarnessTurnResponse:
+    """What the runtime produced for one turn, after the process exited."""
+
+    final_text: str
+    output_json: str | None = None
+
+
+class BaseHarnessConfig(ABC, Generic[OptionsT]):
+    """Declares what a harness is and validates a session before anything starts."""
+
+    harness: ClassVar[Harness]
+    options_type: type[OptionsT]
+    capabilities: ClassVar[Capabilities]
+    # CLI runtimes call a per-session model endpoint; in-process ones call LiteLLM directly.
+    uses_model_endpoint: ClassVar[bool] = True
+
+    def get_options(self, ctx: SessionContext) -> OptionsT:
+        """ctx.options, or this harness's default options."""
+        options = ctx.options
+        if options is None:
+            return self.options_type()
+        if not isinstance(options, self.options_type):
+            raise OptionsMismatch(
+                f"{type(options).__name__} cannot be used with Harness.{self.harness.name}; "
+                f"use {self.options_type.__name__}"
+            )
+        return options
+
+    def validate_environment(self, ctx: SessionContext) -> None:
+        """Static checks on the session. Raise OptionsMismatch / ValueError early."""
+        self.get_options(ctx)
+
+
+class BaseCLIHarnessConfig(BaseHarnessConfig[OptionsT], Generic[OptionsT, StreamStateT]):
+    """A runtime driven as a subprocess that prints one JSON event per line."""
+
+    @abstractmethod
+    def get_binary(self) -> str:
+        """Executable that must be on the sandbox's PATH."""
+
+    @abstractmethod
+    def get_install_hint(self) -> str:
+        """How to install the binary; shown in HarnessInstallFailed."""
+
+    @abstractmethod
+    def transform_session_setup(self, ctx: SessionContext, private_dir: str) -> HarnessSessionSetup:
+        """Config files, env and persisted dirs for the session."""
+
+    @abstractmethod
+    def transform_turn_request(
+        self,
+        ctx: SessionContext,
+        setup: HarnessSessionSetup,
+        private_dir: str,
+        prompt: str,
+        native_session_id: str | None,
+    ) -> HarnessTurnRequest:
+        """argv / env / stdin for one turn. native_session_id is set after the first turn."""
+
+    @abstractmethod
+    def create_stream_state(self) -> StreamStateT:
+        """Fresh per-turn parser state."""
+
+    @abstractmethod
+    def transform_stream_line(self, line: Mapping[str, Any], state: StreamStateT) -> Sequence[Event]:
+        """One decoded JSON line from stdout to zero or more events. Pure."""
+
+    @abstractmethod
+    def get_native_session_id(self, state: StreamStateT) -> str | None:
+        """The runtime's own session / thread id, once the stream has reported it."""
+
+    @abstractmethod
+    def transform_turn_response(
+        self,
+        ctx: SessionContext,
+        state: StreamStateT,
+        exit_code: int,
+        stderr_tail: Sequence[str],
+    ) -> HarnessTurnResponse:
+        """Final text and structured output, or raise HarnessTurnError."""

litellm/llms/base_llm/harness/utils.py

@@ -0,0 +1,116 @@
+"""Pure helpers shared by harness configs."""
+
+from __future__ import annotations
+
+import itertools
+import json
+import os
+from collections.abc import Iterator, Mapping, Sequence
+from types import MappingProxyType
+from typing import Any, Final, TypeAlias
+
+from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+
+# A decoded JSON document: what json.loads / model_json_schema() produce.
+JSONValue: TypeAlias = "dict[str, JSONValue] | list[JSONValue] | str | int | float | bool | None"
+
+SKILL_MANIFEST: Final = "SKILL.md"
+_JSON_DECODER: Final = json.JSONDecoder()
+
+
+def normalize_tool_name(native_name: str, mapping: Mapping[str, str]) -> str:
+    """Normalized tool name (read, write, edit, bash, ...) or the native name if unmapped."""
+    return mapping.get(native_name, native_name)
+
+
+def native_tool_names(normalized: Sequence[str], mapping: Mapping[str, Sequence[str]]) -> Sequence[str]:
+    """Native names for normalized tool names, de-duplicated, order kept."""
+    expanded: Final = itertools.chain.from_iterable(mapping.get(name, (name,)) for name in normalized)
+    return list(dict.fromkeys(expanded))  # mutable-ok: public helper whose callers/tests compare against list literals
+
+
+def last_json_object(text: str) -> str | None:
+    """The last top-level `{...}` in text that parses as a JSON object, re-serialized."""
+    last: str | None = None
+    index = text.find("{")
+    while index != -1:
+        try:
+            obj, end = _JSON_DECODER.raw_decode(text, index)
+        except json.JSONDecodeError:
+            index = text.find("{", index + 1)
+            continue
+        if isinstance(obj, dict):
+            last = json.dumps(obj)
+        index = text.find("{", end)
+    return last
+
+
+def structured_output_instruction(schema: Mapping[str, Any]) -> str:
+    return (
+        "When you have finished, your final message must be a single JSON object that "
+        "matches this JSON schema, with no other text before or after it:\n"
+        f"{json.dumps(schema)}"
+    )
+
+
+def strict_json_schema(schema: JSONValue, depth: int = 0) -> JSONValue:
+    """Make a JSON schema acceptable to OpenAI strict structured outputs.
+
+    Every object gets `additionalProperties: false` and all of its properties required,
+    recursively. Keywords strict mode rejects next to $ref are dropped. Nesting deeper than
+    DEFAULT_MAX_RECURSE_DEPTH raises instead of recursing further.
+    """
+    if depth > DEFAULT_MAX_RECURSE_DEPTH:
+        raise ValueError(f"output schema is nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels")
+    if isinstance(schema, list):
+        return [strict_json_schema(entry, depth + 1) for entry in schema]  # mutable-ok: JSON document output
+    if not isinstance(schema, dict):
+        return schema
+    entries: Final = ((key, strict_json_schema(value, depth + 1)) for key, value in schema.items())
+    result = dict(entries)  # mutable-ok: JSON document; "default" is popped below
+    if "$ref" in result:
+        return {"$ref": result["$ref"]}  # mutable-ok: JSONValue output is a plain JSON document
+    result.pop("default", None)
+    properties = result.get("properties")
+    if result.get("type") == "object" or isinstance(properties, dict):
+        props = properties if isinstance(properties, dict) else {}  # mutable-ok: JSONValue object member
+        required: Final[list[JSONValue]] = list(props)  # mutable-ok: JSON array in the output schema
+        strict: Final[Mapping[str, JSONValue]] = MappingProxyType(
+            {"properties": props, "required": required, "additionalProperties": False}
+        )
+        result = {**result, **strict}  # mutable-ok: JSON document output
+    return result
+
+
+def decode_json_line(line: bytes | str) -> Mapping[str, Any] | None:
+    """One JSONL line as a dict, or None for blank / non-JSON / non-object lines."""
+    text = line.strip()
+    if not text:
+        return None
+    try:
+        obj = json.loads(text)
+    except json.JSONDecodeError:
+        return None
+    return obj if isinstance(obj, dict) else None
+
+
+def stderr_tail_text(stderr_tail: Sequence[str]) -> str:
+    return "\n".join(line for line in stderr_tail if line.strip())
+
+
+def _read_bytes(path: str) -> bytes:
+    with open(path, "rb") as fh:
+        return fh.read()
+
+
+def _walk_files(root: str) -> Iterator[str]:
+    for dirpath, _dirnames, filenames in os.walk(root):
+        yield from (os.path.join(dirpath, filename) for filename in sorted(filenames))
+
+
+def read_skill_files(skill_dir: str) -> tuple[tuple[str, bytes], ...]:
+    """(relative path, bytes) for every file under a local skill folder."""
+    root: Final = os.path.realpath(os.fspath(skill_dir))
+    if not os.path.isfile(os.path.join(root, SKILL_MANIFEST)):
+        raise ValueError(f"skill folder {skill_dir!r} has no {SKILL_MANIFEST}")
+    return tuple((os.path.relpath(path, root), _read_bytes(path)) for path in _walk_files(root))

litellm/llms/bedrock/chat/converse_transformation.py

@@ -1448,7 +1448,7 @@ class AmazonConverseConfig(BaseConfig):
         )
 
     def _converted_text_blocks(self, message: ChatCompletionSystemMessage) -> tuple[ChatCompletionTextObject, ...]:
-        content: Final = message["content"]
+        content: Final = message.get("content")
         if isinstance(content, str):
             return (self._converted_text_block(content, message.get("cache_control")),) if content else ()
         parts: Final[Sequence[object]] = content or ()
@@ -1483,13 +1483,14 @@ class AmazonConverseConfig(BaseConfig):
         for message in hoisted:
             if message["role"] != "system":
                 continue
-            if isinstance(message["content"], str) and message["content"]:
-                system_content_blocks.append(SystemContentBlock(text=message["content"]))
+            content = message.get("content")
+            if isinstance(content, str) and content:
+                system_content_blocks.append(SystemContentBlock(text=content))
                 cache_block = self.get_cache_point_block(message, block_type="system", model=model)
                 if cache_block:
                     system_content_blocks.append(cache_block)
-            elif isinstance(message["content"], list):
-                for m in message["content"]:
+            elif isinstance(content, list):
+                for m in content:
                     if m.get("type") == "text" and m.get("text"):
                         system_content_blocks.append(SystemContentBlock(text=m["text"]))
                         cache_block = self.get_cache_point_block(m, block_type="system", model=model)

litellm/llms/claude_code/harness/transformation.py

@@ -0,0 +1,389 @@
+"""
+Claude Code harness config: `claude -p --output-format stream-json`, once per turn.
+
+Every model call goes to the per-session endpoint with the per-session token. The CLI gets
+a private CLAUDE_CONFIG_DIR and only the `user` setting source (that private dir), so
+neither the user's login, keychain, nor a repo's `.claude/settings.json` can swap the base
+URL or credentials. Verified against Claude Code 2.1.285.
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Final
+
+from litellm.harness.errors import HarnessError, OptionsMismatch
+from litellm.harness.options import ClaudeCodeOptions
+from litellm.harness.types import (
+    Capabilities,
+    Compaction,
+    Event,
+    Harness,
+    Reasoning,
+    Text,
+    ToolCall,
+    ToolResult,
+)
+from litellm.llms.base_llm.harness.transformation import (
+    BaseCLIHarnessConfig,
+    HarnessSessionSetup,
+    HarnessTurnError,
+    HarnessTurnRequest,
+    HarnessTurnResponse,
+    event_list,
+)
+from litellm.llms.base_llm.harness.utils import (
+    last_json_object,
+    native_tool_names,
+    normalize_tool_name,
+    stderr_tail_text,
+)
+
+if TYPE_CHECKING:
+    from litellm.harness.context import SessionContext
+
+CLAUDE_BINARY: Final = "claude"
+SYNTHETIC_MODEL: Final = "<synthetic>"
+
+BASE_COMMAND: Final = ("-p", "--output-format", "stream-json", "--verbose", "--input-format", "text")
+
+PERMISSION_MODES: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "read-only": "plan",
+        "ask": "default",
+        "edit": "acceptEdits",
+        "full": "bypassPermissions",
+    }
+)
+
+NATIVE_TO_NORMALIZED: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "Read": "read",
+        "Write": "write",
+        "Edit": "edit",
+        "MultiEdit": "edit",
+        "Bash": "bash",
+        "Glob": "glob",
+        "Grep": "grep",
+        "WebSearch": "web_search",
+        "LS": "ls",
+    }
+)
+
+NORMALIZED_TO_NATIVE: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
+    {
+        "read": ("Read",),
+        "write": ("Write",),
+        "edit": ("Edit", "MultiEdit"),
+        "bash": ("Bash",),
+        "glob": ("Glob",),
+        "grep": ("Grep",),
+        "web_search": ("WebSearch",),
+        "ls": ("LS",),
+    }
+)
+
+# Env the config owns; ClaudeCodeOptions.env may not override these.
+MANAGED_ENV_KEYS: Final = frozenset(
+    {
+        "ANTHROPIC_BASE_URL",
+        "ANTHROPIC_AUTH_TOKEN",
+        "ANTHROPIC_API_KEY",
+        "ANTHROPIC_MODEL",
+        "ANTHROPIC_SMALL_FAST_MODEL",
+        "CLAUDE_CONFIG_DIR",
+    }
+)
+
+# Claude Code settings.json keys LiteLLM manages (or that could reroute model calls or credentials).
+MANAGED_CONFIG_KEYS: Final = frozenset(
+    {"env", "apiKeyHelper", "model", "permissions", "awsAuthRefresh", "awsCredentialExport", "forceLoginMethod"}
+)
+
+STATIC_ENV: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "DISABLE_TELEMETRY": "1",
+        "DISABLE_ERROR_REPORTING": "1",
+        "DISABLE_AUTOUPDATER": "1",
+        "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
+    }
+)
+
+STRUCTURED_OUTPUT_INSTRUCTION: Final = (
+    "When you have finished the task, end your final reply with only a single JSON "
+    "object (no code fences, no prose after it) that matches this JSON schema:\n{schema}"
+)
+
+
+@dataclass
+class ClaudeCodeStreamState:
+    """What the parser has learned from one turn's stream-json output."""
+
+    session_id: str | None = None
+    text_parts: list[str] = field(default_factory=list)  # mutable-ok: parser appends text deltas
+    result_seen: bool = False
+    result_text: str | None = None
+    is_error: bool = False
+    errors: Sequence[str] = ()
+    structured_output: Any | None = None
+
+    @property
+    def final_text(self) -> str:
+        if self.result_text is not None:
+            return self.result_text
+        return "".join(self.text_parts)
+
+
+def stringify_tool_output(content: object) -> str:
+    """tool_result content is a string or a list of content blocks."""
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        return "\n".join(_stringify_block(block) for block in content)
+    return json.dumps(content, ensure_ascii=False)
+
+
+def _stringify_block(block: object) -> str:
+    if isinstance(block, dict) and block.get("type") == "text":
+        return str(block.get("text", ""))
+    if isinstance(block, str):
+        return block
+    return json.dumps(block, ensure_ascii=False)
+
+
+def _message_blocks(event: Mapping[str, Any]) -> Sequence[Any]:
+    message: Final = event.get("message")
+    content: Final = message.get("content") if isinstance(message, Mapping) else None
+    if isinstance(content, str):
+        return ({"type": "text", "text": content},)  # mutable-ok: JSON content block, like the stream's
+    return content if isinstance(content, list) else ()
+
+
+def _assistant_block_events(block: Mapping[str, Any], state: ClaudeCodeStreamState) -> tuple[Event, ...]:
+    kind = block.get("type")
+    if kind == "text" and block.get("text"):
+        state.text_parts.append(block["text"])
+        return (Text(delta=block["text"]),)
+    if kind == "thinking" and block.get("thinking"):
+        return (Reasoning(delta=block["thinking"]),)
+    if kind == "tool_use":
+        native = str(block.get("name", ""))
+        return (
+            ToolCall(
+                id=str(block.get("id", "")),
+                name=normalize_tool_name(native, NATIVE_TO_NORMALIZED),
+                native_name=native,
+                input=block.get("input")
+                or {},  # mutable-ok: ToolCall.input is a dict field; empty default for a missing input
+                builtin=not native.startswith("mcp__"),
+            ),
+        )
+    return ()
+
+
+def _assistant_events(event: Mapping[str, Any], state: ClaudeCodeStreamState) -> Sequence[Event]:
+    if event.get("parent_tool_use_id"):
+        return event_list()  # subagent traffic
+    message: Final = event.get("message")
+    if isinstance(message, Mapping) and message.get("model") == SYNTHETIC_MODEL:
+        return event_list()  # CLI-generated error text; surfaced via the result event
+    blocks: Final = (block for block in _message_blocks(event) if isinstance(block, dict))
+    return event_list(*itertools.chain.from_iterable(_assistant_block_events(block, state) for block in blocks))
+
+
+def _is_tool_result(block: object) -> bool:
+    return isinstance(block, dict) and block.get("type") == "tool_result"
+
+
+def _user_events(event: Mapping[str, Any]) -> Sequence[Event]:
+    if event.get("parent_tool_use_id"):
+        return event_list()
+    return event_list(
+        *(
+            ToolResult(
+                id=str(block.get("tool_use_id", "")),
+                output=stringify_tool_output(block.get("content")),
+                is_error=bool(block.get("is_error", False)),
+            )
+            for block in _message_blocks(event)
+            if _is_tool_result(block)
+        )
+    )
+
+
+def _system_events(event: Mapping[str, Any], state: ClaudeCodeStreamState) -> Sequence[Event]:
+    subtype = event.get("subtype")
+    if subtype == "init" and event.get("session_id"):
+        state.session_id = str(event["session_id"])
+        return event_list()
+    if subtype == "compact_boundary":
+        meta: Final = event.get("compact_metadata")
+        pre_tokens: Final = meta.get("pre_tokens") if isinstance(meta, Mapping) else None
+        return event_list(Compaction(tokens_before=pre_tokens, tokens_after=None))
+    return event_list()
+
+
+def _record_result(event: Mapping[str, Any], state: ClaudeCodeStreamState) -> Sequence[Event]:
+    state.result_seen = True
+    state.is_error = bool(event.get("is_error", False))
+    result = event.get("result")
+    state.result_text = result if isinstance(result, str) else None
+    state.errors = [str(e) for e in event.get("errors") or ()]  # mutable-ok: mirrors the JSON errors array
+    state.structured_output = event.get("structured_output")
+    if event.get("session_id"):
+        state.session_id = str(event["session_id"])
+    return event_list()
+
+
+def turn_error_message(state: ClaudeCodeStreamState, exit_code: int, stderr_tail: Sequence[str]) -> str | None:
+    """None if the turn succeeded, else the message for HarnessTurnError."""
+    if exit_code == 0 and state.result_seen and not state.is_error:
+        return None
+    reason = state.result_text or "; ".join(state.errors)
+    if not reason:
+        reason = "no result event" if not state.result_seen else "unknown error"
+    message = f"claude exited with code {exit_code}: {reason}"
+    tail = stderr_tail_text(stderr_tail)
+    return f"{message}\nstderr:\n{tail}" if tail else message
+
+
+def build_system_prompt(instructions: str | None, output_schema: Mapping[str, Any] | None) -> str | None:
+    schema_part: Final = (
+        STRUCTURED_OUTPUT_INSTRUCTION.format(schema=json.dumps(output_schema)) if output_schema is not None else None
+    )
+    parts: Final = tuple(part for part in (instructions, schema_part) if part)
+    return "\n\n".join(parts) if parts else None
+
+
+class ClaudeCodeHarnessConfig(BaseCLIHarnessConfig):
+    harness = Harness.CLAUDE_CODE
+    options_type = ClaudeCodeOptions
+    capabilities = Capabilities(
+        structured_output=True,
+        tool_approval=False,
+        tool_filtering=True,
+        history=False,
+        custom_tools=False,
+        skills=True,
+        resume=True,
+        permission_modes=frozenset({"read-only", "edit", "full"}),
+    )
+
+    def get_binary(self) -> str:
+        return CLAUDE_BINARY
+
+    def get_install_hint(self) -> str:
+        return "npm install -g @anthropic-ai/claude-code"
+
+    def validate_environment(self, ctx: SessionContext) -> None:
+        options: ClaudeCodeOptions = self.get_options(ctx)
+        clashing = sorted(MANAGED_ENV_KEYS.intersection(options.env))
+        if clashing:
+            raise OptionsMismatch(f"ClaudeCodeOptions.env may not set {', '.join(clashing)}; LiteLLM manages it")
+        managed = sorted(MANAGED_CONFIG_KEYS.intersection(options.config))
+        if managed:
+            raise OptionsMismatch(
+                f"ClaudeCodeOptions.config may not set {', '.join(managed)}; "
+                "use the matching agent() argument (model=, permissions=) instead"
+            )
+
+    def transform_session_setup(self, ctx: SessionContext, private_dir: str) -> HarnessSessionSetup:
+        if ctx.endpoint is None or not ctx.endpoint.token:
+            raise HarnessError("Claude Code needs the session model endpoint")
+        options: ClaudeCodeOptions = self.get_options(ctx)
+        model = ctx.model
+        # Background calls (titles, summaries) use the same model group, like OpenCode.
+        model_env: Final = (
+            MappingProxyType({"ANTHROPIC_MODEL": model, "ANTHROPIC_SMALL_FAST_MODEL": model})
+            if model
+            else MappingProxyType({})
+        )
+        env: Final = MappingProxyType(
+            {
+                **options.env,
+                **STATIC_ENV,
+                "ANTHROPIC_BASE_URL": ctx.sandbox.host_url(ctx.endpoint.port),
+                "ANTHROPIC_AUTH_TOKEN": ctx.endpoint.token,
+                "ANTHROPIC_API_KEY": "",
+                "CLAUDE_CONFIG_DIR": private_dir,
+                **model_env,
+            }
+        )
+        return HarnessSessionSetup(
+            persisted_dirs=[("projects", "claude_code/projects")],  # mutable-ok: tests compare to a list
+            skills_dir="skills",
+            env=env,
+        )
+
+    def transform_turn_request(
+        self,
+        ctx: SessionContext,
+        setup: HarnessSessionSetup,
+        private_dir: str,
+        prompt: str,
+        native_session_id: str | None,
+    ) -> HarnessTurnRequest:
+        options: ClaudeCodeOptions = self.get_options(ctx)
+        schema = ctx.output.model_json_schema() if ctx.output is not None else None
+        system_prompt: Final = build_system_prompt(ctx.instructions, schema)
+        disallowed: Final = native_tool_names(ctx.disable_tools, NORMALIZED_TO_NATIVE)
+        config: Final = dict(options.config)  # mutable-ok: json.dumps needs a plain dict
+        settings: Final = json.dumps(config) if config else None
+        argv: Final = (
+            CLAUDE_BINARY,
+            *BASE_COMMAND,
+            "--permission-mode",
+            PERMISSION_MODES[ctx.permissions],
+            *(("--model", ctx.model) if ctx.model else ()),
+            # Only read settings from the private CLAUDE_CONFIG_DIR, never the repo's .claude/.
+            "--setting-sources",
+            "user",
+            *(("--settings", settings) if settings else ()),
+            *(("--append-system-prompt", system_prompt) if system_prompt else ()),
+            *(("--max-turns", str(ctx.max_turns)) if ctx.max_turns is not None else ()),
+            *(("--disallowedTools", ",".join(disallowed)) if disallowed else ()),
+            *(("--resume", native_session_id) if native_session_id else ()),
+        )
+        return HarnessTurnRequest(argv=argv, env=setup.env, stdin=prompt)
+
+    def create_stream_state(self) -> ClaudeCodeStreamState:
+        return ClaudeCodeStreamState()
+
+    def transform_stream_line(self, line: Mapping[str, Any], state: ClaudeCodeStreamState) -> Sequence[Event]:
+        kind = line.get("type")
+        if kind == "assistant":
+            return _assistant_events(line, state)
+        if kind == "user":
+            return _user_events(line)
+        if kind == "system":
+            return _system_events(line, state)
+        if kind == "result":
+            return _record_result(line, state)
+        return event_list()
+
+    def get_native_session_id(self, state: ClaudeCodeStreamState) -> str | None:
+        return state.session_id
+
+    def transform_turn_response(
+        self,
+        ctx: SessionContext,
+        state: ClaudeCodeStreamState,
+        exit_code: int,
+        stderr_tail: Sequence[str],
+    ) -> HarnessTurnResponse:
+        error = turn_error_message(state, exit_code, stderr_tail)
+        if error is not None:
+            raise HarnessTurnError(error)
+        output_json: str | None = None
+        if ctx.output is not None:
+            if isinstance(state.structured_output, dict):
+                output_json = json.dumps(state.structured_output)
+            else:
+                output_json = last_json_object(state.final_text)
+        return HarnessTurnResponse(final_text=state.final_text, output_json=output_json)

litellm/llms/codex/harness/transformation.py

@@ -0,0 +1,349 @@
+"""
+Codex harness config: `codex exec --json` (JSONL events), once per turn.
+
+Every model call goes to one custom provider (`litellm`, wire_api=responses) pointing at the
+per-session endpoint. The bearer token only travels in the LITELLM_HARNESS_TOKEN env var,
+never in argv. CODEX_HOME is the private session dir so the user's own Codex config and
+auth are never read. Verified against codex-cli 0.135.0.
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+import re
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass, field
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Final
+
+from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+from litellm.harness.errors import HarnessError, OptionsMismatch
+from litellm.harness.options import CodexOptions
+from litellm.harness.types import (
+    Capabilities,
+    Event,
+    Harness,
+    Reasoning,
+    Text,
+    ToolCall,
+    ToolResult,
+)
+from litellm.llms.base_llm.harness.transformation import (
+    BaseCLIHarnessConfig,
+    HarnessSessionSetup,
+    HarnessTurnError,
+    HarnessTurnRequest,
+    HarnessTurnResponse,
+    event_list,
+)
+from litellm.llms.base_llm.harness.utils import stderr_tail_text, strict_json_schema
+
+if TYPE_CHECKING:
+    from litellm.harness.context import SessionContext
+
+CODEX_BINARY: Final = "codex"
+CODEX_PROVIDER_ID: Final = "litellm"
+CODEX_TOKEN_ENV: Final = "LITELLM_HARNESS_TOKEN"
+CODEX_SCHEMA_FILENAME: Final = "output_schema.json"
+# Top-level config keys LiteLLM sets itself; users may not override them via options.config.
+MANAGED_CONFIG_KEYS: Final = frozenset(
+    {
+        "model",
+        "model_provider",
+        "model_providers",
+        "approval_policy",
+        "sandbox_mode",
+        "mcp_servers",
+        "developer_instructions",
+        "web_search",
+    }
+)
+_BARE_TOML_KEY: Final = re.compile(r"^[A-Za-z0-9_-]+$")
+_TOOL_ITEM_TYPES: Final = frozenset({"command_execution", "file_change", "web_search", "mcp_tool_call"})
+
+
+@dataclass
+class CodexStreamState:
+    """What the parser has learned from one turn's JSONL events."""
+
+    thread_id: str | None = None
+    final_text: str = ""
+    error: str | None = None
+    failed: bool = False
+    started: set[str] = field(default_factory=set)  # mutable-ok: parser records announced tool items
+
+
+def _tool_input(item: Mapping[str, Any]) -> tuple[str, str, Mapping[str, Any], bool]:
+    """(normalized name, native name, input, builtin) for a tool-like item."""
+    item_type = item.get("type")
+    if item_type == "command_execution":
+        return "bash", "command_execution", MappingProxyType({"command": item.get("command", "")}), True
+    if item_type == "file_change":
+        changes: Final = list(item.get("changes") or ())  # mutable-ok: JSON array, as codex reports it
+        return "edit", "apply_patch", MappingProxyType({"changes": changes}), True
+    if item_type == "web_search":
+        return "web_search", "web_search", MappingProxyType({"query": item.get("query", "")}), True
+    server = str(item.get("server") or "")
+    tool = str(item.get("tool") or "")
+    arguments = item.get("arguments")
+    tool_args = arguments if isinstance(arguments, dict) else MappingProxyType({"arguments": arguments})
+    name = f"{server}.{tool}" if server else tool
+    return name, tool, tool_args, False
+
+
+def _tool_output(item: Mapping[str, Any]) -> tuple[str, bool]:
+    """(output text, is_error) for a completed tool-like item."""
+    item_type = item.get("type")
+    status = item.get("status")
+    if item_type == "command_execution":
+        exit_code = item.get("exit_code")
+        is_error = status == "failed" or (exit_code is not None and exit_code != 0)
+        return str(item.get("aggregated_output") or ""), is_error
+    if item_type == "file_change":
+        lines = (f"{c.get('kind', '')} {c.get('path', '')}".strip() for c in item.get("changes") or ())
+        return "\n".join(lines), status == "failed"
+    if item_type == "web_search":
+        return "", status == "failed"
+    error = item.get("error")
+    if error:
+        message = error.get("message") if isinstance(error, dict) else error
+        return str(message), True
+    result = item.get("result")
+    if result is None:
+        return "", status == "failed"
+    if isinstance(result, str):
+        return result, status == "failed"
+    return json.dumps(result), status == "failed"
+
+
+def _tool_item_events(
+    item_id: str, item: Mapping[str, Any], completed: bool, state: CodexStreamState
+) -> Iterator[Event]:
+    if item_id not in state.started:
+        state.started.add(item_id)
+        name, native_name, tool_input, builtin = _tool_input(item)
+        yield ToolCall(id=item_id, name=name, native_name=native_name, input=tool_input, builtin=builtin)
+    if completed:
+        output, is_error = _tool_output(item)
+        yield ToolResult(id=item_id, output=output, is_error=is_error)
+
+
+def _item_events(event_type: str, item: Mapping[str, Any], state: CodexStreamState) -> Sequence[Event]:
+    item_type = item.get("type")
+    item_id = str(item.get("id") or "")
+    completed = event_type == "item.completed"
+    if item_type == "agent_message":
+        if not completed:
+            return event_list()
+        text = str(item.get("text") or "")
+        state.final_text = text
+        return event_list(Text(delta=text)) if text else event_list()
+    if item_type == "reasoning":
+        text = str(item.get("text") or "")
+        return event_list(Reasoning(delta=text)) if completed and text else event_list()
+    if item_type not in _TOOL_ITEM_TYPES:
+        return event_list()
+    return event_list(*_tool_item_events(item_id, item, completed, state))
+
+
+def toml_value(value: object, depth: int = 0) -> str:
+    """Encode a Python value as a TOML value for `codex -c key=value`."""
+    if depth > DEFAULT_MAX_RECURSE_DEPTH:
+        raise OptionsMismatch(f"CodexOptions.config is nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels")
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if isinstance(value, (int, float)):
+        return repr(value)
+    if isinstance(value, str):
+        return json.dumps(value)
+    if isinstance(value, Mapping):
+        pairs = ", ".join(f"{toml_key(k)} = {toml_value(v, depth + 1)}" for k, v in value.items())
+        return "{" + pairs + "}"
+    if isinstance(value, (list, tuple)):
+        return "[" + ", ".join(toml_value(v, depth + 1) for v in value) + "]"
+    raise OptionsMismatch(f"CodexOptions.config value of type {type(value).__name__} cannot be passed to codex")
+
+
+def toml_key(key: object) -> str:
+    text = str(key)
+    return text if _BARE_TOML_KEY.match(text) else json.dumps(text)
+
+
+def _config_override(key: object, value: object) -> str:
+    dotted = str(key)
+    if not dotted or "=" in dotted:
+        raise OptionsMismatch(f"Invalid CodexOptions.config key: {dotted!r}")
+    if dotted.split(".", 1)[0] in MANAGED_CONFIG_KEYS:
+        raise OptionsMismatch(
+            f"CodexOptions.config[{dotted!r}] is managed by LiteLLM; use the matching agent() argument instead"
+        )
+    return f"{dotted}={toml_value(value)}"
+
+
+def config_overrides(config: Mapping[str, Any]) -> Sequence[str]:
+    """`-c` override strings for CodexOptions.config, rejecting managed keys."""
+    overrides: Final = (_config_override(key, value) for key, value in config.items())
+    return list(overrides)  # mutable-ok: public helper; tests compare to a list
+
+
+def _flag_pairs(flag: str, values: Sequence[str]) -> tuple[str, ...]:
+    return tuple(itertools.chain.from_iterable((flag, value) for value in values))
+
+
+class CodexHarnessConfig(BaseCLIHarnessConfig):
+    harness = Harness.CODEX
+    options_type = CodexOptions
+    capabilities = Capabilities(
+        structured_output=True,
+        tool_approval=False,
+        tool_filtering=False,
+        history=False,
+        custom_tools=False,
+        skills=True,
+        resume=True,
+        permission_modes=frozenset({"read-only", "full"}),
+    )
+
+    def get_binary(self) -> str:
+        return CODEX_BINARY
+
+    def get_install_hint(self) -> str:
+        return "npm install -g @openai/codex (or brew install codex)"
+
+    def validate_environment(self, ctx: SessionContext) -> None:
+        options: CodexOptions = self.get_options(ctx)
+        config_overrides(options.config)
+
+    def transform_session_setup(self, ctx: SessionContext, private_dir: str) -> HarnessSessionSetup:
+        if ctx.endpoint is None:
+            raise HarnessError("Codex needs the session model endpoint")
+        options: CodexOptions = self.get_options(ctx)
+        files: Final = (
+            MappingProxyType(
+                {CODEX_SCHEMA_FILENAME: json.dumps(strict_json_schema(ctx.output.model_json_schema())).encode("utf-8")}
+            )
+            if ctx.output is not None
+            else MappingProxyType({})
+        )
+        return HarnessSessionSetup(
+            files=files,
+            persisted_dirs=[("sessions", "codex/sessions")],  # mutable-ok: tests compare to a list
+            skills_dir="skills",
+            env=MappingProxyType({**options.env, CODEX_TOKEN_ENV: ctx.endpoint.token, "CODEX_HOME": private_dir}),
+        )
+
+    def transform_turn_request(
+        self,
+        ctx: SessionContext,
+        setup: HarnessSessionSetup,
+        private_dir: str,
+        prompt: str,
+        native_session_id: str | None,
+    ) -> HarnessTurnRequest:
+        if ctx.endpoint is None:
+            raise HarnessError("Codex needs the session model endpoint")
+        options: CodexOptions = self.get_options(ctx)
+        head: Final = (
+            (CODEX_BINARY, "exec", "resume", native_session_id) if native_session_id else (CODEX_BINARY, "exec")
+        )
+        argv: Final = (
+            *head,
+            "--json",
+            "--skip-git-repo-check",
+            *(("-m", ctx.model) if ctx.model else ()),
+            *_flag_pairs("-c", self._provider_overrides(ctx)),
+            *self._permission_args(ctx, native_session_id),
+            *_flag_pairs("-c", self._feature_overrides(ctx, options)),
+            *_flag_pairs("-c", config_overrides(options.config)),
+            *(
+                ("--output-schema", f"{private_dir}/{CODEX_SCHEMA_FILENAME}")
+                if CODEX_SCHEMA_FILENAME in setup.files
+                else ()
+            ),
+            *(() if native_session_id else ("-C", ctx.sandbox.workdir)),
+            "-",
+        )
+        return HarnessTurnRequest(argv=argv, env=setup.env, stdin=prompt, cwd=ctx.sandbox.workdir)
+
+    def _provider_overrides(self, ctx: SessionContext) -> tuple[str, ...]:
+        assert ctx.endpoint is not None
+        base_url = ctx.sandbox.host_url(ctx.endpoint.port).rstrip("/") + "/v1"
+        prefix = f"model_providers.{CODEX_PROVIDER_ID}"
+        return (
+            f"model_provider={CODEX_PROVIDER_ID}",
+            f"{prefix}.name={CODEX_PROVIDER_ID}",
+            f"{prefix}.base_url={toml_value(base_url)}",
+            f"{prefix}.env_key={CODEX_TOKEN_ENV}",
+            f"{prefix}.wire_api=responses",
+            "approval_policy=never",
+        )
+
+    @staticmethod
+    def _permission_args(ctx: SessionContext, native_session_id: str | None) -> tuple[str, ...]:
+        if ctx.permissions == "read-only":
+            mode = "read-only"
+        elif getattr(ctx.sandbox, "is_container", False):
+            # The container is already the boundary; nested sandboxing fails in containers.
+            return ("--dangerously-bypass-approvals-and-sandbox",)
+        else:
+            mode = "workspace-write"
+        # `codex exec resume` has no --sandbox flag; the config key works for both.
+        if native_session_id:
+            return ("-c", f"sandbox_mode={toml_value(mode)}")
+        return ("--sandbox", mode)
+
+    @staticmethod
+    def _feature_overrides(ctx: SessionContext, options: CodexOptions) -> tuple[str, ...]:
+        reasoning: Final = (
+            (
+                f"model_reasoning_effort={options.reasoning_effort}",
+                "model_reasoning_summary=auto",
+                "model_supports_reasoning_summaries=true",
+            )
+            if options.reasoning_effort
+            else ()
+        )
+        instructions: Final = (f"developer_instructions={toml_value(ctx.instructions)}",) if ctx.instructions else ()
+        return (f"web_search={'live' if options.web_search else 'disabled'}", *reasoning, *instructions)
+
+    def create_stream_state(self) -> CodexStreamState:
+        return CodexStreamState()
+
+    def transform_stream_line(self, line: Mapping[str, Any], state: CodexStreamState) -> Sequence[Event]:
+        """turn.completed usage is ignored on purpose: the session endpoint accounts it."""
+        event_type = line.get("type")
+        if event_type == "thread.started":
+            if line.get("thread_id"):
+                state.thread_id = str(line["thread_id"])
+            return event_list()
+        if event_type in ("item.started", "item.updated", "item.completed"):
+            item = line.get("item")
+            return _item_events(str(event_type), item, state) if isinstance(item, dict) else event_list()
+        if event_type == "error":
+            state.error = str(line.get("message") or "codex reported an error")
+            return event_list()
+        if event_type == "turn.failed":
+            error = line.get("error")
+            message = error.get("message") if isinstance(error, dict) else error
+            state.error = str(message or state.error or "codex turn failed")
+            state.failed = True
+        return event_list()
+
+    def get_native_session_id(self, state: CodexStreamState) -> str | None:
+        return state.thread_id
+
+    def transform_turn_response(
+        self,
+        ctx: SessionContext,
+        state: CodexStreamState,
+        exit_code: int,
+        stderr_tail: Sequence[str],
+    ) -> HarnessTurnResponse:
+        if state.failed:
+            raise HarnessTurnError(f"codex turn failed: {state.error}")
+        if exit_code != 0:
+            detail = stderr_tail_text(stderr_tail) or state.error or "no output"
+            raise HarnessTurnError(f"codex exited with code {exit_code}: {detail}")
+        output_json = state.final_text if ctx.output is not None else None
+        return HarnessTurnResponse(final_text=state.final_text, output_json=output_json)

litellm/llms/deepagents/harness/sandbox_backend.py

@@ -0,0 +1,559 @@
+"""Deep Agents pieces that subclass optional-dependency bases.
+
+Only imported by `litellm.harness.handlers.deepagents_handler.load_deps()`, so `deepagents`,
+`langchain` and `langchain-core` are never imported unless Harness.DEEPAGENTS is used.
+
+`SandboxBackend` implements deepagents' `SandboxBackendProtocol` on top of a litellm
+`Sandbox`. The agent sees virtual paths rooted at the sandbox workdir (`/src/a.py` is
+`<workdir>/src/a.py`); file bytes move through `Sandbox.read/write`, and ls/glob/grep/
+delete/execute run plain POSIX commands through `Sandbox.run`, so the same code serves the
+local and docker sandboxes (no python3 needed inside the sandbox) and inherits the sandbox's
+env scrubbing and path confinement.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import itertools
+import posixpath
+import re
+import shlex
+import uuid
+from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
+from types import MappingProxyType
+from typing import Any, Final, TypeVar
+
+from deepagents.backends.protocol import (
+    DeleteResult,
+    EditResult,
+    ExecuteResponse,
+    FileData,
+    FileDownloadResponse,
+    FileInfo,
+    FileUploadResponse,
+    GlobResult,
+    GrepMatch,
+    GrepResult,
+    LsResult,
+    ReadResult,
+    SandboxBackendProtocol,
+    WriteResult,
+)
+from deepagents.backends.utils import (
+    InvalidGlobPatternError,
+    compile_grep_include_glob,
+    perform_string_replacement,
+    slice_read_response,
+)
+from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse, ToolCallRequest
+from langchain.agents.middleware.types import ModelCallResult
+from langchain_core.callbacks import AsyncCallbackHandler
+from langchain_core.messages import ToolMessage
+from langchain_core.outputs import LLMResult
+from langchain_core.tools import BaseTool
+from langgraph.types import Command
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.constants import HARNESS_SNAPSHOT_SKIP_DIRS
+from litellm.harness.context import SessionContext
+from litellm.harness.errors import SandboxError
+from litellm.harness.sandbox.base import CompletedRun, Sandbox
+
+T = TypeVar("T")
+
+# Module alias so ruff recognises `.exception()` as logging the swallowed error (BLE001).
+_logger = verbose_logger
+# Models cost_per_token could not price: logged once, then skipped (always 0.0).
+_UNPRICED_MODELS: set[str] = set()  # mutable-ok: process-wide log-once memo, grown as unpriced models are seen
+
+DEEPAGENTS_EXECUTE_TIMEOUT_SECONDS: Final = 120.0
+DEEPAGENTS_FS_TIMEOUT_SECONDS: Final = 60.0
+DEEPAGENTS_MAX_OUTPUT_BYTES: Final = 100_000
+_EXIT_NOT_FOUND: Final = 3
+_EXIT_NOT_DIR: Final = 4
+_EXIT_TIMEOUT: Final = 124
+_READ_ONLY_ERROR: Final = "Error: this session is read-only; files cannot be changed"
+_NO_EXECUTE_ERROR: Final = "Error: shell execution is disabled for this session"
+# $1 = directory. Prints "d/<name>" or "f/<name>" per entry ("/" never appears in a name).
+_LS_SCRIPT: Final = (
+    '[ -e "$1" ] || exit 3; [ -d "$1" ] || exit 4; cd "$1" || exit 5; '
+    'for f in * .[!.]* ..?*; do if [ -e "$f" ] || [ -L "$f" ]; then '
+    'if [ -d "$f" ]; then printf "d/%s\\n" "$f"; else printf "f/%s\\n" "$f"; fi; fi; done'
+)
+_DELETE_SCRIPT: Final = '[ -e "$1" ] || [ -L "$1" ] || exit 3; rm -rf -- "$1"'
+# $1 = path. Prints the symlink-resolved absolute path of its nearest existing ancestor (the
+# path itself when it exists). New files and new directories (`a/b/new.py`) resolve through
+# whatever part already exists, so a symlinked ancestor is still caught.
+_REALPATH_SCRIPT: Final = (
+    'p="$1"; while [ ! -e "$p" ] && [ ! -L "$p" ]; do q=$(dirname -- "$p"); '
+    '[ "$q" = "$p" ] && exit 3; p="$q"; done; realpath -- "$p"'
+)
+_GREP_LINE: Final = re.compile(r"^(.+?):(\d+):(.*)$")
+_FILTER_MIDDLEWARE_NAME: Final = "LiteLLMHarnessToolFilter"
+
+
+def _decode(data: bytes) -> str:
+    return data.decode("utf-8", errors="replace")
+
+
+def _ls_entries(base: str, stdout: str) -> Iterator[FileInfo]:
+    for line in stdout.splitlines():
+        kind, _, name = line.partition("/")
+        if name:
+            is_dir = kind == "d"
+            yield FileInfo(path=f"{base}/{name}" + ("/" if is_dir else ""), is_dir=is_dir)
+
+
+class SandboxBackend(SandboxBackendProtocol):  # pyright: ignore[reportUntypedBaseClass]  # deepagents/langchain are optional and not installed for type checking
+    """deepagents backend whose files and shell live in a litellm Sandbox."""
+
+    def __init__(
+        self,
+        sandbox: Sandbox,
+        *,
+        loop: asyncio.AbstractEventLoop,
+        writable: bool = True,
+        allow_execute: bool = True,
+    ) -> None:
+        self._sandbox = sandbox
+        self._loop = loop
+        self._root = posixpath.normpath(sandbox.workdir)
+        self._real_root: str | None = None
+        self._writable = writable
+        self._allow_execute = allow_execute
+        self._id = f"litellm-harness-{uuid.uuid4().hex[:8]}"
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    # -- paths --------------------------------------------------------------
+
+    def to_real(self, path: str) -> str:
+        """Sandbox path for a virtual path (or an absolute path already under workdir)."""
+        normalized = posixpath.normpath("/" + path.lstrip("/"))
+        if ".." in normalized.split("/"):
+            raise ValueError(f"path traversal not allowed: {path}")
+        if normalized == self._root or normalized.startswith(self._root + "/"):
+            return normalized
+        if normalized == "/":
+            return self._root
+        return self._root + normalized
+
+    async def to_confined(self, path: str) -> str:
+        """to_real, then resolve symlinks inside the sandbox and refuse anything outside workdir.
+
+        A repo can contain `link -> ~/.aws/credentials`; without this, read/grep/glob would
+        follow it and read host secrets even in read-only mode.
+        """
+        real = self.to_real(path)
+        done = await self._run(("sh", "-c", _REALPATH_SCRIPT, "sh", real))
+        resolved = done.stdout.strip()
+        if done.exit_code != 0 or not resolved:
+            raise ValueError(f"path not found: {path}")
+        root = await self._resolved_root()
+        if resolved != root and not resolved.startswith(root + "/"):
+            raise ValueError(f"path resolves outside the workspace: {path}")
+        return real
+
+    async def _resolved_root(self) -> str:
+        if self._real_root is None:
+            done = await self._run(("realpath", "--", self._root))
+            self._real_root = done.stdout.strip() if done.exit_code == 0 and done.stdout.strip() else self._root
+        return self._real_root
+
+    def to_virtual(self, real: str) -> str:
+        if real == self._root:
+            return "/"
+        if real.startswith(self._root + "/"):
+            return real[len(self._root) :]
+        return real
+
+    # -- sync bridge (deepagents only calls these outside the event loop) ---
+
+    def _sync(self, coro: Coroutine[Any, Any, T]) -> T:
+        try:
+            running = asyncio.get_running_loop()
+        except RuntimeError:
+            running = None
+        if running is self._loop:
+            coro.close()
+            raise RuntimeError("SandboxBackend sync methods cannot run on the event loop thread")
+        return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
+
+    async def _run(self, cmd: Sequence[str], timeout: float | None = DEEPAGENTS_FS_TIMEOUT_SECONDS) -> CompletedRun:
+        return await self._sandbox.run(cmd, timeout=timeout)
+
+    # -- ls -----------------------------------------------------------------
+
+    async def als(self, path: str) -> LsResult:
+        try:
+            real = await self.to_confined(path)
+            done = await self._run(("sh", "-c", _LS_SCRIPT, "sh", real))
+        except (ValueError, SandboxError) as e:
+            return LsResult(error=f"Path '{path}': {e}")
+        if done.exit_code == _EXIT_NOT_FOUND:
+            return LsResult(error=f"Path '{path}': path_not_found")
+        if done.exit_code == _EXIT_NOT_DIR:
+            return LsResult(error=f"Path '{path}': not_a_directory")
+        if done.exit_code != 0:
+            return LsResult(error=f"Path '{path}': {done.stderr.strip() or 'ls failed'}")
+        base = self.to_virtual(real).rstrip("/")
+        entries = sorted(_ls_entries(base, done.stdout), key=lambda e: e["path"])
+        return LsResult(entries=entries)
+
+    def ls(self, path: str) -> LsResult:
+        return self._sync(self.als(path))
+
+    # -- read / write / edit ------------------------------------------------
+
+    async def _read_bytes(self, path: str) -> bytes:
+        return await self._sandbox.read(await self.to_confined(path))
+
+    async def aread(self, file_path: str, offset: int = 0, limit: int = 2000) -> ReadResult:
+        try:
+            data = await self._read_bytes(file_path)
+        except ValueError as e:
+            return ReadResult(error=f"Error reading file '{file_path}': {e}")
+        except SandboxError:
+            return ReadResult(error=f"File '{file_path}' not found")
+        try:
+            text = data.decode("utf-8")
+        except UnicodeDecodeError:
+            encoded = base64.standard_b64encode(data).decode("ascii")
+            return ReadResult(file_data=FileData(content=encoded, encoding="base64"))
+        return slice_read_response(FileData(content=text, encoding="utf-8"), offset, limit)
+
+    def read(self, file_path: str, offset: int = 0, limit: int = 2000) -> ReadResult:
+        return self._sync(self.aread(file_path, offset, limit))
+
+    async def awrite(self, file_path: str, content: str) -> WriteResult:
+        if not self._writable:
+            return WriteResult(error=_READ_ONLY_ERROR)
+        try:
+            await self._sandbox.write(await self.to_confined(file_path), content.encode("utf-8"))
+        except (ValueError, SandboxError) as e:
+            return WriteResult(error=f"Error writing file '{file_path}': {e}")
+        return WriteResult(path=file_path)
+
+    def write(self, file_path: str, content: str) -> WriteResult:
+        return self._sync(self.awrite(file_path, content))
+
+    async def aedit(
+        self,
+        file_path: str,
+        old_string: str,
+        new_string: str,
+        replace_all: bool = False,
+    ) -> EditResult:
+        if not self._writable:
+            return EditResult(error=_READ_ONLY_ERROR)
+        try:
+            content = _decode(await self._read_bytes(file_path))
+        except ValueError as e:
+            return EditResult(error=f"Error editing file '{file_path}': {e}")
+        except SandboxError:
+            return EditResult(error=f"Error: File '{file_path}' not found")
+        old = old_string.replace("\r\n", "\n")
+        new = new_string.replace("\r\n", "\n")
+        replaced = perform_string_replacement(content.replace("\r\n", "\n"), old, new, replace_all)
+        if isinstance(replaced, str):
+            return EditResult(error=replaced)
+        new_content, occurrences = replaced
+        try:
+            await self._sandbox.write(await self.to_confined(file_path), new_content.encode("utf-8"))
+        except SandboxError as e:
+            return EditResult(error=f"Error editing file '{file_path}': {e}")
+        return EditResult(path=file_path, occurrences=int(occurrences))
+
+    def edit(
+        self,
+        file_path: str,
+        old_string: str,
+        new_string: str,
+        replace_all: bool = False,
+    ) -> EditResult:
+        return self._sync(self.aedit(file_path, old_string, new_string, replace_all))
+
+    async def adelete(self, file_path: str) -> DeleteResult:
+        if not self._writable:
+            return DeleteResult(error=_READ_ONLY_ERROR)
+        try:
+            real = await self.to_confined(file_path)
+            if real == self._root:
+                return DeleteResult(error="Error: refusing to delete the workspace root")
+            done = await self._run(("sh", "-c", _DELETE_SCRIPT, "sh", real))
+        except (ValueError, SandboxError) as e:
+            return DeleteResult(error=f"Error deleting '{file_path}': {e}")
+        if done.exit_code == _EXIT_NOT_FOUND:
+            return DeleteResult(error=f"Error: '{file_path}' not found")
+        if done.exit_code != 0:
+            return DeleteResult(error=f"Error deleting '{file_path}': {done.stderr.strip()}")
+        return DeleteResult(path=file_path)
+
+    def delete(self, file_path: str) -> DeleteResult:
+        return self._sync(self.adelete(file_path))
+
+    # -- glob / grep --------------------------------------------------------
+
+    def _find_cmd(self, root: str) -> tuple[str, ...]:
+        prune = tuple(
+            itertools.chain.from_iterable(
+                ("-o", "-name", name) if index else ("-name", name)
+                for index, name in enumerate(sorted(HARNESS_SNAPSHOT_SKIP_DIRS))
+            )
+        )
+        # -P: never follow symlinks, so a repo link to ~/.aws cannot pull host files in.
+        return ("find", "-P", root, "(", *prune, ")", "-prune", "-o", "-type", "f", "-print")
+
+    def _grep_cmd(self, pattern: str, root: str) -> tuple[str, ...]:
+        # grep only the regular files `find -P -type f` lists: symlinks are never followed,
+        # whatever grep implementation (GNU -R vs BSD -r) the sandbox has.
+        find_cmd = " ".join(shlex.quote(part) for part in self._find_cmd(root))
+        return ("sh", "-c", f'{find_cmd} | tr "\\n" "\\0" | xargs -0 grep -nHFI -e "$1" --', "sh", pattern)
+
+    async def aglob(self, pattern: str, path: str | None = None) -> GlobResult:
+        try:
+            matcher = compile_grep_include_glob(pattern)
+            root = await self.to_confined(path or "/")
+            done = await self._run(self._find_cmd(root))
+        except (InvalidGlobPatternError, ValueError, SandboxError) as e:
+            return GlobResult(error=str(e), matches=None)
+        if done.exit_code != 0 and not done.stdout:
+            return GlobResult(matches=[])  # mutable-ok: deepagents GlobResult.matches is typed list[FileInfo]
+        matches = sorted(
+            (
+                FileInfo(path=self.to_virtual(real), is_dir=False)
+                for real in done.stdout.splitlines()
+                if matcher(posixpath.relpath(real, root))
+            ),
+            key=lambda m: m["path"],
+        )
+        return GlobResult(matches=matches, truncated=done.exit_code != 0)
+
+    def _grep_matches(self, stdout: str, root: str, include: Callable[[str], bool] | None) -> Iterator[GrepMatch]:
+        for line in stdout.splitlines():
+            parsed = _GREP_LINE.match(line)
+            if parsed is None:
+                continue
+            real = parsed.group(1)
+            if include is None or include(posixpath.relpath(real, root)):
+                yield GrepMatch(path=self.to_virtual(real), line=int(parsed.group(2)), text=parsed.group(3))
+
+    def glob(self, pattern: str, path: str | None = None) -> GlobResult:
+        return self._sync(self.aglob(pattern, path))
+
+    async def agrep(
+        self,
+        pattern: str,
+        path: str | None = None,
+        glob: str | None = None,
+        *,
+        max_count: int | None = None,
+    ) -> GrepResult:
+        try:
+            include = compile_grep_include_glob(glob) if glob else None
+            root = await self.to_confined(path or "/")
+            done = await self._run(self._grep_cmd(pattern, root))
+        except (InvalidGlobPatternError, ValueError, SandboxError) as e:
+            return GrepResult(error=f"Path '{path or '/'}': {e}")
+        if done.exit_code not in (0, 1) and not done.stdout:
+            return GrepResult(error=f"Path '{path or '/'}': {done.stderr.strip() or 'grep failed'}")
+        matches = list(  # mutable-ok: GrepResult.matches is list[GrepMatch]
+            self._grep_matches(done.stdout, root, include)
+        )
+        if max_count is not None and len(matches) > max_count:
+            return GrepResult(matches=matches[:max_count], truncated=True)
+        return GrepResult(matches=matches)
+
+    def grep(
+        self,
+        pattern: str,
+        path: str | None = None,
+        glob: str | None = None,
+        *,
+        max_count: int | None = None,
+    ) -> GrepResult:
+        return self._sync(self.agrep(pattern, path, glob, max_count=max_count))
+
+    # -- upload / download --------------------------------------------------
+
+    async def _upload_one(self, path: str, data: bytes) -> FileUploadResponse:
+        if not self._writable:
+            return FileUploadResponse(path=path, error="permission_denied")
+        try:
+            await self._sandbox.write(await self.to_confined(path), data)
+        except ValueError:
+            return FileUploadResponse(path=path, error="invalid_path")
+        except SandboxError as e:
+            return FileUploadResponse(path=path, error=str(e))
+        return FileUploadResponse(path=path)
+
+    async def aupload_files(
+        self,
+        files: list[tuple[str, bytes]],  # mutable-ok: signature fixed by deepagents BackendProtocol
+    ) -> list[FileUploadResponse]:  # mutable-ok: return type fixed by deepagents BackendProtocol
+        return [  # mutable-ok: BackendProtocol returns a list
+            await self._upload_one(path, data) for path, data in files
+        ]
+
+    def upload_files(
+        self,
+        files: list[tuple[str, bytes]],  # mutable-ok: signature fixed by deepagents BackendProtocol
+    ) -> list[FileUploadResponse]:  # mutable-ok: return type fixed by deepagents BackendProtocol
+        return self._sync(self.aupload_files(files))
+
+    async def _download_one(self, path: str) -> FileDownloadResponse:
+        try:
+            return FileDownloadResponse(path=path, content=await self._read_bytes(path))
+        except ValueError:
+            return FileDownloadResponse(path=path, error="invalid_path")
+        except SandboxError:
+            return FileDownloadResponse(path=path, error="file_not_found")
+
+    async def adownload_files(
+        self,
+        paths: list[str],  # mutable-ok: signature fixed by deepagents BackendProtocol
+    ) -> list[FileDownloadResponse]:  # mutable-ok: return type fixed by deepagents BackendProtocol
+        return [await self._download_one(path) for path in paths]  # mutable-ok: BackendProtocol returns a list
+
+    def download_files(
+        self,
+        paths: list[str],  # mutable-ok: signature fixed by deepagents BackendProtocol
+    ) -> list[FileDownloadResponse]:  # mutable-ok: return type fixed by deepagents BackendProtocol
+        return self._sync(self.adownload_files(paths))
+
+    # -- execute ------------------------------------------------------------
+
+    async def aexecute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse:
+        if not self._allow_execute:
+            return ExecuteResponse(output=_NO_EXECUTE_ERROR, exit_code=1)
+        limit = float(timeout) if timeout else DEEPAGENTS_EXECUTE_TIMEOUT_SECONDS
+        try:
+            done = await self._run(("sh", "-c", command), timeout=limit)
+        except SandboxError as e:
+            return ExecuteResponse(output=f"Error: {e}", exit_code=_EXIT_TIMEOUT)
+        output = done.stdout
+        if done.stderr:
+            output = f"{output}\n{done.stderr}" if output else done.stderr
+        truncated = len(output.encode("utf-8")) > DEEPAGENTS_MAX_OUTPUT_BYTES
+        if truncated:
+            output = output.encode("utf-8")[:DEEPAGENTS_MAX_OUTPUT_BYTES].decode("utf-8", errors="ignore")
+        return ExecuteResponse(output=output, exit_code=done.exit_code, truncated=truncated)
+
+    def execute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse:
+        return self._sync(self.aexecute(command, timeout=timeout))
+
+
+def _tool_name(tool: BaseTool | Mapping[str, object]) -> str | None:
+    name = tool.get("name") if isinstance(tool, Mapping) else tool.name
+    return name if isinstance(name, str) else None
+
+
+def _blocked_message(request: ToolCallRequest, blocked: frozenset[str]) -> ToolMessage | None:
+    name = request.tool_call["name"]
+    if name not in blocked:
+        return None
+    return ToolMessage(
+        content=f"Error: {name} is disabled for this session.",
+        tool_call_id=request.tool_call["id"] or "",
+        name=name,
+        status="error",
+    )
+
+
+class ToolFilterMiddleware(AgentMiddleware):  # pyright: ignore[reportUntypedBaseClass]  # deepagents/langchain are optional and not installed for type checking
+    """Hide tools from the model and refuse calls to them (disable_tools / permissions)."""
+
+    def __init__(self, blocked: frozenset[str]) -> None:
+        super().__init__()
+        self._blocked = blocked
+
+    @property
+    def name(self) -> str:
+        return _FILTER_MIDDLEWARE_NAME
+
+    def _filtered(self, request: ModelRequest) -> ModelRequest:
+        return request.override(
+            tools=[  # mutable-ok: ModelRequest.tools is a list
+                t for t in request.tools if _tool_name(t) not in self._blocked
+            ]
+        )
+
+    def wrap_model_call(
+        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
+    ) -> ModelCallResult:
+        return handler(self._filtered(request))
+
+    async def awrap_model_call(
+        self, request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]]
+    ) -> ModelCallResult:
+        return await handler(self._filtered(request))
+
+    def wrap_tool_call(
+        self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage | Command]
+    ) -> ToolMessage | Command:
+        return _blocked_message(request, self._blocked) or handler(request)
+
+    async def awrap_tool_call(
+        self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]]
+    ) -> ToolMessage | Command:
+        return _blocked_message(request, self._blocked) or await handler(request)
+
+
+def _number(value: object) -> float | None:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        return None
+    return float(value)
+
+
+def message_cost(message: object, cost_model: str | None, input_tokens: int, output_tokens: int) -> float:
+    """Cost of one model call: response_cost reported by litellm, else cost_per_token, else 0."""
+    metadata = getattr(message, "response_metadata", None) or MappingProxyType({})
+    reported = _number(metadata.get("response_cost"))
+    if reported is not None:
+        return reported
+    if not cost_model or cost_model in _UNPRICED_MODELS:
+        return 0.0
+    try:
+        prompt_cost, completion_cost = litellm.cost_per_token(
+            model=cost_model,
+            prompt_tokens=input_tokens,
+            completion_tokens=output_tokens,
+        )
+        return float(prompt_cost) + float(completion_cost)
+    except Exception:  # litellm raises plain Exception for unmapped models
+        _UNPRICED_MODELS.add(cost_model)
+        _logger.exception("harness deepagents: no cost for %s; counting its calls as 0", cost_model)
+        return 0.0
+
+
+def record_llm_usage(ctx: SessionContext, cost_model: str | None, response: LLMResult) -> None:
+    """Add one model call's tokens and cost to the session counters. Never raises."""
+    try:
+        ctx.calls += 1
+        for generations in response.generations:
+            for generation in generations:
+                message = getattr(generation, "message", None)
+                usage = getattr(message, "usage_metadata", None) or MappingProxyType({})
+                input_tokens = int(usage.get("input_tokens") or 0)
+                output_tokens = int(usage.get("output_tokens") or 0)
+                ctx.input_tokens += input_tokens
+                ctx.output_tokens += output_tokens
+                ctx.cost += message_cost(message, cost_model, input_tokens, output_tokens)
+    except Exception:
+        # Usage accounting must never fail a turn; log with traceback and move on.
+        _logger.exception("harness deepagents: usage accounting failed")
+
+
+class UsageCallback(AsyncCallbackHandler):  # pyright: ignore[reportUntypedBaseClass]  # deepagents/langchain are optional and not installed for type checking
+    """Counts every model call in the graph, subagents and summarization included."""
+
+    def __init__(self, ctx: SessionContext, cost_model: str | None) -> None:
+        self._ctx = ctx
+        self._cost_model = cost_model
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: object) -> None:
+        record_llm_usage(self._ctx, self._cost_model, response)

litellm/llms/deepagents/harness/transformation.py

@@ -0,0 +1,313 @@
+"""
+Deep Agents harness config: LangChain `deepagents` running in your Python process.
+
+Pure translation only: model kwargs (gateway mode uses `litellm_proxy/<group>` with the same
+attribution headers the CLI endpoint adds), permission and tool filtering, and LangGraph
+stream chunks to events. `litellm/harness/handlers/deepagents_handler.py` builds the agent,
+streams it, answers approvals and counts usage.
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Final
+
+from litellm.harness.options import DeepAgentsOptions
+from litellm.harness.types import (
+    Capabilities,
+    Event,
+    Harness,
+    Reasoning,
+    Text,
+    ToolCall,
+    ToolResult,
+)
+from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
+
+if TYPE_CHECKING:
+    from litellm.harness.context import SessionContext
+
+INSTALL_HINT: Final = "Deep Agents is not installed. Run: pip install deepagents langchain-litellm"
+SKILLS_DIR: Final = ".deepagents/skills"
+# Graph supersteps per agent turn (model node, tools node, middleware hooks) for recursion_limit.
+DEEPAGENTS_STEPS_PER_TURN: Final = 6
+DEEPAGENTS_BASE_RECURSION_LIMIT: Final = 25
+DEEPAGENTS_DEFAULT_RECURSION_LIMIT: Final = 1000
+_MODEL_NODE: Final = "model"
+# Only these graph nodes produce new messages; middleware hooks may re-emit history.
+_EVENT_NODES: Final = frozenset({"model", "tools"})
+
+NORMALIZED_TO_NATIVE: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "read": "read_file",
+        "write": "write_file",
+        "edit": "edit_file",
+        "bash": "execute",
+    }
+)
+NATIVE_TO_NORMALIZED: Final[Mapping[str, str]] = MappingProxyType({v: k for k, v in NORMALIZED_TO_NATIVE.items()})
+BUILTIN_TOOLS: Final = frozenset(
+    {
+        "ls",
+        "read_file",
+        "write_file",
+        "edit_file",
+        "delete",
+        "glob",
+        "grep",
+        "execute",
+        "write_todos",
+        "task",
+    }
+)
+WRITE_TOOLS: Final = frozenset({"write_file", "edit_file", "delete"})
+EXECUTE_TOOLS: Final = frozenset({"execute"})
+APPROVAL_TOOLS: Final = WRITE_TOOLS | EXECUTE_TOOLS
+_APPROVAL_DECISIONS: Final = ("approve", "reject")
+
+
+# ---------------------------------------------------------------------------
+# Pure helpers (unit tested directly; kept module-level so they port cleanly)
+# ---------------------------------------------------------------------------
+
+
+def gateway_headers(
+    ctx: SessionContext,
+) -> dict[str, str]:  # mutable-ok: ChatLiteLLM.extra_headers is a pydantic dict field
+    """Same attribution headers the session endpoint adds for CLI harnesses."""
+    metadata = ctx.metadata
+    metadata_json = json.dumps(dict(metadata), default=str) if metadata else None  # mutable-ok: for json.dumps
+    metadata_header = (("x-litellm-spend-logs-metadata", metadata_json),) if metadata_json is not None else ()
+    return dict(  # mutable-ok: ChatLiteLLM.extra_headers is a pydantic dict field
+        (("x-litellm-tags", f"harness,{ctx.harness.value}"), *metadata_header)
+    )
+
+
+def chat_model_kwargs(
+    ctx: SessionContext,
+) -> dict[str, Any]:  # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs
+    """ChatLiteLLM constructor kwargs for gateway or SDK mode."""
+    if not ctx.model:
+        raise ValueError("Harness.DEEPAGENTS needs model=")
+    if ctx.gateway is not None:
+        return {  # mutable-ok: ChatLiteLLM constructor kwargs, splatted as **kwargs
+            "model": f"litellm_proxy/{ctx.model}",
+            "api_base": ctx.gateway.api_base,
+            "api_key": ctx.gateway.api_key,
+            "extra_headers": gateway_headers(ctx),
+        }
+    return {"model": ctx.model, "api_key": ctx.api_key, "api_base": ctx.api_base}  # mutable-ok: ChatLiteLLM kwargs
+
+
+def native_tool_name(name: str) -> str:
+    return NORMALIZED_TO_NATIVE.get(name, name)
+
+
+def normalized_tool_name(native: str) -> str:
+    return NATIVE_TO_NORMALIZED.get(native, native)
+
+
+def blocked_tools(permissions: str, disable_tools: Sequence[str]) -> frozenset[str]:
+    """Native tool names the model must not see or call."""
+    disabled = frozenset(native_tool_name(name) for name in disable_tools)
+    if permissions == "read-only":
+        return disabled | WRITE_TOOLS | EXECUTE_TOOLS
+    if permissions == "edit":
+        return disabled | EXECUTE_TOOLS
+    return disabled
+
+
+def interrupt_config(
+    permissions: str, blocked: frozenset[str]
+) -> dict[str, Any] | None:  # mutable-ok: deepagents create_deep_agent(interrupt_on=) takes a dict
+    """interrupt_on for permissions='ask': approve/reject every mutating built-in."""
+    if permissions != "ask":
+        return None
+    return {  # mutable-ok: deepagents interrupt_on config (dict of InterruptOnConfig with list allowed_decisions)
+        name: {"allowed_decisions": list(_APPROVAL_DECISIONS)}  # mutable-ok: deepagents InterruptOnConfig shape
+        for name in sorted(APPROVAL_TOOLS - blocked)
+    }
+
+
+def recursion_limit(ctx: SessionContext) -> int:
+    options = ctx.options if isinstance(ctx.options, DeepAgentsOptions) else None
+    if options is not None and options.recursion_limit is not None:
+        return options.recursion_limit
+    if ctx.max_turns is not None:
+        return DEEPAGENTS_BASE_RECURSION_LIMIT + ctx.max_turns * DEEPAGENTS_STEPS_PER_TURN
+    return DEEPAGENTS_DEFAULT_RECURSION_LIMIT
+
+
+def content_text(content: object) -> str:
+    """Plain text of a LangChain message content (str or content blocks)."""
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, list):
+        return ""
+    return "".join(
+        block if isinstance(block, str) else block.get("text", "") for block in content if _is_text_block(block)
+    )
+
+
+def _is_text_block(block: object) -> bool:
+    return isinstance(block, str) or (isinstance(block, dict) and block.get("type") == "text")
+
+
+def reasoning_text(message: object) -> str:
+    """Reasoning deltas from additional_kwargs or reasoning/thinking content blocks."""
+    extra = getattr(message, "additional_kwargs", None) or MappingProxyType({})
+    reasoning = extra.get("reasoning_content")
+    if isinstance(reasoning, str) and reasoning:
+        return reasoning
+    content = getattr(message, "content", None)
+    if not isinstance(content, list):
+        return ""
+    return "".join(
+        str(block.get("reasoning") or block.get("thinking") or "")
+        for block in content
+        if isinstance(block, dict) and block.get("type") in ("reasoning", "thinking")
+    )
+
+
+def stream_events(
+    message: object,
+) -> list[Event]:  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    """Text / Reasoning deltas for one streamed message chunk."""
+    if getattr(message, "type", None) not in ("AIMessageChunk", "ai"):
+        return []  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    reasoning = reasoning_text(message)
+    text = content_text(getattr(message, "content", ""))
+    reasoning_events: tuple[Event, ...] = (Reasoning(delta=reasoning),) if reasoning else ()
+    text_events: tuple[Event, ...] = (Text(delta=text),) if text else ()
+    return [  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+        *reasoning_events,
+        *text_events,
+    ]
+
+
+def tool_call_event(call: Mapping[str, Any]) -> ToolCall:
+    native = str(call.get("name") or "")
+    args = call.get("args")
+    return ToolCall(
+        id=str(call.get("id") or ""),
+        name=normalized_tool_name(native),
+        native_name=native,
+        input=dict(args) if isinstance(args, Mapping) else {"args": args},  # mutable-ok: ToolCall.input is a dict
+        builtin=native in BUILTIN_TOOLS,
+    )
+
+
+def _node_messages(update: Mapping[Any, Any]) -> Iterator[object]:
+    for node, delta in update.items():
+        if node not in _EVENT_NODES or not isinstance(delta, Mapping):
+            continue
+        messages = delta.get("messages")
+        if isinstance(messages, list):
+            yield from messages
+
+
+def update_events(
+    update: object, skip_tools: frozenset[str]
+) -> list[Event]:  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    """ToolCall / ToolResult events from one `updates` stream chunk (node -> state delta)."""
+    if not isinstance(update, Mapping):
+        return []  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    return list(  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+        itertools.chain.from_iterable(_message_events(message, skip_tools) for message in _node_messages(update))
+    )
+
+
+def _message_events(message: object, skip_tools: frozenset[str]) -> tuple[Event, ...]:
+    kind = getattr(message, "type", None)
+    if kind == "ai":
+        calls = getattr(message, "tool_calls", None) or ()
+        return tuple(tool_call_event(call) for call in calls if call.get("name") not in skip_tools)
+    if kind == "tool" and getattr(message, "name", None) not in skip_tools:
+        return (
+            ToolResult(
+                id=str(getattr(message, "tool_call_id", "") or ""),
+                output=content_text(getattr(message, "content", "")),
+                is_error=getattr(message, "status", None) == "error",
+            ),
+        )
+    return ()
+
+
+def interrupts_in(
+    update: object,
+) -> list[Any]:  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    if not isinstance(update, Mapping):
+        return []  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    found = update.get("__interrupt__")
+    items = tuple(found) if isinstance(found, (list, tuple)) else ()
+    return list(items)  # mutable-ok: list return; callers/tests compare to lists
+
+
+def final_ai_text(messages: Sequence[Any]) -> str:
+    for message in reversed(messages):
+        if getattr(message, "type", None) == "ai":
+            text = content_text(getattr(message, "content", ""))
+            if text:
+                return text
+    return ""
+
+
+def structured_json(value: object) -> str | None:
+    if value is None:
+        return None
+    dump = getattr(value, "model_dump_json", None)
+    if callable(dump):
+        return str(dump())
+    return json.dumps(value, default=str)
+
+
+def approval_requests(
+    interrupt_value: object,
+) -> list[Mapping[str, Any]]:  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    """action_requests of a HumanInTheLoopMiddleware interrupt payload."""
+    if not isinstance(interrupt_value, Mapping):
+        return []  # mutable-ok: returns a list; existing callers/tests compare it to list literals
+    requests = interrupt_value.get("action_requests")
+    kept = tuple(r for r in requests if isinstance(r, Mapping)) if isinstance(requests, list) else ()
+    return list(kept)  # mutable-ok: list return; callers/tests compare to lists
+
+
+def decision(allowed: bool, reason: str) -> dict[str, Any]:  # mutable-ok: LangGraph resume payload (HITL decision dict)
+    if allowed:
+        return {"type": "approve"}  # mutable-ok: LangGraph resume payload (HITL decision dict)
+    return {  # mutable-ok: LangGraph HITL decision
+        "type": "reject",
+        "message": reason or "The user denied this tool call.",
+    }
+
+
+@dataclass
+class TurnState:
+    """Mutable state across the stream passes of one turn."""
+
+    interrupts: tuple[Any, ...] = ()
+
+
+class DeepAgentsHarnessConfig(BaseHarnessConfig):
+    harness = Harness.DEEPAGENTS
+    options_type = DeepAgentsOptions
+    uses_model_endpoint = False
+    capabilities = Capabilities(
+        structured_output=True,
+        tool_approval=True,
+        tool_filtering=True,
+        history=True,
+        custom_tools=True,
+        skills=True,
+        resume=True,
+        permission_modes=frozenset({"read-only", "ask", "edit", "full"}),
+    )
+
+    def validate_environment(self, ctx: SessionContext) -> None:
+        self.get_options(ctx)
+        if not ctx.model:
+            raise ValueError("Harness.DEEPAGENTS needs model=")

litellm/llms/opencode/harness/transformation.py

@@ -0,0 +1,427 @@
+"""
+OpenCode harness config: `opencode run --format json`, once per turn.
+
+Every model call goes to one custom provider (`litellm`, `@ai-sdk/openai-compatible`,
+bundled in the binary) whose baseURL is the per-session endpoint. The config travels in
+OPENCODE_CONFIG_CONTENT, which opencode applies after global and project config, so a
+repo's own opencode.json cannot redirect model calls. The token is never in argv or env:
+the config references it with `{file:<private_dir>/token}`. XDG dirs point at a persisted
+LiteLLM-owned root so the user's opencode config and auth are never read, and the session
+DB outlives a session for resume. Verified against opencode 1.14.41.
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Final
+
+from litellm.harness.errors import CapabilityUnsupported, HarnessError, OptionsMismatch
+from litellm.harness.options import OpenCodeOptions
+from litellm.harness.types import (
+    Capabilities,
+    Event,
+    Harness,
+    PermissionMode,
+    Reasoning,
+    Text,
+    ToolCall,
+    ToolResult,
+)
+from litellm.llms.base_llm.harness.transformation import (
+    BaseCLIHarnessConfig,
+    HarnessSessionSetup,
+    HarnessTurnError,
+    HarnessTurnRequest,
+    HarnessTurnResponse,
+    event_list,
+)
+from litellm.llms.base_llm.harness.utils import (
+    last_json_object,
+    native_tool_names,
+    normalize_tool_name,
+    stderr_tail_text,
+    structured_output_instruction,
+)
+
+if TYPE_CHECKING:
+    from litellm.harness.context import SessionContext
+
+OPENCODE_BINARY: Final = "opencode"
+OPENCODE_PROVIDER_ID: Final = "litellm"
+OPENCODE_PROVIDER_NPM: Final = "@ai-sdk/openai-compatible"
+# A fixed title skips opencode's extra title-generation model call on the first turn.
+OPENCODE_SESSION_TITLE: Final = "litellm-harness"
+TOKEN_FILENAME: Final = "token"
+INSTRUCTIONS_FILENAME: Final = "instructions.md"
+XDG_DIRNAME: Final = "xdg"
+XDG_SUBDIRS: Final = ("config", "data", "state", "cache")
+
+# Env that keeps opencode off the network (except the endpoint) and away from ~/.claude.
+OPENCODE_ISOLATION_ENV: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "OPENCODE_DISABLE_AUTOUPDATE": "1",
+        "OPENCODE_DISABLE_MODELS_FETCH": "1",
+        "OPENCODE_DISABLE_LSP_DOWNLOAD": "1",
+        "OPENCODE_DISABLE_SHARE": "1",
+        "OPENCODE_DISABLE_DEFAULT_PLUGINS": "1",
+        "OPENCODE_DISABLE_CLAUDE_CODE": "1",
+        "OPENCODE_DISABLE_EXTERNAL_SKILLS": "1",
+        # Blank (falsy to opencode) so an inherited value can't add config, auth or rules.
+        "OPENCODE_CONFIG": "",
+        "OPENCODE_CONFIG_DIR": "",
+        "OPENCODE_PERMISSION": "",
+        "OPENCODE_AUTH_CONTENT": "",
+    }
+)
+
+MANAGED_CONFIG_KEYS: Final = frozenset(
+    {
+        "provider",
+        "model",
+        "small_model",
+        "permission",
+        "tools",
+        "enabled_providers",
+        "disabled_providers",
+        # plugins run arbitrary code as the host user; runs always use --pure
+        "plugin",
+    }
+)
+AGENT_MANAGED_KEYS: Final = frozenset({"permission", "tools", "model"})
+
+# Later keys win in opencode, so disable_tools denies go last. `opencode run` auto-rejects
+# anything left at "ask", so no mode leaves a tool on ask.
+PERMISSION_RULES: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
+    {
+        "read-only": MappingProxyType({"edit": "deny", "bash": "deny", "webfetch": "deny"}),
+        "edit": MappingProxyType({"edit": "allow", "bash": "deny", "webfetch": "allow"}),
+        "full": MappingProxyType({"*": "allow"}),
+    }
+)
+
+# opencode gates write, edit and apply_patch with the single `edit` permission.
+NORMALIZED_TO_NATIVE: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
+    {
+        "read": ("read",),
+        "write": ("edit",),
+        "edit": ("edit",),
+        "bash": ("bash",),
+        "glob": ("glob",),
+        "grep": ("grep",),
+        "ls": ("list",),
+        "web_search": ("webfetch", "websearch"),
+    }
+)
+
+NATIVE_TO_NORMALIZED: Final[Mapping[str, str]] = MappingProxyType(
+    {
+        "read": "read",
+        "write": "write",
+        "edit": "edit",
+        "multiedit": "edit",
+        "patch": "edit",
+        "apply_patch": "edit",
+        "bash": "bash",
+        "glob": "glob",
+        "grep": "grep",
+        "list": "ls",
+        "webfetch": "web_search",
+        "websearch": "web_search",
+    }
+)
+
+OPENCODE_BUILTIN_TOOLS: Final = frozenset(
+    {
+        *NATIVE_TO_NORMALIZED,
+        "task",
+        "todowrite",
+        "todoread",
+        "skill",
+        "invalid",
+        "question",
+        "lsp",
+        "codesearch",
+        "plan_enter",
+        "plan_exit",
+    }
+)
+
+
+@dataclass
+class OpenCodeStreamState:
+    """What the parser has learned from one `opencode run`."""
+
+    session_id: str | None = None
+    final_text: str = ""
+    error: str | None = None
+    step_texts: Sequence[str] = ()
+
+
+def _as_dict(value: object) -> Mapping[str, Any]:
+    return value if isinstance(value, dict) else MappingProxyType({})
+
+
+def _tool_events(part: Mapping[str, Any]) -> Sequence[Event]:
+    native = str(part.get("tool") or "")
+    call_id = str(part.get("callID") or part.get("id") or "")
+    state = _as_dict(part.get("state"))
+    tool_input = state.get("input")
+    call = ToolCall(
+        id=call_id,
+        name=normalize_tool_name(native, NATIVE_TO_NORMALIZED),
+        native_name=native,
+        input=tool_input if isinstance(tool_input, dict) else MappingProxyType({"input": tool_input}),
+        builtin=native in OPENCODE_BUILTIN_TOOLS,
+    )
+    if state.get("status") == "error":
+        message = str(state.get("error") or state.get("output") or "tool failed")
+        return event_list(call, ToolResult(id=call_id, output=message, is_error=True))
+    output = state.get("output")
+    text = output if isinstance(output, str) else json.dumps(output)
+    # The `invalid` pseudo-tool is how opencode reports a call to an unavailable tool.
+    return event_list(call, ToolResult(id=call_id, output=text, is_error=native == "invalid"))
+
+
+def _error_message(error: object) -> str:
+    if not isinstance(error, dict):
+        return str(error or "opencode reported an error")
+    data = error.get("data")
+    if isinstance(data, dict) and data.get("message"):
+        return str(data["message"])
+    return str(error.get("name") or "opencode reported an error")
+
+
+def validate_user_config(config: Mapping[str, Any]) -> None:
+    """Reject OpenCodeOptions.config keys LiteLLM manages (or that bypass permissions)."""
+    for key in config:
+        if key in MANAGED_CONFIG_KEYS:
+            raise OptionsMismatch(
+                f"OpenCodeOptions.config[{key!r}] is managed by LiteLLM; use the matching "
+                "agent() argument (model=, permissions=, disable_tools=) instead"
+            )
+    for section in ("agent", "mode"):
+        entries = config.get(section)
+        if entries is None:
+            continue
+        if not isinstance(entries, Mapping):
+            raise OptionsMismatch(f"OpenCodeOptions.config[{section!r}] must be a mapping")
+        for name, agent in entries.items():
+            managed = AGENT_MANAGED_KEYS & frozenset(agent or ())
+            if managed:
+                raise OptionsMismatch(
+                    f"OpenCodeOptions.config[{section!r}][{name!r}] sets {sorted(managed)}, "
+                    "which LiteLLM manages; use permissions=/disable_tools=/model= instead"
+                )
+
+
+def permission_rules(permissions: PermissionMode, disable_tools: Sequence[str]) -> Mapping[str, str]:
+    """opencode `permission` config for a mode plus denies for disable_tools."""
+    if permissions not in PERMISSION_RULES:
+        raise CapabilityUnsupported(
+            f"Harness.OPENCODE does not support permissions={permissions!r} (supported: {sorted(PERMISSION_RULES)})"
+        )
+    denied: Final = native_tool_names(disable_tools, NORMALIZED_TO_NATIVE)
+    # Denies go last (later keys win in opencode), so drop them from the mode rules first.
+    kept: Final = ((key, value) for key, value in PERMISSION_RULES[permissions].items() if key not in denied)
+    rules: Final = itertools.chain(kept, ((native, "deny") for native in denied))
+    return dict(rules)  # mutable-ok: opencode config JSON
+
+
+def build_opencode_config(
+    *,
+    model: str,
+    base_url: str,
+    token_path: str,
+    permissions: PermissionMode,
+    disable_tools: Sequence[str] = (),
+    user_config: Mapping[str, Any] | None = None,
+    instructions_path: str | None = None,
+    skills_path: str | None = None,
+) -> Mapping[str, Any]:
+    """The full opencode config: user config underneath, LiteLLM-managed keys on top."""
+    user: Final = user_config or MappingProxyType({})
+    validate_user_config(user)
+    qualified = f"{OPENCODE_PROVIDER_ID}/{model}"
+    extra_instructions: Final = (instructions_path,) if instructions_path else ()
+    instructions: Final = [*(user.get("instructions") or ()), *extra_instructions]  # mutable-ok: opencode config JSON
+    user_skills = _as_dict(user.get("skills"))
+    extra_skills: Final = (skills_path,) if skills_path else ()
+    skill_paths: Final = [*(user_skills.get("paths") or ()), *extra_skills]  # mutable-ok: opencode config JSON
+    options: Final = {"baseURL": base_url, "apiKey": "{file:" + token_path + "}"}  # mutable-ok: opencode config JSON
+    models: Final[dict[str, Any]] = {model: {}}  # mutable-ok: opencode config JSON
+    provider: Final = {  # mutable-ok: opencode config JSON
+        "npm": OPENCODE_PROVIDER_NPM,
+        "name": "LiteLLM",
+        "options": options,
+        "models": models,
+    }
+    managed: Final = {  # mutable-ok: opencode config JSON
+        "provider": {OPENCODE_PROVIDER_ID: provider},  # mutable-ok: opencode config JSON
+        "enabled_providers": [OPENCODE_PROVIDER_ID],  # mutable-ok: opencode config JSON
+        "model": qualified,
+        "small_model": qualified,
+        "permission": permission_rules(permissions, disable_tools),
+        "autoupdate": False,
+        "share": "disabled",
+    }
+    skills: Final = {**user_skills, "paths": skill_paths}  # mutable-ok: opencode config JSON
+    optional: Final = (("instructions", instructions), ("skills", skills if skill_paths else None))
+    present: Final = ((key, value) for key, value in optional if value)
+    return {**user, **managed, **dict(present)}  # mutable-ok: opencode config JSON
+
+
+def build_instructions(ctx: SessionContext) -> str | None:
+    schema_part: Final = (
+        structured_output_instruction(ctx.output.model_json_schema()) if ctx.output is not None else None
+    )
+    sections: Final = tuple(section for section in (ctx.instructions, schema_part) if section)
+    return "\n\n".join(sections) if sections else None
+
+
+def turn_prompt(ctx: SessionContext, prompt: str) -> str:
+    """Repeat the schema instruction in the user turn; system instructions alone are too weak."""
+    if ctx.output is None:
+        return prompt
+    return f"{prompt}\n\n{structured_output_instruction(ctx.output.model_json_schema())}"
+
+
+class OpenCodeHarnessConfig(BaseCLIHarnessConfig):
+    harness = Harness.OPENCODE
+    options_type = OpenCodeOptions
+    capabilities = Capabilities(
+        structured_output=True,
+        tool_approval=False,
+        tool_filtering=True,
+        history=False,
+        custom_tools=False,
+        skills=True,
+        resume=True,
+        permission_modes=frozenset({"read-only", "edit", "full"}),
+    )
+
+    def get_binary(self) -> str:
+        return OPENCODE_BINARY
+
+    def get_install_hint(self) -> str:
+        return "npm install -g opencode-ai (or brew install sst/tap/opencode)"
+
+    def validate_environment(self, ctx: SessionContext) -> None:
+        options: OpenCodeOptions = self.get_options(ctx)
+        validate_user_config(options.config)
+
+    def transform_session_setup(self, ctx: SessionContext, private_dir: str) -> HarnessSessionSetup:
+        if ctx.endpoint is None:
+            raise HarnessError("OpenCode needs the session model endpoint")
+        model = ctx.model or ctx.endpoint.model
+        if not model:
+            raise ValueError("Harness.OPENCODE needs model= (a gateway model group or litellm model)")
+        options: OpenCodeOptions = self.get_options(ctx)
+        instructions = build_instructions(ctx)
+        token: Final = ctx.endpoint.token.encode("utf-8")
+        files: Final = (
+            MappingProxyType({TOKEN_FILENAME: token, INSTRUCTIONS_FILENAME: instructions.encode("utf-8")})
+            if instructions is not None
+            else MappingProxyType({TOKEN_FILENAME: token})
+        )
+        config = build_opencode_config(
+            model=model,
+            base_url=ctx.sandbox.host_url(ctx.endpoint.port).rstrip("/") + "/v1",
+            token_path=f"{private_dir}/{TOKEN_FILENAME}",
+            permissions=ctx.permissions,
+            disable_tools=ctx.disable_tools,
+            user_config=options.config,
+            instructions_path=f"{private_dir}/{INSTRUCTIONS_FILENAME}" if instructions is not None else None,
+            skills_path=f"{private_dir}/skills" if ctx.skills else None,
+        )
+        xdg: Final = MappingProxyType(
+            {f"XDG_{sub.upper()}_HOME": f"{private_dir}/{XDG_DIRNAME}/{sub}" for sub in XDG_SUBDIRS}
+        )
+        return HarnessSessionSetup(
+            files=files,
+            persisted_dirs=((XDG_DIRNAME, "opencode"),),
+            skills_dir="skills",
+            env=MappingProxyType(
+                {**OPENCODE_ISOLATION_ENV, **options.env, **xdg, "OPENCODE_CONFIG_CONTENT": json.dumps(config)}
+            ),
+        )
+
+    def transform_turn_request(
+        self,
+        ctx: SessionContext,
+        setup: HarnessSessionSetup,
+        private_dir: str,
+        prompt: str,
+        native_session_id: str | None,
+    ) -> HarnessTurnRequest:
+        options: OpenCodeOptions = self.get_options(ctx)
+        model = ctx.model or (ctx.endpoint.model if ctx.endpoint else None)
+        # --pure: never load plugins. A repo's .opencode/plugin/*.js would otherwise run as the
+        # host user at startup, before any tool permission applies.
+        argv: Final = (
+            OPENCODE_BINARY,
+            "run",
+            "--pure",
+            "--format",
+            "json",
+            "--thinking",
+            "-m",
+            f"{OPENCODE_PROVIDER_ID}/{model}",
+            *(("--agent", options.agent) if options.agent else ()),
+            *(("--session", native_session_id) if native_session_id else ("--title", OPENCODE_SESSION_TITLE)),
+        )
+        # The prompt goes on stdin; opencode appends non-TTY stdin to the message.
+        return HarnessTurnRequest(argv=argv, env=setup.env, stdin=turn_prompt(ctx, prompt), cwd=ctx.sandbox.workdir)
+
+    def create_stream_state(self) -> OpenCodeStreamState:
+        return OpenCodeStreamState()
+
+    def transform_stream_line(self, line: Mapping[str, Any], state: OpenCodeStreamState) -> Sequence[Event]:
+        """step_finish token counts are ignored on purpose: the session endpoint accounts usage."""
+        session_id = line.get("sessionID")
+        if session_id and state.session_id is None:
+            state.session_id = str(session_id)
+        event_type = line.get("type")
+        part = _as_dict(line.get("part"))
+        if event_type == "step_start":
+            state.step_texts = ()
+            return event_list()
+        if event_type == "text":
+            text = str(part.get("text") or "")
+            if not text:
+                return event_list()
+            state.step_texts = (*state.step_texts, text)
+            state.final_text = "\n\n".join(state.step_texts)
+            return event_list(Text(delta=text))
+        if event_type == "reasoning":
+            text = str(part.get("text") or "")
+            return event_list(Reasoning(delta=text)) if text else event_list()
+        if event_type == "tool_use":
+            return _tool_events(part)
+        if event_type == "error":
+            message = _error_message(line.get("error"))
+            state.error = f"{state.error}\n{message}" if state.error else message
+        return event_list()
+
+    def get_native_session_id(self, state: OpenCodeStreamState) -> str | None:
+        return state.session_id
+
+    def transform_turn_response(
+        self,
+        ctx: SessionContext,
+        state: OpenCodeStreamState,
+        exit_code: int,
+        stderr_tail: Sequence[str],
+    ) -> HarnessTurnResponse:
+        # opencode exits 0 after an `error` event, so check state first.
+        if state.error:
+            raise HarnessTurnError(f"opencode turn failed: {state.error}")
+        if exit_code != 0:
+            raise HarnessTurnError(
+                f"opencode exited with code {exit_code}: {stderr_tail_text(stderr_tail) or 'no output'}"
+            )
+        output_json = last_json_object(state.final_text) if ctx.output is not None else None
+        return HarnessTurnResponse(final_text=state.final_text, output_json=output_json)

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -1947,10 +1947,41 @@ class MCPRequestHandler:
         scope, the tool union and billing attribution are all just different reads of this one
         answer — computing it separately per consumer is how they drift (a throttle map scoped by
         roster instead of by grant charged unrelated teams' buckets)."""
-        return [
+        grants: Final = [
             (source, set(await MCPRequestHandler.get_allowed_mcp_servers(source, keyless_source=True)))
             for source in await MCPRequestHandler.admitted_subject_sources(auth, allowed_team_ids=allowed_team_ids)
         ]
+        scope: Final = await MCPRequestHandler._toolset_scope(auth)
+        if scope is None:
+            return grants
+        return [(source, granted & frozenset(scope)) for source, granted in grants]
+
+    @staticmethod
+    async def _toolset_scope(auth: UserAPIKeyAuth) -> dict[str, list[str]] | None:
+        """The ``server_id -> tools`` a namespaced toolset route pinned this subject to via
+        ``mcp_toolset_id``, or None on the aggregate scope. Every source's servers and tools are
+        intersected with it, so the route narrows a team grant exactly as it narrows the user's own."""
+        if auth.mcp_toolset_id is None:
+            return None
+        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+            global_mcp_server_manager,
+        )
+
+        return await global_mcp_server_manager.resolve_toolset_tool_permissions(
+            toolset_ids=[auth.mcp_toolset_id], requires_fresh_policy=auth.requires_fresh_policy
+        )
+
+    @staticmethod
+    async def _narrow_tools_to_toolset(
+        tools: list[str] | None,
+        server_id: str,
+        auth: UserAPIKeyAuth,
+    ) -> list[str] | None:
+        scope: Final = await MCPRequestHandler._toolset_scope(auth)
+        if scope is None:
+            return tools
+        scoped: Final = frozenset(scope.get(server_id, ()))
+        return sorted(scoped if tools is None else scoped & frozenset(tools))
 
     @staticmethod
     async def resolve_admitted_subject_servers(
@@ -2048,9 +2079,9 @@ class MCPRequestHandler:
                 continue
             tools = await MCPRequestHandler.get_allowed_tools_for_server(server_id, source, keyless_source=True)
             if tools is None:
-                return None
+                return await MCPRequestHandler._narrow_tools_to_toolset(None, server_id, auth)
             allowed.update(tools)
-        return sorted(allowed)
+        return await MCPRequestHandler._narrow_tools_to_toolset(sorted(allowed), server_id, auth)
 
     @staticmethod
     def _get_key_object_permission(
@@ -2067,6 +2098,16 @@ class MCPRequestHandler:
 
         return user_api_key_auth.object_permission
 
+    @staticmethod
+    async def team_object_permission(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None:
+        return await MCPRequestHandler._get_team_object_permission(user_api_key_auth)
+
+    @staticmethod
+    async def key_object_permission_hydrated(
+        user_api_key_auth: UserAPIKeyAuth,
+    ) -> LiteLLM_ObjectPermissionTable | None:
+        return await MCPRequestHandler._key_object_permission_hydrated(user_api_key_auth)
+
     @staticmethod
     async def _get_team_object_permission(
         user_api_key_auth: UserAPIKeyAuth | None = None,

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -915,11 +915,9 @@ if MCP_AVAILABLE:
     ) -> UserAPIKeyAuth:
         """The one credential this tools request acts as.
 
-        A toolset name narrows the caller's own credential to that toolset; otherwise a dashboard
-        session is swapped for its admitted subject. The two are mutually exclusive by construction,
-        which is why they share an owner: the admitted subject resolves per grant source and a team
-        source deliberately carries none of the caller's ``object_permission``, so a toolset
-        narrowing layered on top would evaporate on every team-granted server."""
+        A toolset name pins the acting principal to that toolset through ``_apply_toolset_scope``,
+        which itself swaps a dashboard session for its admitted subject; otherwise the swap happens
+        here so both shapes resolve as the same identity."""
         if not toolset_name:
             return await acting_user_auth(user_api_key_dict)
 

litellm/proxy/_experimental/mcp_server/server.py

@@ -66,7 +66,13 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import (
     get_route_relative_request_path,
     well_known_root_suffix,
 )
-from litellm.proxy._experimental.mcp_server.ui_session_utils import is_ui_session_credential
+from litellm.proxy._experimental.mcp_server.ui_session_utils import (
+    ActingUser,
+    GrantedToolsetIds,
+    acting_user_auth,
+    granted_toolset_ids,
+    is_ui_session_credential,
+)
 from litellm.proxy._experimental.mcp_server.utils import (
     LITELLM_MCP_SERVER_DESCRIPTION,
     LITELLM_MCP_SERVER_NAME,
@@ -1515,16 +1521,21 @@ if MCP_AVAILABLE:
     async def _apply_toolset_scope(
         user_api_key_auth: UserAPIKeyAuth,
         toolset_id: str,
+        acting_user: ActingUser = acting_user_auth,
+        granted: GrantedToolsetIds = granted_toolset_ids,
     ) -> UserAPIKeyAuth:
         """
-        Restrict a key's MCP permissions to a single toolset.
+        Pin a principal's MCP permissions to a single toolset for /toolset/{name}/mcp.
 
-        When a request arrives via /toolset/{name}/mcp we override the key's
-        object_permission so that only the toolset's tools are visible.
+        A virtual key (and an admin session) has its object_permission rewritten to
+        the toolset's servers and tools. A keyless subject resolves per grant source,
+        so a non-admin dashboard session first becomes its admitted user and the
+        toolset rides along as ``mcp_toolset_id``, which every source's grant is
+        intersected with; a team-granted toolset is served without the user's own
+        row capping it.
 
-        Raises HTTPException(403) if the key has an explicit toolset grant list
-        that does not include toolset_id (i.e. mcp_toolsets is set but empty,
-        or set to a list that omits this toolset).  Admin keys always pass.
+        Raises HTTPException(403) unless the principal holds toolset_id through one
+        of its grant sources. Admins always pass.
         """
         from litellm.proxy._types import LiteLLM_ObjectPermissionTable
         from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
@@ -1540,26 +1551,31 @@ if MCP_AVAILABLE:
                 detail="API key is scoped to no MCP servers; toolset access is denied.",
             )
 
-        # Access control: non-admin keys must have this toolset in their grant list.
-        # Use _user_has_admin_view so that PROXY_ADMIN_VIEW_ONLY is also treated as admin.
-        is_admin: Final = _user_has_admin_view(user_api_key_auth)
-        if not is_admin:
-            op: Final = user_api_key_auth.object_permission
-            granted: Final = getattr(op, "mcp_toolsets", None) if op else None
-            # granted=None → key has no explicit toolset grants → deny (same semantics as
-            # fetch_mcp_toolsets which returns [] for non-admin keys with no grants configured).
-            # granted=[] or list without toolset_id → also deny.
-            if granted is None or toolset_id not in granted:
+        acting: Final = await acting_user(user_api_key_auth)
+        is_admin: Final = _user_has_admin_view(acting)
+        if not is_admin and toolset_id not in await granted(acting):
+            raise HTTPException(
+                status_code=403,
+                detail=f"API key does not have access to toolset '{toolset_id}'.",
+            )
+        if _is_mcp_admitted_user_subject(acting):
+            resource_server_id: Final = acting.mcp_session_resource_server_id
+            if resource_server_id is not None and resource_server_id not in (
+                await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(
+                    toolset_ids=[toolset_id], requires_fresh_policy=acting.requires_fresh_policy
+                )
+            ):
                 raise HTTPException(
                     status_code=403,
                     detail=f"API key does not have access to toolset '{toolset_id}'.",
                 )
+            return acting.model_copy(update={"mcp_toolset_id": toolset_id})
 
         tool_permissions = await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(
             toolset_ids=[toolset_id]
         )
         server_ids: Final = list(tool_permissions.keys())
-        existing_op: Final = user_api_key_auth.object_permission
+        existing_op: Final = acting.object_permission
         if existing_op is not None:
             updated_op = existing_op.model_copy(
                 update={
@@ -1576,7 +1592,12 @@ if MCP_AVAILABLE:
                 mcp_servers=server_ids,
                 mcp_tool_permissions=tool_permissions,
             )
-        return user_api_key_auth.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id})
+        return acting.model_copy(update={"object_permission": updated_op, "mcp_toolset_id": toolset_id})
+
+    async def _toolset_server_ids(toolset_id: str) -> set[str]:
+        return set(
+            await operations.global_mcp_server_manager.resolve_toolset_tool_permissions(toolset_ids=[toolset_id])
+        )
 
     async def _raise_preemptive_401_for_unauthenticated_servers(
         scope: Scope,
@@ -2007,8 +2028,7 @@ if MCP_AVAILABLE:
             toolset_allowed_server_ids: set[str] | None = None
             if active_toolset_id and user_api_key_auth is not None:
                 user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id)
-                op: Final = user_api_key_auth.object_permission
-                toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set()
+                toolset_allowed_server_ids = await _toolset_server_ids(active_toolset_id)
 
             # https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
             # Must run after toolset scoping so the challenge set is derived
@@ -2355,8 +2375,7 @@ if MCP_AVAILABLE:
             toolset_allowed_server_ids: set[str] | None = None
             if active_toolset_id and user_api_key_auth is not None:
                 user_api_key_auth = await _apply_toolset_scope(user_api_key_auth, active_toolset_id)
-                op: Final = user_api_key_auth.object_permission
-                toolset_allowed_server_ids = set(op.mcp_servers or []) if op else set()
+                toolset_allowed_server_ids = await _toolset_server_ids(active_toolset_id)
 
             # https://datatracker.ietf.org/doc/html/rfc9728#name-www-authenticate-response
             # Must run after toolset scoping so the challenge set is derived

litellm/proxy/_experimental/mcp_server/ui_session_utils.py

@@ -2,14 +2,23 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
-from typing import Final
+import asyncio
+from collections.abc import Awaitable, Callable, Sequence
+from itertools import chain
+from typing import Final, TypeAlias
 
 from fastapi import HTTPException
 
 from litellm._logging import verbose_logger
 from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
+
+EffectiveAuthContexts: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[Sequence[UserAPIKeyAuth]]]
+TeamObjectPermission: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_ObjectPermissionTable | None]]
+OwnObjectPermission: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_ObjectPermissionTable | None]]
+AdmittedContext: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[UserAPIKeyAuth | None]]
+ActingUser: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[UserAPIKeyAuth]]
+GrantedToolsetIds: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[frozenset[str]]]
 
 
 def clone_user_api_key_auth_with_team(
@@ -109,10 +118,10 @@ async def acting_user_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth:
     both surfaces. An admin session keeps its operator view and any caller-passed credential is
     returned unchanged, never widened.
 
-    Do not combine this with a narrowing that rewrites a single credential's ``object_permission``
-    (toolset scope): the admitted subject resolves per grant source and a team source deliberately
-    carries none of the caller's own grants, so the narrowing would silently evaporate on every
-    team-granted server. A request carrying such a scope keeps the caller's own credential."""
+    A toolset narrowing is never applied to the admitted subject by rewriting its ``object_permission``:
+    it resolves per grant source and a team source deliberately carries none of the caller's own grants,
+    so the rewrite would evaporate on every team-granted server. The route pins ``mcp_toolset_id``
+    instead, which every source's grant is intersected with."""
 
     if not is_ui_session_credential(user_api_key_auth):
         return user_api_key_auth
@@ -153,3 +162,135 @@ async def can_access_mcp_server(
         if server_id in await allowed_servers(context):
             return True
     return False
+
+
+def _restricts_mcp(permission: LiteLLM_ObjectPermissionTable | None) -> bool:
+    return permission is not None and bool(
+        permission.mcp_servers
+        or permission.mcp_toolsets
+        or permission.mcp_tool_permissions
+        or permission.mcp_access_groups
+    )
+
+
+def is_keyless_mcp_subject(user_api_key_auth: UserAPIKeyAuth) -> bool:
+    """A principal with no virtual key to declare MCP access on: the dashboard's own session token or a
+    gateway-admitted user. Its grants are resolved per source, never through a key row."""
+
+    return is_ui_session_credential(user_api_key_auth) or user_api_key_auth.mcp_admitted_user_subject is True
+
+
+async def toolset_grant_contexts(
+    user_api_key_auth: UserAPIKeyAuth,
+    admitted_context: AdmittedContext = admitted_user_context,
+    admitted_sources: EffectiveAuthContexts | None = None,
+) -> Sequence[UserAPIKeyAuth]:
+    """The grant sources a toolset is looked up through. A virtual key is its own single source. A keyless
+    subject fans out exactly as the aggregate /mcp resolution does: its own user row plus every team whose
+    live roster still lists it, so a membership that only survives in the user's cached team list grants
+    nothing."""
+
+    if not is_keyless_mcp_subject(user_api_key_auth):
+        return (user_api_key_auth,)
+    from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+        MCPRequestHandler,
+    )
+
+    load_sources: Final = admitted_sources or MCPRequestHandler.admitted_subject_sources
+    acting: Final = await admitted_context(user_api_key_auth)
+    return tuple(await load_sources(acting if acting is not None else user_api_key_auth))
+
+
+async def _own_toolset_ids(
+    context: UserAPIKeyAuth,
+    load_own_permission: OwnObjectPermission,
+) -> Sequence[str] | None:
+    """The source's own toolsets, or None when it declares no MCP grant of its own. A source that names an
+    ``object_permission_id`` is a known restriction even when the row is unhydrated, unreadable or gone,
+    so it is loaded rather than read as unrestricted, and grants nothing when it cannot be read."""
+    if context.object_permission is None and not context.object_permission_id:
+        return None
+    try:
+        own: Final = await load_own_permission(context)
+    except Exception as exc:  # noqa: BLE001  # a named but unreadable own grant must deny, not widen to the team
+        verbose_logger.warning(
+            "MCP toolset grants: object permission %s unreadable, granting nothing through it: %s",
+            context.object_permission_id,
+            exc,
+        )
+        return ()
+    if own is None:
+        return ()
+    if not _restricts_mcp(own):
+        return None
+    return own.mcp_toolsets or ()
+
+
+async def _inherited_toolset_ids(
+    context: UserAPIKeyAuth,
+    load_team_permission: TeamObjectPermission,
+) -> Sequence[str]:
+    try:
+        team: Final = await load_team_permission(context)
+    except Exception as exc:  # noqa: BLE001  # an unreadable team grants nothing through this source and must not fail the caller's other sources
+        verbose_logger.warning(
+            "MCP toolset grants: team %s unreadable, inheriting nothing from it: %s",
+            context.team_id,
+            exc,
+        )
+        return ()
+    return () if team is None else (team.mcp_toolsets or ())
+
+
+async def _context_toolset_ids(
+    context: UserAPIKeyAuth,
+    inherits_team: bool,
+    load_team_permission: TeamObjectPermission,
+    load_own_permission: OwnObjectPermission,
+) -> Sequence[str]:
+    own: Final = await _own_toolset_ids(context, load_own_permission)
+    if own is not None:
+        return own
+    if not inherits_team or not context.team_id:
+        return ()
+    return await _inherited_toolset_ids(context, load_team_permission)
+
+
+async def granted_toolset_ids(
+    user_api_key_auth: UserAPIKeyAuth,
+    effective_contexts: EffectiveAuthContexts = toolset_grant_contexts,
+    team_object_permission: TeamObjectPermission | None = None,
+    require_key_access: bool | None = None,
+    own_object_permission: OwnObjectPermission | None = None,
+) -> frozenset[str]:
+    """Toolset ids the principal holds, resolved per grant source with the key/team rule the aggregate
+    /mcp listing applies: a source that declares any MCP grant of its own is scoped to its own toolsets and
+    never reads its team, one that declares none inherits its team's, except a virtual key under
+    ``require_key_mcp_access_defined``, which inherits nothing. A keyless subject's team sources always
+    inherit. A team that cannot be read contributes nothing while every other source still counts, and an
+    own grant that is named but cannot be read grants nothing. No grant anywhere yields the empty set."""
+    from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+        MCPRequestHandler,
+    )
+    from litellm.proxy.proxy_server import general_settings
+
+    require: Final = (
+        require_key_access
+        if require_key_access is not None
+        else bool(
+            general_settings.get(  # pyright: ignore[reportUnknownArgumentType]  # general_settings is an untyped dict; truthiness must match the /mcp path's read of this flag
+                "require_key_mcp_access_defined", False
+            )
+        )
+    )
+    inherits_team: Final = is_keyless_mcp_subject(user_api_key_auth) or not require
+    load_team_permission: Final = team_object_permission or MCPRequestHandler.team_object_permission
+    load_own_permission: Final = own_object_permission or MCPRequestHandler.key_object_permission_hydrated
+    contexts: Final = await effective_contexts(user_api_key_auth)
+    per_context: Final = await asyncio.gather(
+        *(
+            _context_toolset_ids(context, inherits_team, load_team_permission, load_own_permission)
+            for context in contexts
+        )
+    )
+    return frozenset(chain.from_iterable(per_context))

litellm/proxy/_lazy_features.py

@@ -3,19 +3,24 @@ Lazy registration for optional feature routers. Each LAZY_FEATURES entry
 imports its module only on the first request matching its path prefix,
 saving ~700 MB at idle for deployments that don't use these features.
 First hit pays the import cost (1-3 s for heavy modules); /openapi.json
-omits each feature's routes until the feature is warmed.
+omits each feature's routes until the feature is warmed. Setting
+LITELLM_DISABLE_LAZY_ROUTES registers every feature at worker startup
+instead, so the route table is complete before the first request.
 """
 
 import asyncio
 import importlib
-from collections.abc import Callable, Mapping, Sequence
+import os
+from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
 from collections.abc import Set as AbstractSet
+from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
+from functools import partial
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final
 
 from starlette.routing import BaseRoute, Match
-from starlette.types import ASGIApp, Receive, Scope, Send
+from starlette.types import ASGIApp, Lifespan, Receive, Scope, Send
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy.route_priority import hot_routes_first
@@ -428,57 +433,127 @@ def _in_registry_order(
 async def _force_load(app: "FastAPI", feat: LazyFeature, features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> bool:
     """Import + register a lazy feature exactly once per (app, module).
     Shared by the middleware and the /lazy/warm endpoint."""
+    async with _lazy_lock(app, feat.module_path):
+        if feat.module_path in _lazy_loaded(app):
+            return False
+        # Import on a thread (heavy modules take 1-3 s). register_fn
+        # mutates app.router.routes, so it stays on the loop thread.
+        imported: Final = asyncio.get_running_loop().run_in_executor(None, importlib.import_module, feat.module_path)
+        await asyncio.wait((imported,))
+        return _install(app, feat, imported.result, features)
+
+
+def _install(
+    app: "FastAPI", feat: LazyFeature, import_module: Callable[[], object], features: tuple[LazyFeature, ...]
+) -> bool:
+    try:
+        _register_feature(app, feat, import_module(), features)
+        return True
+    except Exception as exc:
+        _mark_failed(app, feat, exc)
+        return False
+
+
+def _lazy_loaded(app: "FastAPI") -> set[str]:
     if not hasattr(app.state, "lazy_loaded"):
-        app.state.lazy_loaded = set()
-        app.state.lazy_locks = {}
-    lock: Final = app.state.lazy_locks.setdefault(feat.module_path, asyncio.Lock())
-    async with lock:
-        if feat.module_path in app.state.lazy_loaded:
-            return False
-        try:
-            # Import on a thread (heavy modules take 1-3 s). register_fn
-            # mutates app.router.routes, so it stays on the loop thread.
-            loop: Final = asyncio.get_running_loop()
-            module: Final = await loop.run_in_executor(None, importlib.import_module, feat.module_path)
-            before: Final = len(app.router.routes)
-            feat.register_fn(app, module)
-            previous: Final[Mapping[str, tuple[BaseRoute, ...]]] = (
-                app.state.lazy_routes if hasattr(app.state, "lazy_routes") else MappingProxyType({})
-            )
-            lazy_routes: Final[Mapping[str, tuple[BaseRoute, ...]]] = MappingProxyType(
-                {**previous, feat.module_path: tuple(app.router.routes[before:])}
-            )
-            app.state.lazy_routes = lazy_routes  # rebind-ok: the app owns the record of which routes each feature added
-            app.router.routes[:] = hot_routes_first(  # rebind-ok: the app owns its route table
-                _in_registry_order(app.router.routes, lazy_routes, features, _lazy_slots(app))
-            )
-            app.state.lazy_loaded.add(feat.module_path)
-            app.openapi_schema = None
-            verbose_proxy_logger.info(
-                "Lazy-loaded optional feature %r (module: %s)",
-                feat.name,
-                feat.module_path,
-            )
-            return True
-        except Exception as exc:
-            # Mark loaded anyway so we don't retry on every request.
-            app.state.lazy_loaded.add(feat.module_path)
-            verbose_proxy_logger.warning(
-                "Failed to lazy-load optional feature %r (module: %s): %s. "
-                "This feature's endpoints will return 404 until restart.",
-                feat.name,
-                feat.module_path,
-                exc,
-            )
-            return False
+        app.state.lazy_loaded = set[str]()
+        app.state.lazy_locks = dict[str, asyncio.Lock]()
+    loaded: Final[set[str]] = app.state.lazy_loaded
+    return loaded
 
 
-def attach_lazy_features(app: "FastAPI") -> None:
-    app.include_router(_make_warmup_router(app))
-    app.add_middleware(LazyFeatureMiddleware, fastapi_app=app)
+def _lazy_lock(app: "FastAPI", module_path: str) -> asyncio.Lock:
+    if not hasattr(app.state, "lazy_locks"):
+        app.state.lazy_locks = dict[str, asyncio.Lock]()
+    locks: Final[dict[str, asyncio.Lock]] = app.state.lazy_locks
+    return locks.setdefault(module_path, asyncio.Lock())
 
 
-def _make_warmup_router(app: "FastAPI") -> "APIRouter":
+def _register_feature(app: "FastAPI", feat: LazyFeature, module: object, features: tuple[LazyFeature, ...]) -> None:
+    before: Final = len(app.router.routes)
+    feat.register_fn(app, module)
+    previous: Final[Mapping[str, tuple[BaseRoute, ...]]] = (
+        app.state.lazy_routes if hasattr(app.state, "lazy_routes") else MappingProxyType({})
+    )
+    lazy_routes: Final[Mapping[str, tuple[BaseRoute, ...]]] = MappingProxyType(
+        {**previous, feat.module_path: tuple(app.router.routes[before:])}
+    )
+    app.state.lazy_routes = lazy_routes  # rebind-ok: the app owns the record of which routes each feature added
+    app.router.routes[:] = hot_routes_first(  # rebind-ok: the app owns its route table
+        _in_registry_order(app.router.routes, lazy_routes, features, _lazy_slots(app))
+    )
+    _lazy_loaded(app).add(feat.module_path)
+    app.openapi_schema = None
+    verbose_proxy_logger.info(
+        "Lazy-loaded optional feature %r (module: %s)",
+        feat.name,
+        feat.module_path,
+    )
+
+
+def _mark_failed(app: "FastAPI", feat: LazyFeature, exc: Exception) -> None:
+    # Mark loaded anyway so we don't retry on every request.
+    _lazy_loaded(app).add(feat.module_path)
+    verbose_proxy_logger.warning(
+        "Failed to lazy-load optional feature %r (module: %s): %s. "
+        "This feature's endpoints will return 404 until restart.",
+        feat.name,
+        feat.module_path,
+        exc,
+    )
+
+
+def lazy_routes_disabled() -> bool:
+    return os.getenv("LITELLM_DISABLE_LAZY_ROUTES", "").lower() in ("1", "true", "yes", "on")
+
+
+def register_all_features(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> None:
+    """Register every feature router now, in registry order, so app.routes is
+    complete before the app serves its first request."""
+    for feat in features:
+        _install(app, feat, partial(importlib.import_module, feat.module_path), features)
+
+
+def attach_lazy_features(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> None:
+    if lazy_routes_disabled():
+        app.router.lifespan_context = _register_all_on_startup(app.router.lifespan_context, features)
+        return
+    app.include_router(_make_warmup_router(app, features))
+    app.add_middleware(LazyFeatureMiddleware, fastapi_app=app, features=features)
+
+
+def _register_all_on_startup(inner: "Lifespan[FastAPI]", features: tuple[LazyFeature, ...]) -> "Lifespan[FastAPI]":
+    """Registering at startup, once every route the app defines exists, lands the features
+    where lazy mode splices them: after every eager route (so /mcp/proxy, defined after
+    attach_lazy_features(), still beats the /mcp mount) and before LITELLM_WORKER_STARTUP_HOOKS
+    or an outer lifespan can filter the table. The inner lifespan then adds routes of its own
+    (config pass-through endpoints), so the table is put back in lazy mode's order once it is up."""
+
+    @asynccontextmanager
+    async def lifespan(app: "FastAPI") -> AsyncGenerator[None]:
+        register_all_features(app, features)
+        async with inner(app):
+            _restore_registry_order(app, features)
+            yield
+
+    return lifespan
+
+
+def _restore_registry_order(app: "FastAPI", features: tuple[LazyFeature, ...]) -> None:
+    present: Final = frozenset(id(route) for route in app.router.routes)
+    registered: Final[Mapping[str, tuple[BaseRoute, ...]]] = (
+        app.state.lazy_routes if hasattr(app.state, "lazy_routes") else MappingProxyType({})
+    )
+    still_routed: Final = MappingProxyType(
+        {module_path: tuple(r for r in routes if id(r) in present) for module_path, routes in registered.items()}
+    )
+    app.router.routes[:] = hot_routes_first(  # rebind-ok: the app owns its route table
+        _in_registry_order(app.router.routes, still_routed, features, _lazy_slots(app))
+    )
+    app.openapi_schema = None
+
+
+def _make_warmup_router(app: "FastAPI", features: tuple[LazyFeature, ...] = LAZY_FEATURES) -> "APIRouter":
     """POST /lazy/warm/{name}: load a feature and return its partial openapi
     so the Swagger plugin can merge in-place without a full /openapi.json refetch.
     Requires auth — anyone who can hit the proxy can already trigger the same
@@ -497,13 +572,13 @@ def _make_warmup_router(app: "FastAPI") -> "APIRouter":
         dependencies=[Depends(user_api_key_auth)],
     )
     async def warm(name: str):
-        feat: Final = next((f for f in LAZY_FEATURES if f.name == name), None)
+        feat: Final = next((f for f in features if f.name == name), None)
         if feat is None:
             raise HTTPException(404, f"unknown lazy feature: {name}")
         if feat.persistent_swagger_stub:
             return {"stub_path": None, "paths": {}, "components": {"schemas": {}}}
 
-        await _force_load(app, feat)
+        await _force_load(app, feat, features)
 
         feat_routes: Final = [r for r in app.routes if feat.matches(getattr(r, "path", ""))]
         full: Final = get_openapi(title=app.title, version=app.version, routes=feat_routes)
@@ -524,7 +599,7 @@ def _make_warmup_router(app: "FastAPI") -> "APIRouter":
 
 def loaded_lazy_modules(app: "FastAPI") -> frozenset[str]:
     """The set of lazy feature modules whose routers are actually registered
-    on this app (tracked by _force_load), empty before the middleware ever ran.
+    on this app (tracked by _install), empty until a feature loads or eager startup runs.
     sys.modules is the wrong signal: boot code imports several feature modules
     (mcp_management, cloudzero, vantage, config_overrides) without mounting
     their routers, and their stubs must still be injected."""
@@ -583,6 +658,6 @@ def lazy_tag_to_prefix() -> dict[str, str]:
     because /openapi.json already has full route info."""
     from litellm.proxy._lazy_openapi_snapshot import load_snapshot
 
-    if load_snapshot():
+    if lazy_routes_disabled() or load_snapshot():
         return {}
     return {feat.name: feat.path_prefixes[0] for feat in LAZY_FEATURES if not feat.persistent_swagger_stub}

litellm/proxy/db/prisma_client.py

@@ -984,6 +984,30 @@ class PrismaManager:
                 os.chdir(original_dir)
         return False
 
+    @staticmethod
+    def build_request_log_indexes() -> bool:
+        """Build the request-log indexes the migrations leave out and wait for them, for the
+        migration job (`--skip_server_startup`) after `setup_database` succeeds. False when
+        an index could not be built, so the job exits non-zero and is rerun."""
+        try:
+            from litellm_proxy_extras.utils import ProxyExtrasDBManager
+        except ImportError as e:
+            verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
+            return False
+        return ProxyExtrasDBManager.build_request_log_indexes()
+
+    @staticmethod
+    def start_request_log_index_build() -> None:
+        """Build the request-log indexes on a daemon thread, for a serving proxy that ran the
+        migrations itself (`DISABLE_SCHEMA_UPDATE` unset), so a long build never delays
+        readiness. A build that could not finish is logged and retried on the next boot."""
+        try:
+            from litellm_proxy_extras.utils import ProxyExtrasDBManager
+        except ImportError as e:
+            verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e)
+            return
+        ProxyExtrasDBManager.start_request_log_index_build()
+
 
 def should_update_prisma_schema(
     disable_updates: bool | str | None = None,

litellm/proxy/management_endpoints/mcp_management_endpoints.py

@@ -184,6 +184,7 @@ if MCP_AVAILABLE:
     from litellm.proxy._experimental.mcp_server.ui_session_utils import (
         admitted_user_context,
         build_effective_auth_contexts,
+        granted_toolset_ids,
         is_ui_session_credential,
     )
     from litellm.proxy._types import (
@@ -3332,18 +3333,15 @@ if MCP_AVAILABLE:
     ):
         """Return toolsets the calling key is allowed to access."""
         prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
-        is_admin: Final = _user_has_admin_view(user_api_key_dict)
-        op: Final = user_api_key_dict.object_permission
-        # mcp_toolsets=None or [] both mean "not restricted by toolsets".
-        # For admins: either value → no restriction → return all.
-        # For non-admins: either value → no toolsets explicitly granted → return nothing.
-        # (An admin whose DB row has mcp_toolsets=[] should still see all toolsets.)
-        raw_toolsets: Final = getattr(op, "mcp_toolsets", None) if op else None
-        if not raw_toolsets:
-            if is_admin:
+        if _user_has_admin_view(user_api_key_dict):
+            op: Final = user_api_key_dict.object_permission
+            if op is None or not op.mcp_toolsets:
                 return await list_mcp_toolsets(prisma_client)
+            return await list_mcp_toolsets(prisma_client, toolset_ids=op.mcp_toolsets)
+        granted: Final = await granted_toolset_ids(user_api_key_dict)
+        if not granted:
             return []
-        return await list_mcp_toolsets(prisma_client, toolset_ids=raw_toolsets)
+        return await list_mcp_toolsets(prisma_client, toolset_ids=sorted(granted))
 
     @router.get(
         "/toolset/{toolset_id}",
@@ -3355,15 +3353,13 @@ if MCP_AVAILABLE:
         user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
     ):
         prisma_client: Final = get_prisma_client_or_throw("Database not connected. Connect a database to your proxy")
-        # Non-admin keys may only fetch toolsets they've been explicitly granted.
-        if not _user_has_admin_view(user_api_key_dict):
-            op: Final = user_api_key_dict.object_permission
-            granted: Final = getattr(op, "mcp_toolsets", None) if op else None
-            if granted is None or toolset_id not in granted:
-                raise HTTPException(
-                    status_code=status.HTTP_403_FORBIDDEN,
-                    detail={"error": "API key does not have access to this toolset."},
-                )
+        if not _user_has_admin_view(user_api_key_dict) and toolset_id not in await granted_toolset_ids(
+            user_api_key_dict
+        ):
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail={"error": "API key does not have access to this toolset."},
+            )
         toolset: Final = await get_mcp_toolset(prisma_client, toolset_id)
         if toolset is None:
             raise HTTPException(

litellm/proxy/management_endpoints/team_endpoints.py

@@ -3466,6 +3466,12 @@ async def team_member_add(
         litellm_proxy_admin_name=litellm_proxy_admin_name,
     )
 
+    await delete_cache_team_object(
+        team_id=data.team_id,
+        team_alias=complete_team_data.team_alias,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
     await evict_and_broadcast(
         cache_keys=tuple(sorted(user.user_id for user in updated_users)),
         user_api_key_cache=user_api_key_cache,

litellm/proxy/proxy_cli.py

@@ -1410,10 +1410,15 @@ def run_server(
                             "LiteLLM versions contend for the same DB.\033[0m"
                         )
                     try:
-                        setup_ok: Final = PrismaManager.setup_database(
+                        migrated: Final = PrismaManager.setup_database(
                             use_migrate=not use_prisma_db_push,
                             use_v2_resolver=use_v2_resolver,
                         )
+                        setup_ok: Final = migrated and (
+                            not skip_server_startup or PrismaManager.build_request_log_indexes()
+                        )
+                        if migrated and not skip_server_startup:
+                            PrismaManager.start_request_log_index_build()
                     except RuntimeError as e:
                         # Raised on unrecoverable migration errors: the v2
                         # resolver's non-idempotent failures and permission

litellm/responses/mcp/litellm_proxy_mcp_handler.py

@@ -40,6 +40,7 @@ if TYPE_CHECKING:
     from mcp.types import CallToolResult
     from mcp.types import Tool as MCPTool
 
+    from litellm.proxy._experimental.mcp_server.ui_session_utils import GrantedToolsetIds
     from litellm.proxy._types import UserAPIKeyAuth
     from litellm.proxy.utils import ProxyLogging
 else:
@@ -223,7 +224,9 @@ class LiteLLM_Proxy_MCP_Handler:
                     mcp_servers=all_server_ids,
                     mcp_tool_permissions=tool_permissions,
                 )
-            return user_api_key_auth.model_copy(update={"object_permission": updated_op})
+            return user_api_key_auth.model_copy(
+                update={"object_permission": updated_op, "mcp_explicit_grants_only": True}
+            )
         except Exception as _e:
             verbose_logger.debug("Could not apply toolset permissions: %s", _e)
             return user_api_key_auth
@@ -237,6 +240,7 @@ class LiteLLM_Proxy_MCP_Handler:
         mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
         request_tags: list[str] | None = None,
         raw_headers: dict[str, str] | None = None,
+        granted_toolsets: "GrantedToolsetIds | None" = None,
     ) -> tuple[list[MCPTool], list[str]]:
         """
         Get available tools from the MCP server manager.
@@ -279,23 +283,19 @@ class LiteLLM_Proxy_MCP_Handler:
                     if prisma_client is not None:
                         toolset = await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name)
                         if toolset is not None:
-                            # Access control: only allow if the key explicitly grants this toolset.
                             if user_api_key_auth is not None:
+                                from litellm.proxy._experimental.mcp_server.ui_session_utils import (
+                                    granted_toolset_ids,
+                                )
                                 from litellm.proxy.management_endpoints.common_utils import (
                                     _user_has_admin_view,
                                 )
 
-                                is_admin = _user_has_admin_view(user_api_key_auth)
-                                if not is_admin:
-                                    op = user_api_key_auth.object_permission
-                                    granted = getattr(op, "mcp_toolsets", None) if op else None
-                                    # None means no grants configured → deny (consistent with
-                                    # fetch_mcp_toolsets which returns [] for unconfigured keys)
-                                    if granted is None or toolset.toolset_id not in granted:
-                                        verbose_logger.debug(
-                                            "Key does not have access to toolset '%s', skipping.", name
-                                        )
-                                        continue
+                                if not _user_has_admin_view(user_api_key_auth) and toolset.toolset_id not in (
+                                    await (granted_toolsets or granted_toolset_ids)(user_api_key_auth)
+                                ):
+                                    verbose_logger.debug("Key does not have access to toolset '%s', skipping.", name)
+                                    continue
                             resolved_toolset_ids.append(toolset.toolset_id)
                             # Don't add to resolved_mcp_servers — toolset scope
                             # restricts via object_permission, not server name filter.

litellm/sandbox/__init__.py

@@ -0,0 +1,27 @@
+"""litellm.sandbox: code-interpreter providers (see main.py) plus harness sandboxes.
+
+`sandbox.local(path)` and `sandbox.docker(image, ...)` re-export litellm.harness.sandbox.
+They resolve lazily so `import litellm` does not pull in litellm.harness.
+"""
+
+import importlib
+from typing import Final
+
+_HARNESS_SANDBOX_MODULE: Final = "litellm.harness.sandbox"
+_HARNESS_EXPORTS: Final = frozenset(
+    {
+        "local",
+        "docker",
+        "LocalSandbox",
+        "DockerSandbox",
+        "Sandbox",
+        "Process",
+        "CompletedRun",
+    }
+)
+
+
+def __getattr__(name: str) -> object:
+    if name in _HARNESS_EXPORTS:
+        return getattr(importlib.import_module(_HARNESS_SANDBOX_MODULE), name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

litellm/types/llms/custom_http.py

@@ -36,6 +36,7 @@ class httpxSpecialProvider(str, Enum):
     ModelCostMap = "model_cost_map"
     PasswordBreachCheck = "password_breach_check"
     ASGI = "asgi"
+    AgentHarness = "agent_harness"
 
 
 VerifyTypes = str | bool | ssl.SSLContext

litellm/utils.py

@@ -308,6 +308,7 @@ if TYPE_CHECKING:
         CachingHandlerResponse,
         LLMCachingHandler,
     )
+    from litellm.harness.types import Harness
     from litellm.integrations.custom_logger import CustomLogger
 
     # Type stubs for lazy-loaded functions and classes
@@ -385,6 +386,7 @@ if TYPE_CHECKING:
     from litellm.llms.base_llm.google_genai.transformation import (
         BaseGoogleGenAIGenerateContentConfig,
     )
+    from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
     from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
     from litellm.llms.base_llm.image_generation.transformation import (
         BaseImageGenerationConfig,
@@ -9883,6 +9885,37 @@ class ProviderConfigManager:
             return OpenSandboxSandboxConfig()
         return None
 
+    @staticmethod
+    def get_provider_harness_config(harness: Harness) -> BaseHarnessConfig | None:
+        """
+        Get the agent-harness configuration (Claude Code, Codex, OpenCode, Deep Agents).
+        """
+        from litellm.harness.types import Harness as _Harness
+
+        if harness == _Harness.CLAUDE_CODE:
+            from litellm.llms.claude_code.harness.transformation import (
+                ClaudeCodeHarnessConfig,
+            )
+
+            return ClaudeCodeHarnessConfig()
+        if harness == _Harness.CODEX:
+            from litellm.llms.codex.harness.transformation import CodexHarnessConfig
+
+            return CodexHarnessConfig()
+        if harness == _Harness.OPENCODE:
+            from litellm.llms.opencode.harness.transformation import (
+                OpenCodeHarnessConfig,
+            )
+
+            return OpenCodeHarnessConfig()
+        if harness == _Harness.DEEPAGENTS:
+            from litellm.llms.deepagents.harness.transformation import (
+                DeepAgentsHarnessConfig,
+            )
+
+            return DeepAgentsHarnessConfig()
+        return None
+
     @staticmethod
     def get_provider_text_to_speech_config(
         model: str,

migrations/run.py

@@ -2,7 +2,12 @@
 
 Runs `prisma migrate deploy` against the LiteLLM writer database using the
 recovery logic in `litellm_proxy_extras.ProxyExtrasDBManager.setup_database`
-(P3005 baseline + P3009/P3018 idempotent-error handling, retries, etc.).
+(P3005 baseline + P3009/P3018 idempotent-error handling, retries, etc.), then
+builds the request-log indexes the migrations leave out
+(`litellm_proxy_extras.request_log_indexes`), waiting for them. The job exits
+non-zero when an index could not be built so that it is rerun. A serving proxy
+that runs the migrations itself builds the same indexes in the background once
+it serves.
 
 Env vars:
   DATABASE_URL                  required unless it can be assembled at
@@ -23,10 +28,11 @@ Env vars:
 import os
 import sys
 
-from litellm.proxy.db.db_url_settings import DatabaseURLSettings
 from litellm_proxy_extras._logging import logger
 from litellm_proxy_extras.utils import ProxyExtrasDBManager, str_to_bool
 
+from litellm.proxy.db.db_url_settings import DatabaseURLSettings
+
 
 def main() -> int:
     # Assemble DATABASE_URL from the discrete DATABASE_* env vars, matching
@@ -52,7 +58,7 @@ def main() -> int:
         not use_db_push,
         use_v2,
     )
-    ok = ProxyExtrasDBManager.setup_database(
+    ok = ProxyExtrasDBManager.run_migration_job(
         use_migrate=not use_db_push,
         use_v2_resolver=use_v2,
     )

tests/code_coverage_tests/check_migrations_no_data_rewrites.py

@@ -7,15 +7,16 @@ anything whose cost scales with existing table size turns into downtime. A singl
 plus a doubled heap that plain autovacuum will not give back.
 
 What is banned is the row-rewriting DML behind that, not everything whose cost
-scales that way. A non-concurrent `CREATE INDEX`, an `ALTER COLUMN ... TYPE` that is
-not binary coercible, a volatile `DEFAULT` on a new column, a `CREATE TABLE ... AS
-SELECT` or `SELECT ... INTO` filling a new table from an existing one, the rename
-that pairs with one of those to swap a table out, and a `REFRESH MATERIALIZED VIEW`
-all read the whole table and all pass. That is deliberate: a rule wide enough to
-reach them fires on most ordinary migrations, and a marker everyone adds by reflex
-stops carrying information. The outage this was written for was a backfill.
+scales that way. A non-concurrent `CREATE INDEX` passes except on a request-log
+table, where it blocks writes until the build finishes. An `ALTER COLUMN ... TYPE`
+that is not binary coercible, a volatile `DEFAULT` on a new column, a `CREATE TABLE
+... AS SELECT` or `SELECT ... INTO` filling a new table from an existing one, the
+rename that pairs with one of those to swap a table out, and a `REFRESH MATERIALIZED
+VIEW` all read the whole table and all pass. That is deliberate: a rule wide enough
+to reach them fires on most ordinary migrations, and a marker everyone adds by
+reflex stops carrying information. The outage this was written for was a backfill.
 
-The one schema change banned outright is `ADD COLUMN ... DEFAULT` on a table in
+One column change banned outright is `ADD COLUMN ... DEFAULT` on a table in
 `REQUEST_LOG_TABLES`, the tables that hold a row per request. Postgres 11 stores such
 a default as metadata and touches no rows, but Postgres 10, which is supported,
 rewrites the whole heap and rebuilds every index under an `ACCESS EXCLUSIVE` lock,
@@ -23,6 +24,13 @@ which on a spend-log-sized table is the same outage as a backfill. Every other t
 is small enough that the rewrite is not worth a rule, and a column added to a log
 table without a default is still free on every version.
 
+An index on a request-log table cannot ship as a migration at all. A plain `CREATE
+INDEX` blocks writes to the table until the build finishes, and `CREATE INDEX
+CONCURRENTLY` is refused by Postgres on a partitioned parent, which LiteLLM_SpendLogs
+is wherever the operator ran db_scripts/partition_spend_logs.sql. The migration job builds
+those indexes after `migrate deploy`, concurrently and per partition, from the list in
+litellm_proxy_extras/request_log_indexes.py, so that list is where a new one goes.
+
 Flagged, per statement, by its leading keyword:
 
   UPDATE      rewrites every matching row, and `WHERE` does not bound the scan
@@ -44,6 +52,8 @@ Flagged, per statement, by its leading keyword:
               actions adds a column with a `DEFAULT`. An `ALTER COLUMN ... SET
               DEFAULT` written after the column exists changes metadata alone, so it
               passes, as does an `ADD CONSTRAINT`
+  CREATE      only a `CREATE [UNIQUE] INDEX` on a request-log table, concurrent or
+              not; the migration job builds those
 
 Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a
 statement's leading keyword, so they pass.
@@ -85,7 +95,9 @@ below line up with the statements they exempt.
 Add a column and let the application populate it, or run the rewrite as an opt-in
 batched job outside boot. When a rewrite is genuinely bounded and must ship inside
 the migration, put `-- data-migration-ok: <reason>` on the statement or on the line
-above it, naming what bounds it. The reason is required. A marker sharing a line
+above it, naming what bounds it. The reason is required. A marker never exempts a
+`CREATE INDEX` on a request-log table, since no bound makes that statement safe:
+the migration job is the only place such an index is built. A marker sharing a line
 with the statement it follows exempts that statement alone, so the next statement
 down is still checked rather than picking the marker up as its own. A marker on an
 `EXECUTE` or on the assignment feeding one covers the single-quoted SQL that
@@ -108,6 +120,7 @@ import sys
 from collections.abc import Iterator, Mapping
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Final
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
@@ -115,6 +128,9 @@ MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "
 GRANDFATHERED = frozenset(
     {
         "20250425182129_add_session_id",
+        "20250510142544_add_session_id_index_spend_logs",
+        "20260228100000_add_spend_logs_composite_index",
+        "20250326162113_baseline",
         "20260817000000_shadow_eval_multi_key",
         "20260818000000_add_spend_log_timestamps",
         "20260818224500_add_shadow_eval_stopped_by",
@@ -139,9 +155,7 @@ WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?<![<>!:=])=(?![=>])")
 PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$")
 QUALIFIER_GAP = re.compile(r"[\s.]*")
 EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE)
-DEFINES_A_ROUTINE = re.compile(
-    r"\bCREATE\b(?:\s+OR\s+REPLACE)?\s+(?:FUNCTION|PROCEDURE)\b", re.IGNORECASE
-)
+DEFINES_A_ROUTINE = re.compile(r"\bCREATE\b(?:\s+OR\s+REPLACE)?\s+(?:FUNCTION|PROCEDURE)\b", re.IGNORECASE)
 QUALIFIED_NAME = r"(?:\"[^\"]*\"|[A-Za-z_][A-Za-z0-9_$]*)"
 ROUTINE_NAME = re.compile(rf"\s*(?:{QUALIFIED_NAME}\s*\.\s*)?({QUALIFIED_NAME})")
 TABLE_NAME = ROUTINE_NAME
@@ -204,6 +218,12 @@ statement with the bound spelled out:
     -- data-migration-ok: <what bounds this>
     UPDATE ...
 
+An index on a request-log table is not a migration, and no marker exempts one. Declare it with `@@index` in
+schema.prisma and add it to REQUEST_LOG_INDEXES in
+litellm_proxy_extras/request_log_indexes.py under the name Prisma derives for it; the
+migration job builds it after `migrate deploy`, concurrently on a plain table and per
+partition on a partitioned one, which no single migration statement can do.
+
 On Postgres 10 an `ADD COLUMN ... DEFAULT` on a request-log table rewrites the table
 too. Add the column nullable with no default, then set the default in a separate
 `ALTER COLUMN ... SET DEFAULT`, which never touches existing rows.
@@ -215,10 +235,11 @@ class Violation:
     migration: str
     line: int
     keyword: str
+    consequence: str = "rewrites existing rows at boot"
 
     def render(self) -> str:
         location = f"{MIGRATIONS_DIR.relative_to(REPO_ROOT)}/{self.migration}/migration.sql"
-        return f"{location}:{self.line}: {self.keyword} rewrites existing rows at boot"
+        return f"{location}:{self.line}: {self.keyword} {self.consequence}"
 
 
 @dataclass(frozen=True, slots=True)
@@ -578,6 +599,35 @@ def rewrites_a_log_table(clause: str, region: str, base: int) -> str | None:
     return f"ADD COLUMN ... DEFAULT on {named.group(1)}"
 
 
+def builds_a_log_index(clause: str, region: str, base: int) -> str | None:
+    """The keyword to report when a `CREATE INDEX` targets a request-log table, concurrent or
+    not: a plain build blocks writes for its whole duration, and a concurrent one fails with
+    P3018 on a partitioned parent, so the migration job builds those instead."""
+    created: Final[re.Match[str] | None] = re.match(
+        r"\s*CREATE\s+(?:UNIQUE\s+)?INDEX\b(?:\s+CONCURRENTLY\b)?", clause, re.IGNORECASE
+    )
+    if created is None:
+        return None
+    on: Final[re.Match[str] | None] = re.search(r"\bON\b(?:\s+ONLY\b)?", clause[created.end() :], re.IGNORECASE)
+    if on is None:
+        return None
+    named: Final[re.Match[str] | None] = TABLE_NAME.match(
+        region, skip_comments(region, base + created.end() + on.end())
+    )
+    if named is None or named.group(1).strip('"') not in REQUEST_LOG_TABLES:
+        return None
+    return f"CREATE INDEX on {named.group(1)}"
+
+
+def consequence_of(found: str) -> str:
+    if found.startswith("CREATE INDEX"):
+        return (
+            "blocks writes until the build finishes, or fails on a partitioned table; "
+            "add it to REQUEST_LOG_INDEXES in litellm_proxy_extras/request_log_indexes.py instead"
+        )
+    return "rewrites existing rows at boot"
+
+
 def skip_comments(sql: str, start: int) -> int:
     index = start
     while index < len(sql):
@@ -746,8 +796,7 @@ def read_markers(sql: str) -> Markers:
     return Markers(
         sql,
         tuple(
-            Marker(match.start(), match.end(), alone_on_its_line(sql, match.start()))
-            for match in MARKER.finditer(sql)
+            Marker(match.start(), match.end(), alone_on_its_line(sql, match.start())) for match in MARKER.finditer(sql)
         ),
     )
 
@@ -760,9 +809,7 @@ def scan(sql: str, migration: str, markers: Markers) -> Iterator[Violation]:
     yield from scan_region(sql, sql, migration, markers, 0)
 
 
-def scan_region(
-    document: str, region: str, migration: str, markers: Markers, offset: int
-) -> Iterator[Violation]:
+def scan_region(document: str, region: str, migration: str, markers: Markers, offset: int) -> Iterator[Violation]:
     """Violations in one region of `document`, whose text begins at `offset`. Positions are
     always counted against the whole document, so a statement nested in a dollar-quoted body
     reports its real file line and lines up with the markers read from that file. A single-quoted
@@ -790,13 +837,18 @@ def scan_region(
                             offset + start,
                         )
 
-            keyword = offending_keyword(clause)
-            if exempt:
-                continue
-            found = keyword or rewrites_a_log_table(clause, region, base)
+            index = builds_a_log_index(clause, region, base)
+            found = (
+                index if exempt else offending_keyword(clause) or rewrites_a_log_table(clause, region, base) or index
+            )
             if found is None:
                 continue
-            yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), found)
+            yield Violation(
+                migration,
+                line_of(document, offset + keyword_start(clause, base)),
+                found,
+                consequence_of(found),
+            )
 
     for body in bodies:
         if not runs_when_applied(masked, region, bodies, runnable, identifiers, body):

tests/code_coverage_tests/check_provider_folders_documented.py

@@ -28,6 +28,11 @@ EXCLUDED_FOLDERS = {
     "pass_through",
     "openai_like",  # This is a generic handler, not a specific provider
     "aiohttp_openai",  # Internal implementation detail for async HTTP
+    # Agent-harness configs for litellm.agent(), not LLM providers; documented under docs/harness
+    "claude_code",
+    "codex",
+    "opencode",
+    "deepagents",
 }
 
 

tests/code_coverage_tests/recursive_detector.py

@@ -59,6 +59,8 @@ IGNORE_FUNCTIONS = [
     "_iter_fallback_targets",  # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap.
     "_mergeable_branch",  # max depth set (_MAX_SCHEMA_FLATTEN_DEPTH=32) plus a seen_refs cycle guard; passes the schema through untouched at the cap.
     "json_string_leaves",  # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); fails closed by raising at the cap so nothing goes unscanned.
+    "strict_json_schema",  # harness: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap.
+    "toml_value",  # harness/codex: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising OptionsMismatch at the cap.
     "with_json_string_leaves",  # transitively bounded: only runs on a tree json_string_leaves already walked under the cap.
     "json_unrewritable_labels",  # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks.
     "_flatten_form_field",  # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible).

tests/harness_e2e/conftest.py

@@ -0,0 +1,74 @@
+"""Fixtures for the litellm.agent() end-to-end tests.
+
+These run the real harness runtimes (claude, codex, opencode, deepagents) against a real
+LiteLLM AI Gateway, routed with the `litellm_proxy/` model prefix. They skip unless
+LITELLM_PROXY_API_BASE and LITELLM_PROXY_API_KEY are set. Model groups can be overridden
+per harness with HARNESS_E2E_MODEL_<HARNESS>.
+"""
+
+import importlib.util
+import os
+import shutil
+from collections.abc import Iterator
+from pathlib import Path
+
+import pytest
+
+from litellm import Harness
+
+GATEWAY_BASE = os.environ.get("LITELLM_PROXY_API_BASE", "").strip()
+GATEWAY_KEY = os.environ.get("LITELLM_PROXY_API_KEY", "").strip()
+
+DEFAULT_MODEL_GROUPS = {
+    Harness.CLAUDE_CODE: "claude-haiku-4-5-20251001",
+    Harness.CODEX: "bedrock_mantle/openai.gpt-5.4",
+    Harness.OPENCODE: "claude-haiku-4-5-20251001",
+    Harness.DEEPAGENTS: "claude-haiku-4-5-20251001",
+}
+
+BINARIES = {
+    Harness.CLAUDE_CODE: "claude",
+    Harness.CODEX: "codex",
+    Harness.OPENCODE: "opencode",
+}
+
+requires_gateway = pytest.mark.skipif(
+    not (GATEWAY_BASE and GATEWAY_KEY),
+    reason="LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY not set",
+)
+
+
+def model_for(harness: Harness) -> str:
+    """`litellm_proxy/<group>`: every model call goes through the gateway."""
+    override = os.environ.get(f"HARNESS_E2E_MODEL_{harness.name}", "").strip()
+    return f"litellm_proxy/{override or DEFAULT_MODEL_GROUPS[harness]}"
+
+
+def harness_available(harness: Harness) -> bool:
+    if harness is Harness.DEEPAGENTS:
+        return all(
+            importlib.util.find_spec(m) is not None
+            for m in ("deepagents", "langchain_litellm")
+        )
+    return shutil.which(BINARIES[harness]) is not None
+
+
+def harness_params() -> list:
+    return [
+        pytest.param(
+            h,
+            id=h.value,
+            marks=pytest.mark.skipif(
+                not harness_available(h), reason=f"{h.value} runtime not installed"
+            ),
+        )
+        for h in Harness
+    ]
+
+
+@pytest.fixture
+def workspace(tmp_path: Path) -> Iterator[Path]:
+    repo = tmp_path / "repo"
+    repo.mkdir()
+    (repo / "README.md").write_text("# demo\n")
+    yield repo

tests/harness_e2e/test_harness_e2e.py

@@ -0,0 +1,150 @@
+"""End-to-end: every harness, real runtime, real LiteLLM AI Gateway via litellm_proxy/."""
+
+from pathlib import Path
+
+import pytest
+from pydantic import BaseModel
+
+import litellm
+from litellm import Harness, sandbox
+from litellm.harness import (
+    CapabilityUnsupported,
+    Done,
+    FileChange,
+    State,
+    Text,
+    ToolCall,
+)
+
+from .conftest import harness_params, model_for, requires_gateway
+
+pytestmark = [requires_gateway]
+
+TURN_TIMEOUT = 300
+
+
+class Answer(BaseModel):
+    city: str
+    country: str
+
+
+@pytest.mark.parametrize("harness", harness_params())
+def test_agent_creates_file_and_reports_cost(harness: Harness, workspace: Path) -> None:
+    result = litellm.agent(
+        harness,
+        "Create a file named hello.txt whose entire content is the single word: hi",
+        sandbox=sandbox.local(workspace),
+        model=model_for(harness),
+        timeout=TURN_TIMEOUT,
+    )
+
+    assert result.stop_reason == "done", result.text
+    assert (workspace / "hello.txt").read_text().strip().lower() == "hi"
+    assert [f.path for f in result.files if f.kind == "created"] == ["hello.txt"]
+    assert result.usage.calls >= 1
+    assert result.usage.input_tokens > 0
+    assert result.cost >= 0
+
+
+@pytest.mark.parametrize("harness", harness_params())
+def test_agent_stream_event_order(harness: Harness, workspace: Path) -> None:
+    (workspace / "secret.txt").write_text("The secret word is ZEBRA.\n")
+    events = list(
+        litellm.agent(
+            harness,
+            "Read secret.txt and reply with just the secret word in it.",
+            sandbox=sandbox.local(workspace),
+            model=model_for(harness),
+            timeout=TURN_TIMEOUT,
+            stream=True,
+        )
+    )
+
+    assert isinstance(events[-1], Done)
+    assert sum(isinstance(e, Done) for e in events) == 1
+    assert any(isinstance(e, Text) for e in events)
+    assert any(isinstance(e, ToolCall) for e in events)
+    assert "zebra" in events[-1].result.text.lower()
+
+
+@pytest.mark.parametrize("harness", harness_params())
+def test_agent_structured_output(harness: Harness, workspace: Path) -> None:
+    result = litellm.agent(
+        harness,
+        "What is the capital of France? Do not use any tools.",
+        sandbox=sandbox.local(workspace),
+        model=model_for(harness),
+        output=Answer,
+        permissions="read-only",
+        timeout=TURN_TIMEOUT,
+    )
+
+    assert isinstance(result.output, Answer)
+    assert result.output.city.lower() == "paris"
+
+
+@pytest.mark.parametrize("harness", harness_params())
+def test_agent_session_remembers_previous_turn(
+    harness: Harness, workspace: Path
+) -> None:
+    with litellm.agent_session(
+        harness,
+        sandbox=sandbox.local(workspace),
+        model=model_for(harness),
+        timeout=TURN_TIMEOUT,
+    ) as s:
+        s.run("Remember this code word: PELICAN. Reply with just OK.")
+        second = s.run(
+            "What code word did I ask you to remember? Reply with just the word."
+        )
+        assert "pelican" in second.text.lower()
+        assert s.cost >= second.cost
+
+
+@pytest.mark.parametrize("harness", harness_params())
+def test_agent_detach_and_resume(harness: Harness, workspace: Path) -> None:
+    box = sandbox.local(workspace)
+    s = litellm.agent_session(
+        harness, sandbox=box, model=model_for(harness), timeout=TURN_TIMEOUT
+    )
+    s.run("Remember this number: 4817. Reply with just OK.")
+    raw = s.detach().dumps()
+
+    with litellm.agent_resume(
+        State.loads(raw), sandbox=box, model=model_for(harness)
+    ) as resumed:
+        r = resumed.run(
+            "What number did I ask you to remember? Reply with just the number."
+        )
+    assert "4817" in r.text
+
+
+@pytest.mark.parametrize("harness", harness_params())
+def test_agent_read_only_blocks_writes(harness: Harness, workspace: Path) -> None:
+    result = litellm.agent(
+        harness,
+        "Create a file named blocked.txt containing x. If you cannot, just say you cannot.",
+        sandbox=sandbox.local(workspace),
+        model=model_for(harness),
+        permissions="read-only",
+        timeout=TURN_TIMEOUT,
+    )
+
+    assert not (workspace / "blocked.txt").exists()
+    assert not [f for f in result.files if isinstance(f, FileChange)]
+
+
+def test_string_harness_rejected(workspace: Path) -> None:
+    with pytest.raises(TypeError, match=r"Harness\.CODEX"):
+        litellm.agent("codex", "hi", sandbox=sandbox.local(workspace))  # type: ignore[arg-type]
+
+
+def test_capability_checked_before_start(workspace: Path) -> None:
+    with pytest.raises(CapabilityUnsupported):
+        litellm.agent(
+            Harness.CODEX,
+            "hi",
+            sandbox=sandbox.local(workspace),
+            model=model_for(Harness.CODEX),
+            disable_tools=["bash"],
+        )

tests/integration/_support/mcp_grants.py

@@ -51,12 +51,12 @@ def delete_toolset(gateway: Gateway, identity: str) -> None:
     assert response.status_code in (200, 202, 204), response.text
 
 
-def create_toolset(scenario: Scenario, tools: tuple[tuple[str, str], ...]) -> str:
+def create_toolset(scenario: Scenario, tools: tuple[tuple[str, str], ...], toolset_name: str | None = None) -> str:
     response: Final = scenario.gateway.request(
         "POST",
         "/v1/mcp/toolset",
         {
-            "toolset_name": f"integration-{uuid.uuid4().hex[:10]}",
+            "toolset_name": toolset_name or f"integration-{uuid.uuid4().hex[:10]}",
             "tools": [{"server_id": server_id, "tool_name": tool} for server_id, tool in tools],
         },
     )

tests/integration/_support/process.py

@@ -16,6 +16,10 @@ import httpx
 import psutil
 from integration._support.client import GATEWAY_LIMITS, Gateway
 
+DB_PUSH: Final = ("--use_prisma_db_push",)
+MIGRATE_DEPLOY: Final = ()
+LEGACY_MIGRATE_DEPLOY: Final = ("--use_legacy_migration_resolver",)
+
 
 def proxy_database_environment() -> Mapping[str, str]:
     writer: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL", "")
@@ -73,9 +77,16 @@ def owned_proxy(
     config: Path | None = None,
     remove_environment: tuple[str, ...] = (),
     workers: int = 1,
+    database_setup: tuple[str, ...] = DB_PUSH,
 ) -> Iterator[Gateway]:
     with owned_proxy_process(
-        gateway, directory, overrides, config=config, remove_environment=remove_environment, workers=workers
+        gateway,
+        directory,
+        overrides,
+        config=config,
+        remove_environment=remove_environment,
+        workers=workers,
+        database_setup=database_setup,
     ) as owned:
         yield owned.gateway
 
@@ -133,7 +144,7 @@ def _lost_port_race(launch: _Launch) -> bool:
 
 def _wait_until_ready(launch: _Launch) -> None:
     with httpx.Client(base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False) as client:
-        deadline: Final = time.monotonic() + 70
+        deadline: Final = time.monotonic() + float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS", "70"))
         while launch.process.poll() is None:
             try:
                 if client.get("/health/readiness", timeout=2).status_code == 200:
@@ -171,6 +182,7 @@ def owned_proxy_process(
     config: Path | None = None,
     remove_environment: tuple[str, ...] = (),
     workers: int = 1,
+    database_setup: tuple[str, ...] = DB_PUSH,
 ) -> Iterator[OwnedProxy]:
     root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
     environment: Final = {
@@ -196,7 +208,7 @@ def owned_proxy_process(
         "127.0.0.1",
         "--num_workers",
         str(workers),
-        "--use_prisma_db_push",
+        *database_setup,
         "--enforce_prisma_migration_check",
     )
     launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS)

tests/integration/configuration/test_lazy_routes_flag.py

@@ -0,0 +1,605 @@
+"""Route table contract for the LITELLM_DISABLE_LAZY_ROUTES startup flag.
+
+By default optional feature routers (``LAZY_FEATURES``) are registered on the first
+request to their path prefix, so an operator inspecting the route table right after
+boot cannot see or gate them. With the flag set every feature is registered at worker
+startup, so ``GET /routes`` lists them before any feature request is served and the
+first feature request changes nothing.
+"""
+
+import asyncio
+import json
+import os
+import re
+import uuid
+from collections.abc import Iterator, Mapping
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Final
+
+import anthropic
+import httpx
+import openai
+import psutil
+import pytest
+import yaml
+from pydantic import JsonValue, TypeAdapter
+
+from litellm.proxy._lazy_features import LAZY_FEATURES, LazyFeature
+from tests.integration._support.client import Gateway, eventually, object_value, string_value
+from tests.integration._support.mcp import McpPeer, call_tool, echo_tool, scripted_peer, tool_calls, tool_names
+from tests.integration._support.process import OwnedProxy, owned_proxy_process
+from tests.integration._support.wire import Reply, Request, Wire, wire_server
+
+TICKET_FEATURES: Final = ("mcp_management", "mcp_byok_oauth")
+FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES"
+WARMUP_ROUTE: Final = "/lazy/warm/{name}"
+MCP_WARM_PATH: Final = "/mcp/enabled"
+MARKER: Final = re.compile(rb"lazyroutes-[0-9a-f]{32}")
+FAILED_FEATURE: Final = re.compile(r"Failed to lazy-load optional feature '([a-z_]+)'")
+JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+HOOK_MODULE: Final = "lazy_routes_route_filter_hook"
+HOOK_SOURCE: Final = """from litellm.proxy.proxy_server import app
+
+
+def drop_mcp_routes() -> None:
+    app.router.routes[:] = [
+        route for route in app.router.routes if not getattr(route, "path", "").startswith(("/mcp", "/v1/mcp"))
+    ]
+"""
+
+
+def _paths(candidate: Gateway) -> tuple[str, ...]:
+    routes: Final = candidate.get("/routes")["routes"]
+    assert isinstance(routes, list), routes
+    return tuple(string_value(object_value(route)["path"]) for route in routes)
+
+
+def _routed_features(candidate: Gateway) -> Mapping[str, tuple[str, ...]]:
+    paths: Final = _paths(candidate)
+    return {feature.name: tuple(path for path in paths if feature.matches(path)) for feature in LAZY_FEATURES}
+
+
+def _mcp_paths(candidate: Gateway) -> tuple[str, ...]:
+    return tuple(path for path in _paths(candidate) if path.startswith(("/mcp", "/v1/mcp")))
+
+
+def _route_filter_hook(directory: Path) -> Mapping[str, str]:
+    (directory / f"{HOOK_MODULE}.py").write_text(HOOK_SOURCE)
+    search_path: Final = (str(directory), os.environ.get("PYTHONPATH", ""))
+    return {
+        "PYTHONPATH": os.pathsep.join(entry for entry in search_path if entry),
+        "LITELLM_WORKER_STARTUP_HOOKS": f"{HOOK_MODULE}:drop_mcp_routes",
+    }
+
+
+def test_lazy_routes_are_absent_from_the_route_table_until_first_request(gateway: Gateway, tmp_path: Path) -> None:
+    with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as owned:
+        at_boot: Final = _routed_features(owned.gateway)
+        assert {name: at_boot[name] for name in TICKET_FEATURES} == {name: () for name in TICKET_FEATURES}, at_boot
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+        after_first_request: Final = _routed_features(owned.gateway)
+        assert after_first_request["mcp_management"] != (), "first request did not register the router"
+        assert after_first_request["mcp_byok_oauth"] == (), "only the requested feature is mounted"
+
+
+@pytest.mark.parametrize("workers", (1, 4))
+def test_disable_lazy_routes_flag_registers_every_feature_at_startup(
+    gateway: Gateway, tmp_path: Path, workers: int
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}, workers=workers) as owned:
+        at_boot: Final = tuple(_routed_features(owned.gateway) for _ in range(2 * workers))
+        unregistered: Final = sorted(name for name, paths in at_boot[0].items() if not paths)
+        assert unregistered == [], f"features still missing from /routes at startup: {unregistered}"
+        assert all(table == at_boot[0] for table in at_boot), "workers disagree on the route table"
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+        assert _routed_features(owned.gateway) == at_boot[0], "first feature request changed the route table"
+
+
+def test_startup_hook_cannot_remove_lazy_routes_that_register_after_it_ran(gateway: Gateway, tmp_path: Path) -> None:
+    with owned_proxy_process(gateway, tmp_path, _route_filter_hook(tmp_path), remove_environment=(FLAG,)) as owned:
+        assert _mcp_paths(owned.gateway) == (), "hook should have removed the routes registered before it ran"
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+        assert _mcp_paths(owned.gateway) != (), "first request should have registered the routes the hook never saw"
+
+
+def test_disable_lazy_routes_flag_lets_a_startup_hook_remove_optional_routes_for_good(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    overrides: Final = {**_route_filter_hook(tmp_path), FLAG: "true"}
+    with owned_proxy_process(gateway, tmp_path, overrides) as owned:
+        assert _mcp_paths(owned.gateway) == (), "hook should have seen and removed every MCP route"
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 404, listing.text
+        mounted: Final = owned.gateway.request("POST", "/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
+        assert mounted.status_code == 404, mounted.text
+        guardrails: Final = owned.gateway.request("GET", "/guardrails/list")
+        assert guardrails.status_code == 200, guardrails.text
+        assert _mcp_paths(owned.gateway) == (), "a feature request re-registered routes the hook removed"
+
+
+def _openapi_paths(candidate: Gateway) -> Mapping[str, tuple[str, ...]]:
+    paths: Final = object_value(candidate.get("/openapi.json")["paths"])
+    return {path: tuple(sorted(object_value(operations))) for path, operations in paths.items()}
+
+
+def _published(feature: LazyFeature, paths: Mapping[str, tuple[str, ...]]) -> bool:
+    return any(feature.matches(path) for path in paths)
+
+
+def _warm_every_feature(candidate: Gateway) -> None:
+    warmed: Final = tuple(
+        (feature.name, candidate.request("POST", f"/lazy/warm/{feature.name}")) for feature in LAZY_FEATURES
+    )
+    cold: Final = [
+        (name, response.status_code, response.text) for name, response in warmed if response.status_code != 200
+    ]
+    assert cold == [], cold
+    enabled: Final = candidate.request("GET", MCP_WARM_PATH)
+    assert enabled.status_code == 200, enabled.text
+    unregistered: Final = sorted(name for name, paths in _routed_features(candidate).items() if not paths)
+    assert unregistered == [], f"features still missing after warming every one of them: {unregistered}"
+
+
+def _shadowed_dependency(directory: Path) -> Mapping[str, str]:
+    package: Final = directory / "shadow" / "RestrictedPython"
+    package.mkdir(parents=True)
+    (package / "__init__.py").write_text('raise ImportError("shadowed by the lazy routes audit")\n')
+    search_path: Final = (str(package.parent), os.environ.get("PYTHONPATH", ""))
+    return {"PYTHONPATH": os.pathsep.join(entry for entry in search_path if entry)}
+
+
+def _failed_features(owned: OwnedProxy) -> frozenset[str]:
+    return frozenset(FAILED_FEATURE.findall(owned.log.read_text()))
+
+
+def _marker() -> str:
+    return "lazyroutes-" + uuid.uuid4().hex
+
+
+def _chat_reply(identity: str, stream: bool) -> Reply:
+    if not stream:
+        return Reply(
+            body=json.dumps(
+                {
+                    "id": identity,
+                    "object": "chat.completion",
+                    "created": 1,
+                    "model": "gpt-4o-mini",
+                    "choices": [
+                        {"index": 0, "message": {"role": "assistant", "content": "lazy ok"}, "finish_reason": "stop"}
+                    ],
+                    "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9},
+                }
+            ).encode()
+        )
+    chunk: Final[dict[str, JsonValue]] = {
+        "id": identity,
+        "object": "chat.completion.chunk",
+        "created": 1,
+        "model": "gpt-4o-mini",
+    }
+    deltas: Final[tuple[dict[str, JsonValue], ...]] = (
+        {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "lazy"}}]},
+        {**chunk, "choices": [{"index": 0, "delta": {"content": " ok"}, "finish_reason": "stop"}]},
+        {**chunk, "choices": [], "usage": {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9}},
+    )
+    return Reply(
+        content_type="text/event-stream",
+        chunks=(*(b"data: " + json.dumps(delta).encode() + b"\n\n" for delta in deltas), b"data: [DONE]\n\n"),
+    )
+
+
+def _responses_reply(identity: str, stream: bool) -> Reply:
+    response: Final[dict[str, JsonValue]] = {
+        "id": identity,
+        "object": "response",
+        "created_at": 1,
+        "status": "completed",
+        "model": "gpt-4o-mini",
+        "output": [
+            {
+                "id": "msg_" + identity,
+                "type": "message",
+                "role": "assistant",
+                "status": "completed",
+                "content": [{"type": "output_text", "text": "lazy ok", "annotations": []}],
+            }
+        ],
+        "usage": {"input_tokens": 7, "output_tokens": 2, "total_tokens": 9},
+    }
+    if not stream:
+        return Reply(body=json.dumps(response).encode())
+    events: Final[tuple[dict[str, JsonValue], ...]] = (
+        {"type": "response.created", "sequence_number": 0, "response": {**response, "status": "in_progress"}},
+        {
+            "type": "response.output_text.delta",
+            "sequence_number": 1,
+            "item_id": "msg_" + identity,
+            "output_index": 0,
+            "content_index": 0,
+            "delta": "lazy ok",
+        },
+        {"type": "response.completed", "sequence_number": 2, "response": response},
+    )
+    return Reply(
+        content_type="text/event-stream",
+        chunks=tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events),
+    )
+
+
+def _upstream(request: Request) -> Reply:
+    found: Final = MARKER.search(request.body)
+    if found is None:
+        return Reply(status=404, body=b'{"error":"no marker"}')
+    marker: Final = found.group(0).decode()
+    stream: Final = object_value(JSON.validate_json(request.body)).get("stream") is True
+    if request.target.endswith("/responses"):
+        return _responses_reply(f"resp_{marker}", stream)
+    return _chat_reply(f"chatcmpl-{marker}", stream)
+
+
+@pytest.fixture(scope="module")
+def provider() -> Iterator[Wire]:
+    with wire_server(_upstream) as wire:
+        yield wire
+
+
+async def _stream_chat(base_url: str, key: str, model: str, marker: str) -> tuple[frozenset[str], str]:
+    client: Final = openai.AsyncOpenAI(base_url=base_url + "/v1", api_key=key, max_retries=0)
+    stream: Final = await client.chat.completions.create(
+        model=model, messages=[{"role": "user", "content": marker}], stream=True
+    )
+    chunks: Final = [chunk async for chunk in stream]
+    text: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
+    return frozenset(chunk.id for chunk in chunks), text
+
+
+async def _stream_message(base_url: str, key: str, model: str, marker: str) -> str:
+    client: Final = anthropic.AsyncAnthropic(base_url=base_url, api_key=key, max_retries=0)
+    async with client.messages.stream(
+        model=model, max_tokens=16, messages=[{"role": "user", "content": marker}]
+    ) as stream:
+        return "".join([text async for text in stream.text_stream])
+
+
+def _status(candidate: Gateway, path: str) -> int:
+    return candidate.request("GET", path).status_code
+
+
+def _workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]:
+    return tuple(child for child in psutil.Process(owned.process.pid).children() if _is_worker(child))
+
+
+def _is_worker(child: psutil.Process) -> bool:
+    try:
+        return "spawn_main" in " ".join(child.cmdline()) and child.status() != psutil.STATUS_ZOMBIE
+    except psutil.Error:
+        return False
+
+
+@pytest.mark.parametrize("spelling", ("1", "Yes", "ON"))
+def test_every_truthy_spelling_of_the_flag_registers_the_ticket_features_at_startup(
+    gateway: Gateway, tmp_path: Path, spelling: str
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {FLAG: spelling}) as owned:
+        at_boot: Final = _routed_features(owned.gateway)
+        assert all(at_boot[name] for name in TICKET_FEATURES), {name: at_boot[name] for name in TICKET_FEATURES}
+
+
+@pytest.mark.parametrize(
+    "spelling", ("", "0", "off", "maybe", "x" * 5000), ids=("empty", "zero", "off", "unknown-word", "five-kilobytes")
+)
+def test_a_falsey_or_unknown_flag_value_keeps_the_default_lazy_registration(
+    gateway: Gateway, tmp_path: Path, spelling: str
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {FLAG: spelling}) as owned:
+        at_boot: Final = _routed_features(owned.gateway)
+        assert {name: at_boot[name] for name in TICKET_FEATURES} == {name: () for name in TICKET_FEATURES}, at_boot
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+        assert _routed_features(owned.gateway)["mcp_management"] != (), "first request did not register the router"
+
+
+def test_disable_lazy_routes_flag_publishes_the_live_route_table_in_openapi_before_any_feature_request(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as owned:
+        at_boot: Final = _openapi_paths(owned.gateway)
+        unpublished: Final = sorted(
+            feature.name
+            for feature in LAZY_FEATURES
+            if feature.name in TICKET_FEATURES and not _published(feature, at_boot)
+        )
+        assert unpublished == [], f"ticket features missing from /openapi.json at startup: {unpublished}"
+        assert "get" in at_boot["/v1/mcp/server"], at_boot["/v1/mcp/server"]
+        assert WARMUP_ROUTE not in at_boot
+        stranger: Final = owned.gateway.request("GET", "/v1/mcp/server", key="sk-not-a-real-key")
+        assert stranger.status_code == 401, stranger.text
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+        assert _openapi_paths(owned.gateway) == at_boot, "first feature request changed /openapi.json"
+
+
+def test_disable_lazy_routes_flag_removes_the_warmup_route(gateway: Gateway, tmp_path: Path) -> None:
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as owned:
+        assert WARMUP_ROUTE not in _paths(owned.gateway)
+        warmed: Final = owned.gateway.request("POST", "/lazy/warm/mcp_management")
+        assert warmed.status_code == 404, warmed.text
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+
+
+def test_the_warmup_route_registers_a_feature_on_demand_by_default(gateway: Gateway, tmp_path: Path) -> None:
+    with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as owned:
+        assert WARMUP_ROUTE in _paths(owned.gateway)
+        warmed: Final = owned.gateway.request("POST", "/lazy/warm/mcp_management")
+        assert warmed.status_code == 200, warmed.text
+        assert "/v1/mcp/server" in object_value(object_value(JSON.validate_json(warmed.content))["paths"])
+        assert _routed_features(owned.gateway)["mcp_management"] != (), "warmup did not register the router"
+
+
+def test_disable_lazy_routes_flag_keeps_the_fixed_mcp_proxy_route_ahead_of_the_mcp_mount(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as lazy:
+        control: Final = lazy.gateway.request("POST", "/mcp/proxy", {})
+        assert control.status_code == 400, control.text
+        warmed: Final = _paths(lazy.gateway)
+        assert warmed.count("/mcp") == 2, "expected the fixed /mcp route and the /mcp mount"
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as eager:
+        at_boot: Final = _paths(eager.gateway)
+        assert at_boot.count("/mcp") == 2, "expected the fixed /mcp route and the /mcp mount at startup"
+        mount: Final = max(index for index, path in enumerate(at_boot) if path == "/mcp")
+        assert at_boot.index("/mcp/proxy") < mount, "the /mcp mount shadows /mcp/proxy"
+        proxied: Final = eager.gateway.request("POST", "/mcp/proxy", {})
+        assert (proxied.status_code, proxied.text) == (control.status_code, control.text)
+
+
+def test_disable_lazy_routes_flag_matches_the_fully_warmed_lazy_route_table_and_openapi(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {}, remove_environment=(FLAG,)) as lazy:
+        _warm_every_feature(lazy.gateway)
+        warmed_paths: Final = tuple(path for path in _paths(lazy.gateway) if path != WARMUP_ROUTE)
+        warmed_openapi: Final = _openapi_paths(lazy.gateway)
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as eager:
+        assert _paths(eager.gateway) == warmed_paths
+        assert _openapi_paths(eager.gateway) == warmed_openapi
+
+
+def test_disable_lazy_routes_flag_keeps_registering_after_an_optional_dependency_fails_to_import(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {**_shadowed_dependency(tmp_path), FLAG: "true"}) as owned:
+        at_boot: Final = _routed_features(owned.gateway)
+        unregistered: Final = frozenset(name for name, paths in at_boot.items() if not paths)
+        failed: Final = _failed_features(owned)
+        assert "guardrails" in failed, owned.log.read_text()
+        assert unregistered == failed, (sorted(unregistered), sorted(failed))
+        guardrails: Final = owned.gateway.request("GET", "/guardrails/list")
+        assert guardrails.status_code == 404, guardrails.text
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+        stores: Final = owned.gateway.request("GET", "/vector_store/list")
+        assert stores.status_code == 200, stores.text
+        assert _routed_features(owned.gateway) == at_boot, "feature requests changed the route table"
+
+
+def test_a_broken_optional_dependency_only_404s_its_own_feature_by_default(gateway: Gateway, tmp_path: Path) -> None:
+    with owned_proxy_process(gateway, tmp_path, _shadowed_dependency(tmp_path), remove_environment=(FLAG,)) as owned:
+        guardrails: Final = owned.gateway.request("GET", "/guardrails/list")
+        assert guardrails.status_code == 404, guardrails.text
+        assert "guardrails" in _failed_features(owned), owned.log.read_text()
+        listing: Final = owned.gateway.request("GET", "/v1/mcp/server")
+        assert listing.status_code == 200, listing.text
+        assert _routed_features(owned.gateway)["mcp_management"] != ()
+
+
+def test_disable_lazy_routes_flag_leaves_the_completion_endpoints_serving_every_client(
+    gateway: Gateway, tmp_path: Path, provider: Wire
+) -> None:
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as owned, owned.gateway.scenario() as scenario:
+        model: Final = scenario.model(api_base=provider.url + "/v1")
+        base_url: Final = str(owned.gateway.client.base_url)
+        key: Final = owned.gateway.key
+        markers: Final = tuple(_marker() for _ in range(6))
+
+        completion: Final = openai.OpenAI(
+            base_url=base_url + "/v1", api_key=key, max_retries=0
+        ).chat.completions.create(model=model, messages=[{"role": "user", "content": markers[0]}])
+        assert (completion.id, completion.choices[0].message.content) == (f"chatcmpl-{markers[0]}", "lazy ok")
+
+        assert asyncio.run(_stream_chat(base_url, key, model, markers[1])) == (
+            frozenset({f"chatcmpl-{markers[1]}"}),
+            "lazy ok",
+        )
+
+        message: Final = anthropic.Anthropic(base_url=base_url, api_key=key, max_retries=0).messages.create(
+            model=model, max_tokens=16, messages=[{"role": "user", "content": markers[2]}]
+        )
+        assert [block.text for block in message.content if block.type == "text"] == ["lazy ok"]
+
+        assert asyncio.run(_stream_message(base_url, key, model, markers[3])) == "lazy ok"
+
+        responded: Final = owned.gateway.request("POST", "/v1/responses", {"model": model, "input": markers[4]})
+        assert responded.status_code == 200, responded.text
+        response: Final = object_value(JSON.validate_json(responded.content))
+        assert (response["status"], response["object"]) == ("completed", "response"), responded.text
+        assert "lazy ok" in responded.text, responded.text
+
+        streamed: Final = owned.gateway.request(
+            "POST", "/v1/responses", {"model": model, "input": markers[5], "stream": True}
+        )
+        assert streamed.status_code == 200, streamed.text
+        assert "response.completed" in streamed.text and "lazy ok" in streamed.text, streamed.text
+
+        reached: Final = tuple(request.target for request in provider.drain() if MARKER.search(request.body))
+        assert len(reached) == 6, reached
+
+
+def test_disable_lazy_routes_flag_route_table_survives_a_boot_burst_and_a_killed_worker(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    probes: Final = ("/routes", "/v1/mcp/server", "/openapi.json", "/guardrails/list", "/vector_store/list") * 8
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}, workers=2) as owned:
+        at_boot: Final = _routed_features(owned.gateway)
+        unregistered: Final = sorted(name for name, paths in at_boot.items() if not paths)
+        assert unregistered == [], f"features still missing from /routes at startup: {unregistered}"
+        with ThreadPoolExecutor(max_workers=8) as pool:
+            statuses: Final = tuple(pool.map(partial(_status, owned.gateway), probes))
+        assert statuses == (200,) * len(probes), statuses
+        assert _routed_features(owned.gateway) == at_boot, "the boot burst changed the route table"
+
+        victim: Final = eventually(lambda: _workers(owned), lambda workers: len(workers) == 2)[0]
+        victim.kill()
+        with httpx.Client(base_url=owned.gateway.client.base_url, timeout=15, trust_env=False) as fresh:
+            survivor: Final = Gateway(fresh, owned.gateway.key, owned.gateway.upstream_url)
+            during: Final = tuple(survivor.request("GET", "/v1/mcp/server").status_code for _ in range(10))
+        assert during == (200,) * 10, during
+        respawned: Final = eventually(
+            lambda: frozenset(worker.pid for worker in _workers(owned)),
+            lambda pids: len(pids) == 2 and victim.pid not in pids,
+            seconds=30,
+        )
+        assert f"Child process [{victim.pid}] died" in owned.log.read_text(), respawned
+        tables: Final = tuple(_routed_features(owned.gateway) for _ in range(4))
+        assert all(table == at_boot for table in tables), "the respawned worker disagrees on the route table"
+
+
+def test_disable_lazy_routes_flag_yields_the_same_route_table_after_a_restart(gateway: Gateway, tmp_path: Path) -> None:
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as first:
+        table: Final = _paths(first.gateway)
+    with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}) as second:
+        assert _paths(second.gateway) == table
+        assert all(_routed_features(second.gateway).values()), "a feature is missing after restart"
+
+
+SELF_HOSTED_LANGFUSE: Final = "/self-hosted-langfuse"
+
+
+@dataclass(frozen=True, slots=True)
+class _ConfiguredFeatures:
+    alias: str
+    config: Path
+    policy: Wire
+    langfuse: Wire
+    peer: McpPeer
+
+
+def _allow(request: Request) -> Reply:
+    return Reply(body=json.dumps({"action": "NONE"}).encode())
+
+
+def _langfuse_health(request: Request) -> Reply:
+    return Reply(body=json.dumps({"status": "OK"}).encode())
+
+
+def _config_declaring(directory: Path, alias: str, policy: Wire, langfuse: Wire, peer: McpPeer) -> Path:
+    base: Final = object_value(
+        JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    )
+    config: Final = {
+        **base,
+        "guardrails": [
+            {
+                "guardrail_name": alias,
+                "litellm_params": {
+                    "guardrail": "generic_guardrail_api",
+                    "mode": "pre_call",
+                    "default_on": True,
+                    "api_base": policy.url,
+                    "api_key": "synthetic-guardrail-key",
+                },
+            }
+        ],
+        "mcp_servers": {alias: peer.registration()},
+        "general_settings": {
+            **object_value(base["general_settings"]),
+            "pass_through_endpoints": [
+                {
+                    "path": "/langfuse",
+                    "target": langfuse.url + SELF_HOSTED_LANGFUSE,
+                    "include_subpath": True,
+                    "auth": True,
+                }
+            ],
+        },
+    }
+    path: Final = directory / "configured-features.yaml"
+    path.write_text(yaml.safe_dump(config))
+    return path
+
+
+@contextmanager
+def _configured_features(directory: Path) -> Iterator[_ConfiguredFeatures]:
+    alias: Final = "lazyroutes" + uuid.uuid4().hex[:8]
+    with (
+        wire_server(_allow) as policy,
+        wire_server(_langfuse_health) as langfuse,
+        scripted_peer(echo_tool("add")) as peer,
+    ):
+        config: Final = _config_declaring(directory, alias, policy, langfuse, peer)
+        yield _ConfiguredFeatures(alias, config, policy, langfuse, peer)
+
+
+def _config_server_id(candidate: Gateway, alias: str) -> str:
+    servers: Final = JSON.validate_json(candidate.request("GET", "/v1/mcp/server").content)
+    assert isinstance(servers, list), servers
+    return next(
+        string_value(object_value(server)["server_id"])
+        for server in servers
+        if object_value(server)["server_name"] == alias
+    )
+
+
+def _assert_config_declared_features_serve(owned: OwnedProxy, features: _ConfiguredFeatures, provider: Wire) -> None:
+    marker: Final = _marker()
+    with owned.gateway.scenario() as scenario:
+        model: Final = scenario.model(api_base=provider.url + "/v1")
+        completion: Final = owned.gateway.request(
+            "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]}
+        )
+        assert completion.status_code == 200, completion.text
+        screened: Final = [request for request in features.policy.drain() if marker.encode() in request.body]
+        assert len(screened) == 1, "the config-declared guardrail did not screen the completion"
+        assert len([request for request in provider.drain() if marker.encode() in request.body]) == 1
+
+        identity: Final = _config_server_id(owned.gateway, features.alias)
+        key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
+        tool: Final = tool_names(owned.gateway, key, identity)["add"]
+        features.peer.drain()
+        called: Final = call_tool(owned.gateway, key, identity, tool, {"marker": marker})
+        assert called.status_code == 200, called.text
+        reached_peer: Final = [
+            object_value(object_value(call["body"])["params"]) for call in tool_calls(features.peer.drain())
+        ]
+        assert [(params["name"], params["arguments"]) for params in reached_peer] == [("add", {"marker": marker})]
+
+    forwarded: Final = owned.gateway.request("GET", "/langfuse/api/public/health")
+    assert forwarded.status_code == 200, forwarded.text
+    reached_langfuse: Final = tuple(request.target for request in features.langfuse.drain())
+    assert reached_langfuse == (SELF_HOSTED_LANGFUSE + "/api/public/health",), (
+        f"the config pass-through for /langfuse lost to the built-in Langfuse route: {reached_langfuse}"
+    )
+
+
+def test_disable_lazy_routes_flag_serves_config_declared_features_like_the_warmed_lazy_proxy(
+    gateway: Gateway, tmp_path: Path, provider: Wire
+) -> None:
+    with _configured_features(tmp_path) as features:
+        with owned_proxy_process(gateway, tmp_path, {}, config=features.config, remove_environment=(FLAG,)) as lazy:
+            _assert_config_declared_features_serve(lazy, features, provider)
+            _warm_every_feature(lazy.gateway)
+            warmed_paths: Final = tuple(path for path in _paths(lazy.gateway) if path != WARMUP_ROUTE)
+            warmed_openapi: Final = _openapi_paths(lazy.gateway)
+        with owned_proxy_process(gateway, tmp_path, {FLAG: "true"}, config=features.config) as eager:
+            _assert_config_declared_features_serve(eager, features, provider)
+            assert _paths(eager.gateway) == warmed_paths
+            assert _openapi_paths(eager.gateway) == warmed_openapi

tests/integration/database/test_request_log_indexes_at_boot.py

@@ -0,0 +1,364 @@
+import os
+import shutil
+import subprocess
+import sys
+from collections.abc import Mapping
+from dataclasses import dataclass
+from itertools import product
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import psycopg
+import pytest
+from integration._support.client import Gateway, eventually
+from integration._support.database import scratch_database
+from integration._support.process import LEGACY_MIGRATE_DEPLOY, MIGRATE_DEPLOY, owned_proxy_process
+from psycopg import sql
+from psycopg.rows import class_row
+
+REPO_ROOT: Final = Path(__file__).resolve().parents[3]
+PRISMA_DIR: Final = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras"
+PARTITION_SCRIPT: Final = REPO_ROOT / "db_scripts" / "partition_spend_logs.sql"
+SHIPPED_MIGRATIONS: Final = tuple(sorted(path.name for path in (PRISMA_DIR / "migrations").iterdir() if path.is_dir()))
+API_KEY_INDEX_MIGRATION: Final = "20260823000000_add_spend_logs_api_key_starttime_index"
+CALL_ID_INDEX_MIGRATION: Final = "20260831120001_spend_logs_litellm_call_id_index"
+ORIGINAL_MIGRATION_SQL: Final = MappingProxyType(
+    {
+        API_KEY_INDEX_MIGRATION: (
+            "-- CreateIndex\n"
+            'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" '
+            'ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
+        ),
+        CALL_ID_INDEX_MIGRATION: (
+            "-- CreateIndex\n"
+            'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" '
+            'ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
+        ),
+    }
+)
+MIGRATION_JOB_SECONDS: Final = 300
+INDEXES_IN_PLACE: Final = "Request-log indexes are all in place"
+INDEX_BUILD_LINES: Final = ("Building index", "Attached index")
+BUILD_SECONDS: Final = 60
+SPEND_LOGS_INDEXES: Final = ("LiteLLM_SpendLogs_api_key_startTime_idx", "LiteLLM_SpendLogs_litellm_call_id_idx")
+POPULATED_PARTITIONS: Final = MappingProxyType(
+    {
+        "LiteLLM_SpendLogs_p2026_08": ("2026-08-01", "2026-09-01"),
+        "LiteLLM_SpendLogs_p2026_09": ("2026-09-01", "2026-10-01"),
+    }
+)
+DEFAULT_PARTITION: Final = "LiteLLM_SpendLogs_pdefault"
+ROWS_PER_PARTITION: Final = 500
+PARTITIONED_PARENT_ERROR: Final = 'cannot create index on partitioned table "LiteLLM_SpendLogs" concurrently'
+
+
+@dataclass(frozen=True, slots=True)
+class Resolver:
+    """One migration resolver as the serving proxy selects it (CLI flags) and as the
+    migration job selects it (environment)."""
+
+    proxy_flags: tuple[str, ...]
+    job_environment: Mapping[str, str]
+
+
+V2: Final = Resolver(MIGRATE_DEPLOY, MappingProxyType({"USE_V2_MIGRATION_RESOLVER": "true"}))
+LEGACY: Final = Resolver(LEGACY_MIGRATE_DEPLOY, MappingProxyType({"USE_V2_MIGRATION_RESOLVER": "false"}))
+RESOLVERS: Final = pytest.mark.parametrize("resolver", (V2, LEGACY), ids=("v2", "legacy"))
+
+
+def release_layout(directory: Path, migrations: tuple[str, ...]) -> Path:
+    """The Prisma layout of the release that shipped `migrations`: the two index migrations
+    carry the SQL they shipped with, not the inert files of this build."""
+    (directory / "migrations").mkdir(parents=True)
+    shutil.copy(PRISMA_DIR / "schema.prisma", directory / "schema.prisma")
+    shutil.copy(PRISMA_DIR / "migrations" / "migration_lock.toml", directory / "migrations" / "migration_lock.toml")
+    for name in migrations:
+        shutil.copytree(PRISMA_DIR / "migrations" / name, directory / "migrations" / name)
+    for name, original in ORIGINAL_MIGRATION_SQL.items():
+        if name in migrations:
+            (directory / "migrations" / name / "migration.sql").write_text(original)
+    return directory / "schema.prisma"
+
+
+def migrate_deploy(database_url: str, schema: Path) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(schema)],
+        capture_output=True,
+        text=True,
+        timeout=300,
+        env={**os.environ, "DATABASE_URL": database_url},
+    )
+
+
+def migration_job(database_url: str, resolver: Resolver) -> subprocess.CompletedProcess[str]:
+    """The migrations image entrypoint, as the helm migration Job runs it."""
+    return subprocess.run(
+        [sys.executable, "-I", str(REPO_ROOT / "migrations" / "run.py")],
+        capture_output=True,
+        text=True,
+        timeout=MIGRATION_JOB_SECONDS,
+        cwd=REPO_ROOT,
+        env={**os.environ, "DATABASE_URL": database_url, **resolver.job_environment},
+    )
+
+
+def migration_cli(database_url: str, gateway: Gateway, resolver: Resolver) -> subprocess.CompletedProcess[str]:
+    """The proxy CLI as a migration job: `--skip_server_startup` migrates, builds the
+    indexes and exits by their result."""
+    return subprocess.run(
+        [
+            sys.executable,
+            "-P",
+            "-m",
+            "integration._support.proxy",
+            "--config",
+            "tests/integration/proxy_config.yaml",
+            *resolver.proxy_flags,
+            "--skip_server_startup",
+            "--enforce_prisma_migration_check",
+        ],
+        capture_output=True,
+        text=True,
+        timeout=MIGRATION_JOB_SECONDS,
+        cwd=REPO_ROOT,
+        env={**os.environ, "DATABASE_URL": database_url, "LITELLM_MASTER_KEY": gateway.key},
+    )
+
+
+def deploy_schema_before_the_index_migrations(database_url: str, directory: Path) -> None:
+    older: Final = tuple(name for name in SHIPPED_MIGRATIONS if name < API_KEY_INDEX_MIGRATION)
+    deployed: Final = migrate_deploy(database_url, release_layout(directory / "older-release", older))
+    assert deployed.returncode == 0, deployed.stdout + deployed.stderr
+
+
+def deploy_the_original_index_migrations(database_url: str, directory: Path) -> subprocess.CompletedProcess[str]:
+    """Boot the v1.103.0 layout once: on a partitioned table its CONCURRENTLY call_id index
+    fails with P3018 and leaves the ledger row unfinished, on a plain table both apply."""
+    return migrate_deploy(database_url, release_layout(directory / "v1.103.0", SHIPPED_MIGRATIONS))
+
+
+def fail_the_call_id_index_migration_like_the_shipped_release(database_url: str, directory: Path) -> None:
+    deployed: Final = deploy_the_original_index_migrations(database_url, directory)
+    assert deployed.returncode != 0, deployed.stdout
+    assert "P3018" in deployed.stderr and PARTITIONED_PARENT_ERROR in deployed.stderr, deployed.stderr
+    assert ledger(database_url)[CALL_ID_INDEX_MIGRATION] is False
+
+
+def partition_spend_logs(database_url: str) -> None:
+    with psycopg.connect(database_url, autocommit=True) as connection:
+        connection.execute(PARTITION_SCRIPT.read_bytes())
+        for partition, (start, stop) in POPULATED_PARTITIONS.items():
+            add_partition(connection, partition, start, stop)
+            connection.execute(
+                'INSERT INTO "LiteLLM_SpendLogs" ("request_id", "call_type", "api_key", "startTime", "endTime") '
+                "SELECT %s || n, 'acompletion', 'sk-' || (n %% 7), %s::timestamp + (n * interval '1 minute'), "
+                "%s::timestamp + (n * interval '1 minute') + interval '1 second' FROM generate_series(1, %s) AS n",
+                (partition, start, start, ROWS_PER_PARTITION),
+            )
+        connection.execute(
+            'INSERT INTO "LiteLLM_SpendLogs" ("request_id", "call_type", "startTime", "endTime") '
+            "SELECT 'default-' || n, 'acompletion', '2026-07-01'::timestamp + (n * interval '1 minute'), "
+            "'2026-07-01'::timestamp + (n * interval '1 minute') FROM generate_series(1, %s) AS n",
+            (ROWS_PER_PARTITION,),
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class _LedgerRow:
+    name: str
+    finished: bool
+
+
+@dataclass(frozen=True, slots=True)
+class _IndexRow:
+    index: str
+    valid: bool
+
+
+@dataclass(frozen=True, slots=True)
+class _OidRow:
+    index: str
+    oid: int
+
+
+@dataclass(frozen=True, slots=True)
+class _AttachedRow:
+    partition: str
+    parent_index: str
+    valid: bool
+
+
+def ledger(database_url: str) -> Mapping[str, bool]:
+    """Every migration in the ledger that was not rolled back, mapped to whether it finished."""
+    with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_LedgerRow)) as cursor:
+        rows: Final = cursor.execute(
+            'SELECT migration_name AS name, finished_at IS NOT NULL AS finished FROM "_prisma_migrations" '
+            "WHERE rolled_back_at IS NULL ORDER BY migration_name"
+        ).fetchall()
+    return MappingProxyType({row.name: row.finished for row in rows})
+
+
+def parent_index_validity(database_url: str) -> Mapping[str, bool]:
+    with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_IndexRow)) as cursor:
+        rows: Final = cursor.execute(
+            "SELECT c.relname AS index, x.indisvalid AS valid FROM pg_index x JOIN pg_class c ON c.oid = x.indexrelid "
+            "WHERE x.indrelid = '\"LiteLLM_SpendLogs\"'::regclass AND c.relname = ANY(%s)",
+            (list(SPEND_LOGS_INDEXES),),
+        ).fetchall()
+    return MappingProxyType({row.index: row.valid for row in rows})
+
+
+def index_oids(database_url: str) -> Mapping[str, int]:
+    """index name -> oid for the SpendLogs indexes on the parent or any partition; a rebuild changes the oid."""
+    with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_OidRow)) as cursor:
+        rows: Final = cursor.execute(
+            "SELECT c.relname AS index, c.oid::int AS oid FROM pg_index x JOIN pg_class c ON c.oid = x.indexrelid "
+            "WHERE x.indrelid = '\"LiteLLM_SpendLogs\"'::regclass OR x.indrelid IN "
+            "(SELECT inhrelid FROM pg_inherits WHERE inhparent = '\"LiteLLM_SpendLogs\"'::regclass)"
+        ).fetchall()
+    return MappingProxyType({row.index: row.oid for row in rows})
+
+
+def attached_partition_indexes(database_url: str) -> frozenset[tuple[str, str, bool]]:
+    """(partition, parent index, child is valid) for every child index attached under a SpendLogs parent index."""
+    with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_AttachedRow)) as cursor:
+        rows: Final = cursor.execute(
+            "SELECT part.relname AS partition, parent_index.relname AS parent_index, child.indisvalid AS valid "
+            "FROM pg_inherits attached "
+            "JOIN pg_class parent_index ON parent_index.oid = attached.inhparent "
+            "JOIN pg_index child ON child.indexrelid = attached.inhrelid "
+            "JOIN pg_class part ON part.oid = child.indrelid "
+            "WHERE parent_index.relname = ANY(%s)",
+            (list(SPEND_LOGS_INDEXES),),
+        ).fetchall()
+    return frozenset((row.partition, row.parent_index, row.valid) for row in rows)
+
+
+def expected_attachments(partitions: tuple[str, ...]) -> frozenset[tuple[str, str, bool]]:
+    return frozenset((partition, index, True) for partition, index in product(partitions, SPEND_LOGS_INDEXES))
+
+
+def add_partition(connection: psycopg.Connection[tuple[object, ...]], partition: str, start: str, stop: str) -> None:
+    connection.execute(
+        sql.SQL('CREATE TABLE {} PARTITION OF "LiteLLM_SpendLogs" FOR VALUES FROM ({}) TO ({})').format(
+            sql.Identifier(partition), sql.Literal(start), sql.Literal(stop)
+        )
+    )
+
+
+def assert_ready(booted_gateway: Gateway) -> None:
+    readiness: Final = booted_gateway.request("GET", "/health/readiness")
+    assert readiness.status_code == 200, readiness.text
+    assert readiness.json()["db"] == "connected", readiness.text
+
+
+def assert_both_indexes_cover_every_partition(database_url: str) -> None:
+    """Every populated partition's index is attached and valid, both parents are valid, and
+    a partition created afterwards inherits both indexes."""
+    populated: Final = (*POPULATED_PARTITIONS, DEFAULT_PARTITION)
+    assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
+    assert attached_partition_indexes(database_url) == expected_attachments(populated)
+    assert parent_index_validity(database_url) == {index: True for index in SPEND_LOGS_INDEXES}
+    with psycopg.connect(database_url, autocommit=True) as connection:
+        add_partition(connection, "LiteLLM_SpendLogs_p2026_10", "2026-10-01", "2026-11-01")
+    assert attached_partition_indexes(database_url) == expected_attachments((*populated, "LiteLLM_SpendLogs_p2026_10"))
+
+
+def assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(
+    gateway: Gateway, directory: Path, database_url: str, resolver: Resolver
+) -> None:
+    """The serving proxy applies the inert files, reports ready, and its background build
+    finds every index already there, so it builds nothing and the catalog is untouched."""
+    oids: Final = index_oids(database_url)
+    with owned_proxy_process(
+        gateway, directory, {"DATABASE_URL": database_url}, database_setup=resolver.proxy_flags
+    ) as booted:
+        assert_ready(booted.gateway)
+        log: Final = eventually(
+            lambda: booted.log.read_text(errors="replace"), lambda text: INDEXES_IN_PLACE in text, seconds=BUILD_SECONDS
+        )
+    assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
+    assert not any(line in log for line in INDEX_BUILD_LINES), log[-4000:]
+    assert index_oids(database_url) == oids
+
+
+@RESOLVERS
+def test_the_migration_job_gives_a_partitioned_table_at_the_pre_index_schema_both_indexes_per_partition(
+    gateway: Gateway, tmp_path: Path, resolver: Resolver
+) -> None:
+    with scratch_database() as database_url:
+        deploy_schema_before_the_index_migrations(database_url, tmp_path)
+        partition_spend_logs(database_url)
+        job: Final = migration_job(database_url, resolver)
+        assert job.returncode == 0, job.stdout + job.stderr
+        assert INDEXES_IN_PLACE in job.stderr + job.stdout, job.stdout + job.stderr
+        assert_both_indexes_cover_every_partition(database_url)
+        assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
+
+
+@RESOLVERS
+def test_the_migration_job_heals_a_partitioned_table_left_with_the_failed_call_id_ledger_row(
+    gateway: Gateway, tmp_path: Path, resolver: Resolver
+) -> None:
+    with scratch_database() as database_url:
+        deploy_schema_before_the_index_migrations(database_url, tmp_path)
+        partition_spend_logs(database_url)
+        fail_the_call_id_index_migration_like_the_shipped_release(database_url, tmp_path)
+        job: Final = migration_job(database_url, resolver)
+        assert job.returncode == 0, job.stdout + job.stderr
+        assert_both_indexes_cover_every_partition(database_url)
+        assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
+
+
+@RESOLVERS
+def test_the_migration_job_leaves_a_plain_table_that_applied_the_original_index_migrations_alone(
+    gateway: Gateway, tmp_path: Path, resolver: Resolver
+) -> None:
+    with scratch_database() as database_url:
+        deploy_schema_before_the_index_migrations(database_url, tmp_path)
+        deployed: Final = deploy_the_original_index_migrations(database_url, tmp_path)
+        assert deployed.returncode == 0, deployed.stdout + deployed.stderr
+        before: Final = index_oids(database_url)
+        assert set(SPEND_LOGS_INDEXES) <= set(before), before
+        job: Final = migration_job(database_url, resolver)
+        assert job.returncode == 0, job.stdout + job.stderr
+        assert INDEXES_IN_PLACE in job.stderr + job.stdout, job.stdout + job.stderr
+        assert "Building index" not in job.stderr + job.stdout, job.stdout + job.stderr
+        assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
+        assert index_oids(database_url) == before
+        assert_the_serving_proxy_boots_and_finds_the_indexes_in_place(gateway, tmp_path, database_url, resolver)
+
+
+@RESOLVERS
+def test_a_serving_proxy_that_runs_the_migrations_itself_builds_both_indexes_after_it_is_ready(
+    gateway: Gateway, tmp_path: Path, resolver: Resolver
+) -> None:
+    """A deployment that runs migrate deploy from the serving proxy and never runs the
+    migration job answers readiness with the inert files applied, then its background build
+    puts both indexes on every partition."""
+    with scratch_database() as database_url:
+        deploy_schema_before_the_index_migrations(database_url, tmp_path)
+        partition_spend_logs(database_url)
+        assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in parent_index_validity(database_url)
+        with owned_proxy_process(
+            gateway, tmp_path, {"DATABASE_URL": database_url}, database_setup=resolver.proxy_flags
+        ) as booted:
+            assert_ready(booted.gateway)
+            assert ledger(database_url) == {name: True for name in SHIPPED_MIGRATIONS}
+            log: Final = eventually(
+                lambda: booted.log.read_text(errors="replace"),
+                lambda text: INDEXES_IN_PLACE in text,
+                seconds=BUILD_SECONDS,
+            )
+        assert "Building index" in log and "Attached index" in log, log[-4000:]
+        assert_both_indexes_cover_every_partition(database_url)
+
+
+def test_the_cli_run_as_a_migration_job_builds_both_indexes_before_it_exits(gateway: Gateway, tmp_path: Path) -> None:
+    with scratch_database() as database_url:
+        deploy_schema_before_the_index_migrations(database_url, tmp_path)
+        partition_spend_logs(database_url)
+        job: Final = migration_cli(database_url, gateway, V2)
+        assert job.returncode == 0, job.stdout + job.stderr
+        assert_both_indexes_cover_every_partition(database_url)

tests/integration/mcp/test_mcp_llm_endpoints.py

@@ -9,6 +9,7 @@ import httpx
 import pytest
 from integration._support.client import Gateway, Scenario
 from integration._support.mcp import McpPeer, mcp_peer, register_mcp, tool_calls
+from integration._support.mcp_grants import create_toolset
 from integration._support.wire import Reply, Request, Wire, wire_server
 
 Surface = Literal["chat", "responses", "messages", "messages_bridge"]
@@ -194,9 +195,7 @@ class Rig:
         )
 
     def upstream_tools(self) -> tuple[tuple[str, ...], ...]:
-        return tuple(
-            _tool_names(json.loads(request.body)) for request in self.wire.drain() if request.method == "POST"
-        )
+        return tuple(_tool_names(json.loads(request.body)) for request in self.wire.drain() if request.method == "POST")
 
     def final_text(self, body: Mapping[str, object]) -> str:
         if self.surface == "chat":
@@ -314,6 +313,25 @@ def test_allowed_tools_narrows_the_tool_list_handed_to_the_model(gateway: Gatewa
         assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
 
 
+@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
+def test_toolset_gateway_url_serves_a_team_granted_toolset_to_a_key_without_its_own_grant(
+    gateway: Gateway, surface: Surface
+) -> None:
+    with _rig(gateway, surface) as rig:
+        register_mcp(rig.scenario, rig.peer, "open" + uuid.uuid4().hex[:8], allow_all_keys=True)
+        toolset_name: Final = "ts" + uuid.uuid4().hex[:8]
+        toolset_id: Final = create_toolset(rig.scenario, ((rig.server_id, "add"),), toolset_name=toolset_name)
+        sibling_id: Final = create_toolset(rig.scenario, ((rig.server_id, "multiply"),))
+        team_id: Final = rig.scenario.team(object_permission={"mcp_toolsets": [toolset_id, sibling_id]})
+        key: Final = rig.scenario.key(team_id=team_id)
+        response: Final = rig.send(key, [{**AUTO, "server_url": f"litellm_proxy/mcp/{toolset_name}"}])
+        assert response.status_code == 200, response.text
+        requests: Final = rig.upstream_tools()
+        assert requests, "model was never called"
+        assert all(names == (rig.tool,) for names in requests), requests
+        assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
+
+
 @pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
 def test_server_scoped_gateway_url_exposes_only_that_servers_tools(gateway: Gateway, surface: Surface) -> None:
     with _rig(gateway, surface) as rig, mcp_peer() as other_peer:
@@ -347,3 +365,41 @@ def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway
         assert text == ANSWER, response.text
         assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
         assert len(rig.upstream_tools()) == 2
+
+
+def test_streaming_chat_through_a_toolset_gateway_url_serves_a_team_key_without_its_own_grant(
+    gateway: Gateway,
+) -> None:
+    with _rig(gateway, "chat") as rig:
+        register_mcp(rig.scenario, rig.peer, "open" + uuid.uuid4().hex[:8], allow_all_keys=True)
+        toolset_name: Final = "ts" + uuid.uuid4().hex[:8]
+        toolset_id: Final = create_toolset(rig.scenario, ((rig.server_id, "add"),), toolset_name=toolset_name)
+        key: Final = rig.scenario.key(team_id=rig.scenario.team(object_permission={"mcp_toolsets": [toolset_id]}))
+        response: Final = rig.send(key, [{**AUTO, "server_url": f"litellm_proxy/mcp/{toolset_name}"}], stream=True)
+        assert response.status_code == 200, response.text
+        chunks: Final = tuple(
+            json.loads(line.removeprefix("data: "))
+            for line in response.text.splitlines()
+            if line.startswith("data: ") and line != "data: [DONE]"
+        )
+        text: Final = "".join(
+            str(chunk["choices"][0]["delta"].get("content") or "") for chunk in chunks if chunk.get("choices")
+        )
+        assert text == ANSWER, response.text
+        assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"]
+        requests: Final = rig.upstream_tools()
+        assert len(requests) == 2 and all(names == (rig.tool,) for names in requests), requests
+
+
+@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
+def test_toolset_gateway_url_gives_a_key_of_an_ungranted_team_no_tools_and_never_reaches_the_peer(
+    gateway: Gateway, surface: Surface
+) -> None:
+    with _rig(gateway, surface) as rig:
+        toolset_name: Final = "ts" + uuid.uuid4().hex[:8]
+        create_toolset(rig.scenario, ((rig.server_id, "add"),), toolset_name=toolset_name)
+        key: Final = rig.scenario.key(team_id=rig.scenario.team())
+        response: Final = rig.send(key, [{**AUTO, "server_url": f"litellm_proxy/mcp/{toolset_name}"}])
+        assert _peer_add_calls(rig.peer) == (), "denied caller reached the peer"
+        assert all(rig.tool not in names for names in rig.upstream_tools()), rig.upstream_tools()
+        assert response.status_code in (200, 400, 401, 403), response.text

tests/integration/mcp/test_mcp_oauth_flows.py

@@ -1,25 +1,31 @@
 import base64
 import hashlib
 import secrets
+import time
 import uuid
 from dataclasses import dataclass
 from typing import Final
 from urllib.parse import parse_qs, urlsplit
 
 import httpx
+import jwt
 import pytest
 from integration._support.client import Gateway, eventually
 from integration._support.database import read_rows
 from integration._support.mcp import (
     ENTRY_POINTS,
+    INITIALIZE,
     EntryPoint,
     McpCaller,
     McpPeer,
+    Outcome,
+    _outcome_from_rpc,
     call_tool,
     mcp_peer,
     register_mcp,
     tool_calls,
 )
+from integration._support.mcp_grants import create_toolset
 from integration._support.oauth_server import AuthorizationServer, oauth_server
 
 ADD: Final = {"a": 2, "b": 3}
@@ -383,3 +389,117 @@ def test_dcr_bridge_relays_client_registration_and_advertises_gateway_endpoints(
         assert issuer.json()["authorization_endpoint"] == f"{_base(gateway)}/{alias}/authorize"
         assert issuer.json()["token_endpoint"] == f"{_base(gateway)}/{alias}/token"
         assert "S256" in issuer.json()["code_challenge_methods_supported"]
+
+
+def _ui_session_cookie(gateway: Gateway, user_id: str) -> dict[str, str]:
+    claims: Final = {"user_id": user_id, "login_method": "username_password", "exp": int(time.time()) + 600}
+    return {"token": jwt.encode(claims, gateway.key, algorithm="HS256")}
+
+
+def _gateway_session_bearer(gateway: Gateway, user_id: str, resource: str | None = None) -> str:
+    registered: Final = gateway.client.post(
+        "/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
+    )
+    assert registered.status_code in (200, 201), registered.text
+    client_id: Final = registered.json()["client_id"]
+    pkce: Final = _Pkce(secrets.token_urlsafe(48))
+    cookies: Final = _ui_session_cookie(gateway, user_id)
+    started: Final = gateway.client.get(
+        "/authorize",
+        params={
+            "client_id": client_id,
+            "redirect_uri": CLIENT_REDIRECT,
+            "response_type": "code",
+            "state": "lit6029",
+            "code_challenge": pkce.challenge,
+            "code_challenge_method": "S256",
+            **({} if resource is None else {"resource": resource}),
+        },
+        cookies=cookies,
+    )
+    assert started.status_code == 303, started.text
+    handle: Final = parse_qs(urlsplit(started.headers["location"]).query)["connect_flow"][0]
+    completed: Final = gateway.client.post(
+        "/authorize/complete", data={"flow": handle}, cookies={**cookies, **dict(started.cookies)}
+    )
+    assert completed.status_code == 303, completed.text
+    callback: Final = parse_qs(urlsplit(completed.headers["location"]).query)
+    assert "code" in callback, completed.headers["location"]
+    issued: Final = gateway.client.post(
+        "/token",
+        data={
+            "grant_type": "authorization_code",
+            "code": callback["code"][0],
+            "redirect_uri": CLIENT_REDIRECT,
+            "client_id": client_id,
+            "code_verifier": pkce.verifier,
+        },
+    )
+    assert issued.status_code == 200, issued.text
+    return _issued_token(issued.json())
+
+
+def _toolset_rpc(gateway: Gateway, bearer: str, name: str, method: str, params: dict[str, object]) -> Outcome:
+    def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response:
+        return gateway.client.post(
+            f"/toolset/{name}/mcp",
+            json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params},
+            headers={"Authorization": f"Bearer {bearer}", "Accept": "application/json, text/event-stream"},
+        )
+
+    initialized: Final = _outcome_from_rpc(post("initialize", dict(INITIALIZE)))
+    if not initialized.ok:
+        return initialized
+    return _outcome_from_rpc(post(method, params))
+
+
+def test_gateway_session_bearer_of_a_team_member_is_served_the_team_toolset_on_its_route(gateway: Gateway) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029sess" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_name: Final = "lit6029g" + uuid.uuid4().hex[:8]
+        withheld_name: Final = "lit6029w" + uuid.uuid4().hex[:8]
+        granted_id: Final = create_toolset(scenario, ((server_id, "add"),), toolset_name=granted_name)
+        create_toolset(scenario, ((server_id, "multiply"),), toolset_name=withheld_name)
+        member: Final = scenario.member(scenario.team(object_permission={"mcp_toolsets": [granted_id]}))
+        bearer: Final = _gateway_session_bearer(gateway, member)
+        assert bearer.startswith("llm_session_"), bearer[:16]
+        listed: Final = _toolset_rpc(gateway, bearer, granted_name, "tools/list", {})
+        assert listed.tools == (f"{alias}-add",), listed.raw
+        peer.drain()
+        called: Final = _toolset_rpc(
+            gateway, bearer, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": {"a": 4, "b": 5}}
+        )
+        assert called.ok and called.text == "9", called.raw
+        assert len(tool_calls(peer.drain())) == 1
+        denied: Final = _toolset_rpc(gateway, bearer, withheld_name, "tools/list", {})
+        assert denied.status == 403, denied.raw
+        assert tool_calls(peer.drain()) == ()
+
+
+def test_resource_scoped_session_bearer_opens_a_team_toolset_inside_its_server_and_none_outside(
+    gateway: Gateway,
+) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        inside: Final = "lit6029in" + uuid.uuid4().hex[:6]
+        outside: Final = "lit6029out" + uuid.uuid4().hex[:6]
+        inside_server: Final = register_mcp(scenario, peer, inside)
+        outside_server: Final = register_mcp(scenario, peer, outside)
+        inside_name: Final = "lit6029i" + uuid.uuid4().hex[:8]
+        outside_name: Final = "lit6029o" + uuid.uuid4().hex[:8]
+        inside_id: Final = create_toolset(scenario, ((inside_server, "add"),), toolset_name=inside_name)
+        outside_id: Final = create_toolset(scenario, ((outside_server, "add"),), toolset_name=outside_name)
+        member: Final = scenario.member(scenario.team(object_permission={"mcp_toolsets": [inside_id, outside_id]}))
+        bearer: Final = _gateway_session_bearer(gateway, member, resource=f"{_base(gateway)}/{inside}/mcp")
+        assert bearer.startswith("llm_session_"), bearer[:16]
+        listed: Final = _toolset_rpc(gateway, bearer, inside_name, "tools/list", {})
+        assert listed.tools == (f"{inside}-add",), listed.raw
+        peer.drain()
+        called: Final = _toolset_rpc(
+            gateway, bearer, inside_name, "tools/call", {"name": f"{inside}-add", "arguments": {"a": 4, "b": 5}}
+        )
+        assert called.ok and called.text == "9", called.raw
+        assert len(tool_calls(peer.drain())) == 1
+        refused: Final = _toolset_rpc(gateway, bearer, outside_name, "tools/list", {})
+        assert refused.status == 403, refused.raw
+        assert tool_calls(peer.drain()) == ()

tests/integration/mcp/test_mcp_toolsets.py

@@ -0,0 +1,509 @@
+import secrets
+import uuid
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+import yaml
+from integration._support.client import Gateway, Scenario, object_value
+from integration._support.mcp import (
+    INITIALIZE,
+    Outcome,
+    _outcome_from_rest,
+    _outcome_from_rpc,
+    mcp_peer,
+    register_mcp,
+    tool_calls,
+)
+from integration._support.mcp_grants import create_toolset
+from integration._support.process import owned_proxy
+
+from litellm.models.user import LiteLLM_UserTable
+from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+from litellm.proxy.auth.auth_checks import LITELLM_SESSION_TOKEN_PREFIX, ExperimentalUIJWTToken
+from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token
+
+ADD: Final = {"a": 4, "b": 5}
+
+
+def _dashboard_ui_session_token(user_id: str) -> str:
+    user: Final = LiteLLM_UserTable(user_id=user_id, user_role="internal_user", models=[])
+    return ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token(user)
+
+
+def _toolset(scenario: Scenario, server_id: str, tool: str) -> tuple[str, str]:
+    name: Final = "lit6029_" + uuid.uuid4().hex[:10]
+    return create_toolset(scenario, ((server_id, tool),), toolset_name=name), name
+
+
+def _toolset_rpc(
+    gateway: Gateway, headers: dict[str, str], name: str, method: str, params: dict[str, object]
+) -> Outcome:
+    def post(rpc_method: str, rpc_params: dict[str, object]) -> httpx.Response:
+        return gateway.client.post(
+            f"/toolset/{name}/mcp",
+            json={"jsonrpc": "2.0", "id": 1, "method": rpc_method, "params": rpc_params},
+            headers={**headers, "Accept": "application/json, text/event-stream"},
+        )
+
+    initialized: Final = _outcome_from_rpc(post("initialize", dict(INITIALIZE)))
+    if not initialized.ok:
+        return initialized
+    return _outcome_from_rpc(post(method, params))
+
+
+def _listed_toolset_ids(gateway: Gateway, headers: dict[str, str]) -> tuple[str, ...]:
+    response: Final = gateway.client.get("/v1/mcp/toolset", headers=headers)
+    assert response.status_code == 200, response.text
+    return tuple(toolset["toolset_id"] for toolset in response.json())
+
+
+def _assert_team_grants_only(gateway: Gateway, team_id: str, key: str, toolset_id: str) -> None:
+    team: Final = object_value(gateway.get("/team/info", {"team_id": team_id})["team_info"])
+    assert object_value(team["object_permission"])["mcp_toolsets"] == [toolset_id], team
+    key_info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
+    assert key_info.get("object_permission") is None, f"key must carry no grant of its own: {key_info}"
+
+
+def test_team_granted_toolset_is_listed_and_served_to_a_team_key(gateway: Gateway) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        withheld_id, withheld_name = _toolset(scenario, server_id, "multiply")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        key: Final = scenario.key(team_id=team_id)
+        _assert_team_grants_only(gateway, team_id, key, granted_id)
+        headers: Final = {"Authorization": f"Bearer {key}"}
+
+        assert _listed_toolset_ids(gateway, headers) == (granted_id,)
+        detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
+        assert detail.status_code == 200, detail.text
+        assert detail.json()["toolset_name"] == granted_name, detail.text
+        withheld_detail: Final = gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=headers)
+        assert withheld_detail.status_code == 403, withheld_detail.text
+
+        listed: Final = _toolset_rpc(gateway, headers, granted_name, "tools/list", {})
+        assert listed.ok, listed.raw
+        assert listed.tools == (f"{alias}-add",), listed.raw
+        peer.drain()
+        called: Final = _toolset_rpc(
+            gateway, headers, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": ADD}
+        )
+        assert called.ok and called.text == "9", called.raw
+        assert len(tool_calls(peer.drain())) == 1
+        denied: Final = _toolset_rpc(gateway, headers, withheld_name, "tools/list", {})
+        assert denied.status == 403, denied.raw
+
+
+def test_dashboard_session_of_a_team_member_lists_the_team_granted_toolset(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        _toolset(scenario, server_id, "multiply")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        user_id: Final = scenario.user(user_role="internal_user", teams=[team_id])
+        user: Final = object_value(gateway.get("/user/info", {"user_id": user_id})["user_info"])
+        assert user["teams"] == [team_id], user
+        headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
+
+        assert _listed_toolset_ids(gateway, headers) == (granted_id,)
+        detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
+        assert detail.status_code == 200, detail.text
+        assert detail.json()["toolset_name"] == granted_name, detail.text
+
+
+def test_direct_grants_no_grants_and_admin_listing_are_unchanged_by_team_resolution(gateway: Gateway) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        withheld_id, withheld_name = _toolset(scenario, server_id, "multiply")
+        direct: Final = {"Authorization": f"Bearer {scenario.key(object_permission={'mcp_toolsets': [granted_id]})}"}
+        ungranted_team: Final = scenario.team()
+        no_grant: Final = {"Authorization": f"Bearer {scenario.key(team_id=ungranted_team)}"}
+        admin: Final = {"Authorization": f"Bearer {gateway.key}"}
+
+        assert _listed_toolset_ids(gateway, direct) == (granted_id,)
+        assert _toolset_rpc(gateway, direct, granted_name, "tools/list", {}).tools == (f"{alias}-add",)
+        assert _toolset_rpc(gateway, direct, withheld_name, "tools/list", {}).status == 403
+        assert gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=direct).status_code == 403
+
+        assert _listed_toolset_ids(gateway, no_grant) == ()
+        assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=no_grant).status_code == 403
+        assert _toolset_rpc(gateway, no_grant, granted_name, "tools/list", {}).status == 403
+
+        assert {granted_id, withheld_id} <= set(_listed_toolset_ids(gateway, admin))
+        assert gateway.client.get(f"/v1/mcp/toolset/{withheld_id}", headers=admin).status_code == 200
+        assert _toolset_rpc(gateway, admin, withheld_name, "tools/list", {}).tools == (f"{alias}-multiply",)
+
+
+def test_a_key_with_its_own_toolset_grant_does_not_inherit_the_team_toolset(gateway: Gateway) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        own_id, own_name = _toolset(scenario, server_id, "add")
+        team_only_id, team_only_name = _toolset(scenario, server_id, "multiply")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [own_id, team_only_id]})
+        key: Final = scenario.key(team_id=team_id, object_permission={"mcp_toolsets": [own_id]})
+        headers: Final = {"Authorization": f"Bearer {key}"}
+
+        assert _listed_toolset_ids(gateway, headers) == (own_id,)
+        assert gateway.client.get(f"/v1/mcp/toolset/{team_only_id}", headers=headers).status_code == 403
+        assert _toolset_rpc(gateway, headers, team_only_name, "tools/list", {}).status == 403
+        assert _toolset_rpc(gateway, headers, own_name, "tools/list", {}).tools == (f"{alias}-add",)
+
+
+def _team_member_with_own_grant(scenario: Scenario, team_id: str, own_server_id: str) -> str:
+    user_id: Final = scenario.user(user_role="internal_user", object_permission={"mcp_servers": [own_server_id]})
+    scenario.gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}})
+    return user_id
+
+
+def test_dashboard_session_serves_the_team_toolset_despite_a_disjoint_grant_on_the_user_row(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The member's own row grants a different server outright. That grant must not cap the team's toolset
+    to nothing, and the team's sibling toolset must not leak onto the granted toolset's route."""
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        own_server_id: Final = register_mcp(scenario, peer, "lit6029_own_" + uuid.uuid4().hex[:8])
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        sibling_id, sibling_name = _toolset(scenario, server_id, "multiply")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id, sibling_id]})
+        user_id: Final = _team_member_with_own_grant(scenario, team_id, own_server_id)
+        headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
+
+        assert set(_listed_toolset_ids(gateway, headers)) == {granted_id, sibling_id}
+        listed: Final = _toolset_rpc(gateway, headers, granted_name, "tools/list", {})
+        assert listed.ok, listed.raw
+        assert listed.tools == (f"{alias}-add",), listed.raw
+        assert _toolset_rpc(gateway, headers, sibling_name, "tools/list", {}).tools == (f"{alias}-multiply",)
+        peer.drain()
+        called: Final = _toolset_rpc(
+            gateway, headers, granted_name, "tools/call", {"name": f"{alias}-add", "arguments": ADD}
+        )
+        assert called.ok and called.text == "9", called.raw
+        assert len(tool_calls(peer.drain())) == 1
+        stranger: Final = {
+            "Authorization": f"Bearer {_dashboard_ui_session_token(scenario.user(user_role='internal_user'))}"
+        }
+        assert _toolset_rpc(gateway, stranger, granted_name, "tools/list", {}).status == 403
+
+
+def test_a_member_removed_from_the_team_loses_its_toolset_on_the_dashboard_session(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029_" + uuid.uuid4().hex[:8]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        user_id: Final = scenario.member(team_id)
+        headers: Final = {"Authorization": f"Bearer {_dashboard_ui_session_token(user_id)}"}
+        assert _listed_toolset_ids(gateway, headers) == (granted_id,)
+        assert _toolset_rpc(gateway, headers, granted_name, "tools/list", {}).tools == (f"{alias}-add",)
+
+        gateway.post("/team/member_delete", {"team_id": team_id, "user_id": user_id})
+
+        assert _listed_toolset_ids(gateway, headers) == ()
+        assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
+        assert _toolset_rpc(gateway, headers, granted_name, "tools/list", {}).status == 403
+
+
+def _bearer(token: str) -> dict[str, str]:
+    return {"Authorization": f"Bearer {token}"}
+
+
+def _expired_dashboard_token(user_id: str) -> str:
+    expired: Final = (datetime.now(timezone.utc) - timedelta(minutes=5)).strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3]
+    stale: Final = UserAPIKeyAuth(
+        token="ui-token",
+        key_name="ui-token",
+        key_alias="ui-token",
+        expires=expired + "+00:00",
+        user_id=user_id,
+        team_id="litellm-dashboard",
+        models=[],
+        user_role=LitellmUserRoles.INTERNAL_USER,
+    )
+    return encrypt_bearer_token(stale.model_dump_json(exclude_none=True), prefix=LITELLM_SESSION_TOKEN_PREFIX)
+
+
+def _rest_list(gateway: Gateway, headers: dict[str, str], params: object) -> httpx.Response:
+    return gateway.client.get("/mcp-rest/tools/list", headers=headers, params=params)
+
+
+def _rest_call(gateway: Gateway, headers: dict[str, str], name: str, server_id: str) -> Outcome:
+    return _outcome_from_rest(
+        gateway.client.post(
+            "/mcp-rest/tools/call",
+            headers=headers,
+            json={"name": name, "arguments": dict(ADD), "server_id": server_id},
+        )
+    )
+
+
+def _route_tools(gateway: Gateway, headers: dict[str, str], name: str) -> Outcome:
+    return _toolset_rpc(gateway, headers, name, "tools/list", {})
+
+
+def _route_call(gateway: Gateway, headers: dict[str, str], name: str, tool: str) -> Outcome:
+    return _toolset_rpc(gateway, headers, name, "tools/call", {"name": tool, "arguments": dict(ADD)})
+
+
+def _strict_config(directory: Path) -> Path:
+    base: Final = yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
+    strict: Final = {
+        **base,
+        "general_settings": {**base.get("general_settings", {}), "require_key_mcp_access_defined": True},
+    }
+    path: Final = directory / "require_key_mcp_access.yaml"
+    path.write_text(yaml.safe_dump(strict))
+    return path
+
+
+def test_mcp_rest_toolset_name_narrows_the_list_and_serves_the_call_for_a_team_key_and_a_dashboard_member(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029rest" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        _, withheld_name = _toolset(scenario, server_id, "multiply")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        key: Final = scenario.key(team_id=team_id)
+        member: Final = scenario.member(team_id)
+        for headers in (_bearer(key), _bearer(_dashboard_ui_session_token(member))):
+            listed: Final = _rest_list(gateway, headers, {"toolset_name": granted_name})
+            assert listed.status_code == 200, listed.text
+            listed_names: Final = tuple(tool["name"] for tool in listed.json()["tools"])
+            assert len(listed_names) == 1 and listed_names[0].endswith("add"), listed.text
+            denied: Final = _rest_list(gateway, headers, {"toolset_name": withheld_name})
+            assert denied.status_code == 200 and denied.json()["tools"] == [], denied.text
+            assert "does not have access to toolset" in denied.json()["message"], denied.text
+            peer.drain()
+            called: Final = _rest_call(gateway, headers, listed_names[0], server_id)
+            assert called.ok and called.text == "9", called.raw
+            assert len(tool_calls(peer.drain())) == 1
+
+
+def test_a_key_restricted_to_its_own_servers_does_not_inherit_the_team_toolset(gateway: Gateway) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029own" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        other_id: Final = register_mcp(scenario, peer, "lit6029other" + uuid.uuid4().hex[:6])
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id], "mcp_servers": [other_id]})
+        key: Final = scenario.key(team_id=team_id, object_permission={"mcp_servers": [other_id]})
+        headers: Final = _bearer(key)
+        assert _listed_toolset_ids(gateway, headers) == ()
+        assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
+        assert _route_tools(gateway, headers, granted_name).status == 403
+        assert tool_calls(peer.drain()) == ()
+
+
+def test_a_team_with_an_empty_or_absent_toolset_grant_gives_its_keys_nothing(gateway: Gateway) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        server_id: Final = register_mcp(scenario, peer, "lit6029empty" + uuid.uuid4().hex[:6])
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        teams: Final = (
+            scenario.team(object_permission={"mcp_toolsets": []}),
+            scenario.team(object_permission={"mcp_toolsets": None}),
+            scenario.team(),
+        )
+        for team_id in teams:
+            headers: Final = _bearer(scenario.key(team_id=team_id))
+            assert _listed_toolset_ids(gateway, headers) == (), team_id
+            assert gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers).status_code == 403
+            assert _route_tools(gateway, headers, granted_name).status == 403, team_id
+        assert tool_calls(peer.drain()) == ()
+
+
+def test_an_unknown_or_malformed_toolset_name_is_refused_without_peer_traffic_and_the_route_keeps_serving(
+    gateway: Gateway,
+) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029bad" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        _, withheld_name = _toolset(scenario, server_id, "multiply")
+        headers: Final = _bearer(scenario.key(team_id=scenario.team(object_permission={"mcp_toolsets": [granted_id]})))
+        unknown: Final = "missing" + uuid.uuid4().hex[:8]
+        assert _route_tools(gateway, headers, unknown).status == 404
+        assert _rest_list(gateway, headers, {"toolset_name": unknown}).status_code == 404
+        malformed: Final = (
+            [("toolset_name", granted_name), ("toolset_name", withheld_name)],
+            {"toolset_name": "x" * 5000},
+            {"toolset_name": ""},
+            {"toolset_name": granted_name + "\x00"},
+        )
+        statuses: Final = tuple(_rest_list(gateway, headers, params).status_code for params in malformed)
+        assert all(status < 500 for status in statuses), statuses
+        assert tool_calls(peer.drain()) == ()
+        assert gateway.client.get("/health/liveliness").status_code == 200
+        served: Final = _route_tools(gateway, headers, granted_name)
+        assert served.tools == (f"{alias}-add",), served.raw
+
+
+def test_garbage_expired_and_tampered_credentials_are_refused_on_every_toolset_surface(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        server_id: Final = register_mcp(scenario, peer, "lit6029cred" + uuid.uuid4().hex[:6])
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        member: Final = scenario.member(team_id)
+        forged: Final = (
+            "sk-" + secrets.token_urlsafe(24),
+            _expired_dashboard_token(member),
+            "llm_session_" + secrets.token_urlsafe(32),
+        )
+        for bearer in forged:
+            headers: Final = _bearer(bearer)
+            listed: Final = gateway.client.get("/v1/mcp/toolset", headers=headers)
+            assert listed.status_code == 401, (bearer[:12], listed.text)
+            detail: Final = gateway.client.get(f"/v1/mcp/toolset/{granted_id}", headers=headers)
+            assert detail.status_code == 401, (bearer[:12], detail.text)
+            routed: Final = _route_tools(gateway, headers, granted_name)
+            assert routed.status == 401, (bearer[:12], routed.raw)
+            rest: Final = _rest_list(gateway, headers, {"toolset_name": granted_name})
+            assert rest.status_code == 401, (bearer[:12], rest.text)
+        assert peer.drain() == ()
+
+
+def test_a_dashboard_member_of_a_deleted_team_loses_the_toolset_while_a_direct_grant_survives(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029gone" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        team_id_, team_name = _toolset(scenario, server_id, "add")
+        own_id, own_name = _toolset(scenario, server_id, "multiply")
+        doomed: Final = scenario.gateway.post(
+            "/team/new",
+            {"team_alias": f"integration-{uuid.uuid4().hex}", "object_permission": {"mcp_toolsets": [team_id_]}},
+        )
+        doomed_team: Final = str(doomed["team_id"])
+        member: Final = scenario.user(user_role="internal_user", teams=[doomed_team])
+        granted: Final = scenario.user(
+            user_role="internal_user", teams=[doomed_team], object_permission={"mcp_toolsets": [own_id]}
+        )
+        member_headers: Final = _bearer(_dashboard_ui_session_token(member))
+        granted_headers: Final = _bearer(_dashboard_ui_session_token(granted))
+        assert _listed_toolset_ids(gateway, member_headers) == (team_id_,)
+        assert set(_listed_toolset_ids(gateway, granted_headers)) == {team_id_, own_id}
+        scenario.delete_team(doomed_team)
+        assert _listed_toolset_ids(gateway, member_headers) == ()
+        assert gateway.client.get(f"/v1/mcp/toolset/{team_id_}", headers=member_headers).status_code == 403
+        assert _route_tools(gateway, member_headers, team_name).status == 403
+        assert _listed_toolset_ids(gateway, granted_headers) == (own_id,)
+        assert _route_tools(gateway, granted_headers, team_name).status == 403
+        assert _route_tools(gateway, granted_headers, own_name).tools == (f"{alias}-multiply",)
+        peer.drain()
+        kept: Final = _route_call(gateway, granted_headers, own_name, f"{alias}-multiply")
+        assert kept.ok and kept.text == "20", kept.raw
+        assert [call["body"]["params"]["name"] for call in tool_calls(peer.drain())] == ["multiply"]
+
+
+def test_require_key_mcp_access_defined_stops_key_inheritance_but_not_the_dashboard_member(
+    gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with (
+        owned_proxy(gateway, tmp_path, {}, config=_strict_config(tmp_path), workers=2) as strict,
+        mcp_peer() as peer,
+        strict.scenario() as scenario,
+    ):
+        alias: Final = "lit6029strict" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        inheriting: Final = _bearer(scenario.key(team_id=team_id))
+        own: Final = _bearer(scenario.key(team_id=team_id, object_permission={"mcp_toolsets": [granted_id]}))
+        member: Final = _bearer(_dashboard_ui_session_token(scenario.member(team_id)))
+        assert _listed_toolset_ids(strict, inheriting) == ()
+        assert _route_tools(strict, inheriting, granted_name).status == 403
+        assert _listed_toolset_ids(strict, own) == (granted_id,)
+        assert _route_tools(strict, own, granted_name).tools == (f"{alias}-add",)
+        assert _listed_toolset_ids(strict, member) == (granted_id,)
+        assert _route_tools(strict, member, granted_name).tools == (f"{alias}-add",)
+        peer.drain()
+        called: Final = _route_call(strict, member, granted_name, f"{alias}-add")
+        assert called.ok and called.text == "9", called.raw
+        assert len(tool_calls(peer.drain())) == 1
+
+
+def test_a_key_with_only_a_vector_store_grant_still_inherits_the_team_toolset(gateway: Gateway) -> None:
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029vs" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        granted_id, granted_name = _toolset(scenario, server_id, "add")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        headers: Final = _bearer(
+            scenario.key(team_id=team_id, object_permission={"vector_stores": ["vs-" + uuid.uuid4().hex[:8]]})
+        )
+        assert _listed_toolset_ids(gateway, headers) == (granted_id,)
+        assert _route_tools(gateway, headers, granted_name).tools == (f"{alias}-add",)
+        peer.drain()
+        called: Final = _route_call(gateway, headers, granted_name, f"{alias}-add")
+        assert called.ok and called.text == "9", called.raw
+        assert len(tool_calls(peer.drain())) == 1
+
+
+def test_a_member_added_after_the_team_was_cached_sees_the_toolset_on_every_following_request(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        server_id: Final = register_mcp(scenario, peer, "lit6029cache" + uuid.uuid4().hex[:6])
+        granted_id, _ = _toolset(scenario, server_id, "add")
+        team_id: Final = scenario.team(object_permission={"mcp_toolsets": [granted_id]})
+        user_id: Final = scenario.user(user_role="internal_user")
+        headers: Final = _bearer(_dashboard_ui_session_token(user_id))
+        warm: Final = _bearer(scenario.key(team_id=team_id))
+        assert tuple(_listed_toolset_ids(gateway, warm) for _ in range(4)) == ((granted_id,),) * 4
+        assert tuple(_listed_toolset_ids(gateway, headers) for _ in range(4)) == ((),) * 4
+        added: Final = gateway.request(
+            "POST", "/team/member_add", {"team_id": team_id, "member": {"user_id": user_id, "role": "user"}}
+        )
+        assert added.status_code == 200, added.text
+        listings: Final = tuple(_listed_toolset_ids(gateway, headers) for _ in range(8))
+        assert listings == ((granted_id,),) * 8, listings
+
+
+def test_a_member_of_two_teams_sees_the_union_and_each_route_stays_narrowed_to_its_own_toolset(
+    gateway: Gateway, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setenv("LITELLM_SALT_KEY", "sk-integration-salt")
+    with mcp_peer() as peer, gateway.scenario() as scenario:
+        alias: Final = "lit6029two" + uuid.uuid4().hex[:6]
+        server_id: Final = register_mcp(scenario, peer, alias)
+        first_id, first_name = _toolset(scenario, server_id, "add")
+        second_id, second_name = _toolset(scenario, server_id, "multiply")
+        teams: Final = (
+            scenario.team(object_permission={"mcp_toolsets": [first_id]}),
+            scenario.team(object_permission={"mcp_toolsets": [second_id]}),
+        )
+        headers: Final = _bearer(
+            _dashboard_ui_session_token(scenario.user(user_role="internal_user", teams=list(teams)))
+        )
+        assert set(_listed_toolset_ids(gateway, headers)) == {first_id, second_id}
+        assert _route_tools(gateway, headers, first_name).tools == (f"{alias}-add",)
+        assert _route_tools(gateway, headers, second_name).tools == (f"{alias}-multiply",)
+        peer.drain()
+        crossed: Final = _route_call(gateway, headers, first_name, f"{alias}-multiply")
+        assert not crossed.ok, crossed.raw
+        assert tool_calls(peer.drain()) == ()

tests/integration/providers/test_bedrock_converse_missing_content_wire.py

@@ -0,0 +1,966 @@
+import asyncio
+import base64
+import json
+import os
+import re
+import signal
+import threading
+import uuid
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from queue import SimpleQueue
+from types import MappingProxyType
+from typing import Final, Literal
+from urllib.parse import unquote, urlsplit
+
+import anthropic
+import httpx
+import openai
+import psutil
+import pytest
+import yaml
+from integration._support.client import (
+    Gateway,
+    Scenario,
+    eventually,
+    gateway_from_environment,
+    object_value,
+    string_value,
+)
+from integration._support.database import read_rows
+from integration._support.process import owned_proxy_process
+from integration._support.upstream import _aws_event_frame
+from integration._support.wire import Reply, Request, Wire, wire_server
+from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
+from pydantic import JsonValue, TypeAdapter
+
+_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0"
+_CONVERSE_MODEL: Final = f"bedrock/converse/{_MODEL_ID}"
+_INVOKE_MODEL: Final = f"bedrock/invoke/{_MODEL_ID}"
+_CONVERSE_TARGET: Final = f"/model/{_MODEL_ID}/converse"
+_STREAM_TARGET: Final = f"/model/{_MODEL_ID}/converse-stream"
+_INVOKE_TARGET: Final = f"/model/{_MODEL_ID}/invoke"
+_ANSWER: Final = "bedrock missing content control"
+_RESPONSE: Final = json.dumps(
+    {
+        "output": {"message": {"role": "assistant", "content": [{"text": _ANSWER}]}},
+        "stopReason": "end_turn",
+        "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
+        "metrics": {"latencyMs": 1},
+    }
+).encode()
+_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
+_STREAM_EVENTS: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = (
+    ("messageStart", {"role": "assistant"}),
+    ("contentBlockDelta", {"delta": {"text": _ANSWER}, "contentBlockIndex": 0}),
+    ("messageStop", {"stopReason": "end_turn"}),
+    ("metadata", {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}}),
+)
+_STREAM_BYTES: Final = b"".join(_aws_event_frame(kind, payload, "sc", "u") for kind, payload in _STREAM_EVENTS)
+_DEFAULT_CONTINUE: Final = "Please continue."
+_DEPLOYMENT_CONTINUE: Final = "Deployment says continue."
+_DEPLOYMENT_CONTINUE_MESSAGE: Final[dict[str, JsonValue]] = {"role": "user", "content": _DEPLOYMENT_CONTINUE}
+_NO_NON_SYSTEM_MESSAGE: Final = "bedrock requires at least one non-system message"
+_JSON: Final = TypeAdapter(dict[str, JsonValue])
+_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
+_CALL_INDEX: Final = re.compile(r"call-[0-9a-f]{32}-(\d+)")
+_QUESTION: Final = "What is the capital of France?"
+_ANSWERED: Final = "Paris."
+_FOLLOW_UP: Final = "And the capital of Spain?"
+_QUESTION_TURN: Final[dict[str, JsonValue]] = {"role": "user", "content": _QUESTION}
+_ANSWERED_TURN: Final[dict[str, JsonValue]] = {"role": "assistant", "content": _ANSWERED}
+_FOLLOW_UP_TURN: Final[dict[str, JsonValue]] = {"role": "user", "content": _FOLLOW_UP}
+_NO_CONTENT_USER: Final[dict[str, JsonValue]] = {"role": "user"}
+_NULL_CONTENT_USER: Final[dict[str, JsonValue]] = {"role": "user", "content": None}
+_EMPTY_CONTENT_USER: Final[dict[str, JsonValue]] = {"role": "user", "content": ""}
+_NO_CONTENT_SYSTEM: Final[dict[str, JsonValue]] = {"role": "system"}
+_NULL_CONTENT_SYSTEM: Final[dict[str, JsonValue]] = {"role": "system", "content": None}
+_NO_CONTENT_ASSISTANT: Final[dict[str, JsonValue]] = {"role": "assistant"}
+_TOOL_CALL_TURN: Final[dict[str, JsonValue]] = {
+    "role": "assistant",
+    "content": None,
+    "tool_calls": [
+        {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": '{"city": "Boston"}'}}
+    ],
+}
+_NO_CONTENT_TOOL: Final[dict[str, JsonValue]] = {"role": "tool", "tool_call_id": "call_1"}
+_TOOLS: Final[tuple[dict[str, JsonValue], ...]] = (
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "parameters": {"type": "object", "properties": {"city": {"type": "string"}}},
+        },
+    },
+)
+_CONVERSE_QUESTION: Final[dict[str, JsonValue]] = {"role": "user", "content": [{"text": _QUESTION}]}
+_CONVERSE_ANSWERED: Final[dict[str, JsonValue]] = {"role": "assistant", "content": [{"text": _ANSWERED}]}
+_CONVERSE_FOLLOW_UP: Final[dict[str, JsonValue]] = {"role": "user", "content": [{"text": _FOLLOW_UP}]}
+_CONVERSE_TOOL_USE: Final[dict[str, JsonValue]] = {
+    "role": "assistant",
+    "content": [{"toolUse": {"toolUseId": "call_1", "name": "get_weather", "input": {"city": "Boston"}}}],
+}
+_CONVERSE_EMPTY_TOOL_RESULT: Final[dict[str, JsonValue]] = {
+    "role": "user",
+    "content": [{"toolResult": {"toolUseId": "call_1", "content": []}}],
+}
+_NEUTRALIZED_TOOL_CALL: Final[dict[str, JsonValue]] = {
+    "role": "assistant",
+    "content": [{"text": '[tool call call_1: get_weather({"city": "Boston"})]'}],
+}
+_NEUTRALIZED_TOOL_RESULT: Final[dict[str, JsonValue]] = {
+    "role": "user",
+    "content": [{"text": "[tool result for call_1: <non-text tool result omitted>]"}],
+}
+_EXTRA: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}}
+_SIGNING_KEY: Final = os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt")
+_AWS: Final[dict[str, JsonValue]] = {
+    "aws_access_key_id": "AKIASCRIPTEDPROVIDER",
+    "aws_secret_access_key": "scripted-secret",
+    "aws_region_name": "us-east-1",
+}
+_PLAIN: Final = "bedrock-missing-content-plain"
+_CONTINUE: Final = "bedrock-missing-content-continue"
+
+Endpoint = Literal["chat", "responses"]
+
+
+@dataclass(frozen=True, slots=True)
+class _Call:
+    endpoint: Endpoint
+    stream: bool
+    user: str
+    index: int
+
+
+@dataclass(frozen=True, slots=True)
+class _Served:
+    call: _Call
+    status: int
+    text: str
+
+
+@dataclass(frozen=True, slots=True)
+class _ModifyParamsCell:
+    row: str
+    model: str
+    messages: tuple[dict[str, JsonValue], ...]
+    expected: tuple[dict[str, JsonValue], ...]
+    extra: Mapping[str, JsonValue] = MappingProxyType({})
+
+
+def _continue_turn(text: str) -> dict[str, JsonValue]:
+    return {"role": "user", "content": [{"text": text}]}
+
+
+_MODIFY_PARAMS_CELLS: Final = (
+    _ModifyParamsCell("r11", _PLAIN, (_NO_CONTENT_USER,), (_continue_turn(_DEFAULT_CONTINUE),)),
+    _ModifyParamsCell(
+        "r12",
+        _PLAIN,
+        (_QUESTION_TURN, _ANSWERED_TURN, _NULL_CONTENT_USER),
+        (_CONVERSE_QUESTION, _CONVERSE_ANSWERED, _continue_turn(_DEFAULT_CONTINUE)),
+    ),
+    _ModifyParamsCell(
+        "r13",
+        _PLAIN,
+        (_QUESTION_TURN, _TOOL_CALL_TURN, _NO_CONTENT_TOOL),
+        (_CONVERSE_QUESTION, _CONVERSE_TOOL_USE, _CONVERSE_EMPTY_TOOL_RESULT),
+        MappingProxyType({"tools": list(_TOOLS)}),
+    ),
+    _ModifyParamsCell("r14", _PLAIN, (_NO_CONTENT_SYSTEM, _QUESTION_TURN), (_CONVERSE_QUESTION,)),
+    _ModifyParamsCell("r15", _CONTINUE, (_NO_CONTENT_USER,), (_continue_turn(_DEPLOYMENT_CONTINUE),)),
+)
+
+
+def _converse_peer(request: Request) -> Reply:
+    if unquote(request.target) == _STREAM_TARGET:
+        return Reply(body=_STREAM_BYTES, content_type=_EVENT_STREAM)
+    return Reply(body=_RESPONSE)
+
+
+def _scripted_error(status: int, message: str) -> Reply:
+    return Reply(status=status, body=json.dumps({"message": message}).encode())
+
+
+def _converse_deployment(scenario: Scenario, wire: Wire, **extra: JsonValue) -> str:
+    return scenario.model(model=_CONVERSE_MODEL, api_base=wire.url, **_AWS, **extra)
+
+
+def _auth(gateway: Gateway) -> dict[str, str]:
+    return {"Authorization": f"Bearer {gateway.key}"}
+
+
+def _proxy_url(gateway: Gateway) -> str:
+    return str(gateway.client.base_url).rstrip("/")
+
+
+def _chat(
+    model: str,
+    messages: Sequence[Mapping[str, JsonValue]],
+    *,
+    stream: bool = False,
+    cached: bool = False,
+    **extra: JsonValue,
+) -> dict[str, JsonValue]:
+    return {
+        "model": model,
+        "messages": [dict(message) for message in messages],
+        "max_tokens": 16,
+        "stream": stream,
+        "num_retries": 0,
+        **({} if cached else {"cache": {"no-cache": True}}),
+        **extra,
+    }
+
+
+def _post_chat(gateway: Gateway, body: Mapping[str, JsonValue]) -> httpx.Response:
+    return gateway.request("POST", "/v1/chat/completions", body)
+
+
+def _chat_answer(response: httpx.Response) -> str:
+    assert response.status_code == 200, response.text
+    body: Final = response.json()
+    assert body["choices"][0]["message"]["content"] == _ANSWER, response.text
+    return body["id"]
+
+
+def _only_received(wire: Wire) -> tuple[str, dict[str, JsonValue]]:
+    (request,) = wire.drain()
+    return unquote(request.target), json.loads(request.body)
+
+
+def _sse_payloads(lines: Iterable[str]) -> tuple[dict[str, JsonValue], ...]:
+    return tuple(json.loads(line[6:]) for line in lines if line.startswith("data: ") and line != "data: [DONE]")
+
+
+def _stream_lines(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[str, ...]:
+    with gateway.client.stream("POST", path, json=body, headers=_auth(gateway)) as response:
+        lines: Final = tuple(line for line in response.iter_lines() if line)
+        status_code: Final = response.status_code
+    assert status_code == 200, "\n".join(lines)
+    return lines
+
+
+def _chat_stream_text(chunks: Iterable[dict[str, JsonValue]]) -> str:
+    return "".join(chunk["choices"][0]["delta"].get("content") or "" for chunk in chunks if chunk["choices"])
+
+
+def _spend_row(request_id: str) -> dict[str, JsonValue]:
+    (row,) = eventually(
+        lambda: read_rows(
+            'SELECT request_id, status, call_type, end_user FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+            (request_id,),
+        ),
+        lambda found: len(found) >= 1,
+        seconds=70,
+    )
+    return row
+
+
+def _success_rows(prefix: str, expected: int) -> tuple[dict[str, JsonValue], ...]:
+    rows: Final = eventually(
+        lambda: read_rows(
+            'SELECT request_id, call_type, end_user FROM "LiteLLM_SpendLogs" WHERE end_user LIKE %s AND status=%s',
+            (f"{prefix}%", "success"),
+        ),
+        lambda found: len(found) >= expected,
+        seconds=70,
+    )
+    assert len(rows) == expected, rows
+    assert len({row["request_id"] for row in rows}) == expected, rows
+    return tuple(rows)
+
+
+def _converse_cell(gateway: Gateway, wire: Wire, body: Mapping[str, JsonValue]) -> tuple[str, dict[str, JsonValue]]:
+    identity: Final = _chat_answer(_post_chat(gateway, body))
+    target, received = _only_received(wire)
+    assert target == _CONVERSE_TARGET, target
+    assert _spend_row(identity)["status"] == "success"
+    return identity, received
+
+
+def _owned_config(wire: Wire, directory: Path, *, modify_params: bool) -> Path:
+    base: Final = _JSON.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    deployment: Final[dict[str, JsonValue]] = {
+        "model": _CONVERSE_MODEL,
+        "api_base": wire.url,
+        "api_key": "integration-provider-key",
+        **_AWS,
+    }
+    config: Final[dict[str, JsonValue]] = {
+        **base,
+        "model_list": [
+            {"model_name": _PLAIN, "litellm_params": deployment},
+            {
+                "model_name": _CONTINUE,
+                "litellm_params": {**deployment, "user_continue_message": _DEPLOYMENT_CONTINUE_MESSAGE},
+            },
+        ],
+        "litellm_settings": {**_JSON.validate_python(base["litellm_settings"]), "modify_params": modify_params},
+        "router_settings": {**_JSON.validate_python(base["router_settings"]), "num_retries": 0},
+    }
+    path: Final = directory / f"bedrock-missing-content-{'modify-params' if modify_params else 'plain'}.yaml"
+    path.write_text(yaml.safe_dump(config))
+    return path
+
+
+@pytest.fixture(scope="module")
+def modify_params_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Wire]]:
+    directory: Final = tmp_path_factory.mktemp("bedrock-modify-params")
+    with gateway_from_environment() as gateway, wire_server(_converse_peer) as wire:
+        config: Final = _owned_config(wire, directory, modify_params=True)
+        with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
+            yield owned.gateway, wire
+
+
+def test_r01_openai_sync_lone_user_without_content_sends_no_converse_block(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        with openai.OpenAI(base_url=f"{_proxy_url(gateway)}/v1", api_key=gateway.key, max_retries=0) as client:
+            completion: Final = client.chat.completions.create(
+                model=model, messages=[_NO_CONTENT_USER], max_tokens=16, extra_body=_EXTRA
+            )
+        assert completion.choices[0].message.content == _ANSWER, completion
+        target, received = _only_received(wire)
+        assert target == _CONVERSE_TARGET, target
+        assert received["messages"] == [] and "system" not in received, received
+        assert _spend_row(completion.id)["status"] == "success"
+
+
+async def test_r02_openai_async_user_with_null_content_after_an_assistant_turn_keeps_the_earlier_turns(
+    gateway: Gateway,
+) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        async with openai.AsyncOpenAI(
+            base_url=f"{_proxy_url(gateway)}/v1", api_key=gateway.key, max_retries=0
+        ) as client:
+            completion: Final = await client.chat.completions.create(
+                model=model,
+                messages=[_QUESTION_TURN, _ANSWERED_TURN, _NULL_CONTENT_USER],
+                max_tokens=16,
+                extra_body=_EXTRA,
+            )
+        assert completion.choices[0].message.content == _ANSWER, completion
+        target, received = _only_received(wire)
+        assert target == _CONVERSE_TARGET, target
+        assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED], received
+        assert _spend_row(completion.id)["status"] == "success"
+
+
+def test_r03_httpx_raw_sse_user_without_content_after_an_assistant_turn_streams_to_done(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        lines: Final = _stream_lines(
+            gateway,
+            "/v1/chat/completions",
+            _chat(model, (_QUESTION_TURN, _ANSWERED_TURN, _NO_CONTENT_USER), stream=True),
+        )
+        assert lines[-1] == "data: [DONE]", lines
+        chunks: Final = _sse_payloads(lines)
+        assert _chat_stream_text(chunks) == _ANSWER, lines
+        (identity,) = {chunk["id"] for chunk in chunks}
+        target, received = _only_received(wire)
+        assert target == _STREAM_TARGET, target
+        assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED], received
+        assert _spend_row(identity)["status"] == "success"
+
+
+async def test_r04_openai_async_stream_tool_turn_without_content_sends_an_empty_tool_result(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        async with openai.AsyncOpenAI(
+            base_url=f"{_proxy_url(gateway)}/v1", api_key=gateway.key, max_retries=0
+        ) as client:
+            stream: Final = await client.chat.completions.create(
+                model=model,
+                messages=[_QUESTION_TURN, _TOOL_CALL_TURN, _NO_CONTENT_TOOL],
+                tools=list(_TOOLS),
+                max_tokens=16,
+                stream=True,
+                extra_body=_EXTRA,
+            )
+            chunks: Final = [chunk async for chunk in stream]
+        assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == _ANSWER, chunks
+        (identity,) = {chunk.id for chunk in chunks}
+        target, received = _only_received(wire)
+        assert target == _STREAM_TARGET, target
+        assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_TOOL_USE, _CONVERSE_EMPTY_TOOL_RESULT], received
+        assert received["toolConfig"]["tools"][0]["toolSpec"]["name"] == "get_weather", received
+        assert _spend_row(identity)["status"] == "success"
+
+
+@pytest.mark.parametrize("system_turn", (_NO_CONTENT_SYSTEM, _NULL_CONTENT_SYSTEM), ids=("r05", "r06"))
+def test_r05_r06_leading_system_without_content_is_dropped(gateway: Gateway, system_turn: dict[str, JsonValue]) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        _, received = _converse_cell(gateway, wire, _chat(model, (system_turn, _QUESTION_TURN)))
+        assert "system" not in received, received
+        assert received["messages"] == [_CONVERSE_QUESTION], received
+
+
+def test_r07_mid_conversation_system_without_content_is_dropped(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        _, received = _converse_cell(
+            gateway, wire, _chat(model, (_QUESTION_TURN, _ANSWERED_TURN, _NO_CONTENT_SYSTEM, _FOLLOW_UP_TURN))
+        )
+        assert "system" not in received, received
+        assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED, _CONVERSE_FOLLOW_UP], received
+
+
+def test_r08_assistant_without_content_between_two_user_turns_merges_them(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        _, received = _converse_cell(
+            gateway, wire, _chat(model, (_QUESTION_TURN, _NO_CONTENT_ASSISTANT, _FOLLOW_UP_TURN))
+        )
+        assert received["messages"] == [{"role": "user", "content": [{"text": _QUESTION}, {"text": _FOLLOW_UP}]}], (
+            received
+        )
+
+
+def test_r09_lone_user_with_empty_string_content_sends_no_converse_block(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        _, received = _converse_cell(gateway, wire, _chat(model, (_EMPTY_CONTENT_USER,)))
+        assert received["messages"] == [], received
+
+
+def test_r10_empty_null_and_missing_content_produce_byte_identical_converse_bodies(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        identities: Final = tuple(
+            _chat_answer(_post_chat(gateway, _chat(model, (turn,))))
+            for turn in (_EMPTY_CONTENT_USER, _NULL_CONTENT_USER, _NO_CONTENT_USER)
+        )
+        assert len(set(identities)) == 3, identities
+        received: Final = wire.drain()
+        assert [unquote(request.target) for request in received] == [_CONVERSE_TARGET] * 3, received
+        assert len({request.body for request in received}) == 1, received
+        assert json.loads(received[0].body)["messages"] == [], received
+        for identity in identities:
+            assert _spend_row(identity)["status"] == "success"
+
+
+@pytest.mark.parametrize("cell", _MODIFY_PARAMS_CELLS, ids=lambda cell: cell.row)
+def test_r11_to_r15_modify_params_fills_the_missing_user_content(
+    cell: _ModifyParamsCell, modify_params_proxy: tuple[Gateway, Wire]
+) -> None:
+    gateway, wire = modify_params_proxy
+    _, received = _converse_cell(gateway, wire, _chat(cell.model, cell.messages, **cell.extra))
+    assert received["messages"] == list(cell.expected), received
+    assert "system" not in received, received
+
+
+def test_r16_deployment_user_continue_message_fills_the_missing_content_without_modify_params(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire, user_continue_message=_DEPLOYMENT_CONTINUE_MESSAGE)
+        _, received = _converse_cell(gateway, wire, _chat(model, (_NO_CONTENT_USER,)))
+        assert received["messages"] == [_continue_turn(_DEPLOYMENT_CONTINUE)], received
+
+
+def test_r17_anthropic_sync_lone_user_without_content_is_rejected_before_any_peer_call(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        with anthropic.Anthropic(base_url=_proxy_url(gateway), api_key=gateway.key, max_retries=0) as client:
+            with pytest.raises(anthropic.BadRequestError, match=_NO_NON_SYSTEM_MESSAGE):
+                client.messages.create(model=model, max_tokens=16, messages=[_NO_CONTENT_USER], extra_body=_EXTRA)
+        assert wire.drain() == ()
+
+
+async def test_r18_anthropic_async_stream_user_without_content_after_an_assistant_turn_keeps_the_earlier_turns(
+    gateway: Gateway,
+) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        async with anthropic.AsyncAnthropic(base_url=_proxy_url(gateway), api_key=gateway.key, max_retries=0) as client:
+            async with client.messages.stream(
+                model=model,
+                max_tokens=16,
+                messages=[_QUESTION_TURN, _ANSWERED_TURN, _NO_CONTENT_USER],
+                extra_body=_EXTRA,
+            ) as stream:
+                events: Final = [event async for event in stream]
+                final: Final = await stream.get_final_message()
+        (started,) = tuple(event for event in events if event.type == "message_start")
+        assert final.content[0].text == _ANSWER, final
+        assert final.id == started.message.id, (final.id, started.message.id)
+        target, received = _only_received(wire)
+        assert target == _STREAM_TARGET, target
+        assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED], received
+        assert _spend_row(final.id)["call_type"] == "anthropic_messages"
+
+
+def test_r19_anthropic_native_invoke_forwards_the_turn_verbatim_and_relays_the_scripted_400(gateway: Gateway) -> None:
+    with (
+        wire_server(lambda _: _scripted_error(400, "scripted invoke validation")) as wire,
+        gateway.scenario() as scenario,
+    ):
+        model: Final = scenario.model(
+            model=_INVOKE_MODEL, api_key=None, aws_bedrock_runtime_endpoint=wire.url, api_base=wire.url, **_AWS
+        )
+        with anthropic.Anthropic(base_url=_proxy_url(gateway), api_key=gateway.key, max_retries=0) as client:
+            with pytest.raises(anthropic.BadRequestError, match="scripted invoke validation"):
+                client.messages.create(model=model, max_tokens=16, messages=[_NO_CONTENT_USER], extra_body=_EXTRA)
+        target, received = _only_received(wire)
+        assert target == _INVOKE_TARGET, target
+        assert received["messages"] == [_NO_CONTENT_USER], received
+
+
+def test_r20_responses_lone_input_item_without_content_is_rejected_before_any_peer_call(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        response: Final = gateway.request(
+            "POST", "/v1/responses", {"model": model, "input": [_NO_CONTENT_USER], **_EXTRA}
+        )
+        assert response.status_code == 400, response.text
+        assert _NO_NON_SYSTEM_MESSAGE in response.text, response.text
+        assert wire.drain() == ()
+
+
+def test_r21_responses_stream_input_item_without_content_after_an_assistant_item_keeps_the_earlier_turns(
+    gateway: Gateway,
+) -> None:
+    marker: Final = f"call-{uuid.uuid4().hex}"
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        lines: Final = _stream_lines(
+            gateway,
+            "/v1/responses",
+            {
+                "model": model,
+                "input": [_QUESTION_TURN, _ANSWERED_TURN, _NO_CONTENT_USER],
+                "stream": True,
+                "user": marker,
+                **_EXTRA,
+            },
+        )
+        events: Final = _sse_payloads(lines)
+        (completed,) = tuple(event for event in events if event["type"] == "response.completed")
+        assert completed["response"]["output"][0]["content"][0]["text"] == _ANSWER, lines
+        target, received = _only_received(wire)
+        assert target == _STREAM_TARGET, target
+        assert received["messages"] == [_CONVERSE_QUESTION, _CONVERSE_ANSWERED], received
+        (row,) = _success_rows(marker, 1)
+        assert row["call_type"] == "aresponses" and row["end_user"] == marker, row
+        assert row["request_id"] == _inner_response_id(str(completed["response"]["id"])), (row, completed)
+
+
+def test_r22_passthrough_converse_forwards_a_message_without_content_verbatim(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        deployment: Final = scenario.model(
+            model=f"bedrock/{_MODEL_ID}", api_base=wire.url, aws_bedrock_runtime_endpoint=wire.url, **_AWS
+        )
+        response: Final = gateway.request(
+            "POST", f"/bedrock/model/{deployment}/converse", {"messages": [_NO_CONTENT_USER]}
+        )
+        assert response.status_code == 200, response.text
+        assert response.content == _RESPONSE, response.text
+        (request,) = wire.drain()
+        assert unquote(request.target) == _CONVERSE_TARGET, request.target
+        assert json.loads(request.body)["messages"] == [_NO_CONTENT_USER], request.body
+
+
+def test_r23_tool_turn_without_content_and_without_tools_is_neutralized_to_text(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        _, received = _converse_cell(gateway, wire, _chat(model, (_QUESTION_TURN, _TOOL_CALL_TURN, _NO_CONTENT_TOOL)))
+        assert received["messages"] == [_CONVERSE_QUESTION, _NEUTRALIZED_TOOL_CALL, _NEUTRALIZED_TOOL_RESULT], received
+        assert "toolConfig" not in received, received
+
+
+def test_s01_lone_user_with_integer_content_errors_before_any_peer_call(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        response: Final = _post_chat(gateway, _chat(model, ({"role": "user", "content": 42},)))
+        assert response.status_code >= 400, response.text
+        assert "error" in response.json(), response.text
+        assert wire.drain() == ()
+
+
+def test_s02_lone_user_with_list_content_sends_the_text_block(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        _, received = _converse_cell(
+            gateway, wire, _chat(model, ({"role": "user", "content": [{"type": "text", "text": _QUESTION}]},))
+        )
+        assert received["messages"] == [_CONVERSE_QUESTION], received
+
+
+def test_s03_lone_user_with_a_five_kilobyte_string_reaches_the_peer_whole(gateway: Gateway) -> None:
+    text: Final = "k" * 5120
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        _, received = _converse_cell(gateway, wire, _chat(model, ({"role": "user", "content": text},)))
+        assert received["messages"] == [{"role": "user", "content": [{"text": text}]}], received
+
+
+def test_s04_duplicate_content_keys_in_the_raw_body_let_the_last_value_win(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        raw: Final = (
+            f'{{"model": "{model}", "messages": [{{"role": "user", "content": "first", "content": "second"}}],'
+            ' "max_tokens": 16, "num_retries": 0, "cache": {"no-cache": true}}'
+        )
+        response: Final = gateway.client.post(
+            "/v1/chat/completions",
+            content=raw.encode(),
+            headers={**_auth(gateway), "content-type": "application/json"},
+        )
+        identity: Final = _chat_answer(response)
+        target, received = _only_received(wire)
+        assert target == _CONVERSE_TARGET, target
+        assert received["messages"] == [{"role": "user", "content": [{"text": "second"}]}], received
+        assert _spend_row(identity)["status"] == "success"
+
+
+def test_s05_unauthenticated_content_less_request_never_reaches_the_peer(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        response: Final = gateway.request(
+            "POST", "/v1/chat/completions", _chat(model, (_NO_CONTENT_USER,)), key="sk-integration-bogus"
+        )
+        assert response.status_code == 401, response.text
+        assert wire.drain() == (), response.text
+
+
+@pytest.mark.parametrize(
+    ("peer_status", "message", "expected"),
+    (
+        (400, "ValidationException: scripted validation", 400),
+        (429, "ThrottlingException: scripted throttle", 429),
+        (500, "scripted outage", 503),
+    ),
+    ids=("s06", "s07", "s08"),
+)
+def test_s06_to_s08_peer_errors_on_a_content_less_turn_reach_the_caller_after_one_attempt(
+    gateway: Gateway, peer_status: int, message: str, expected: int
+) -> None:
+    with wire_server(lambda _: _scripted_error(peer_status, message)) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        response: Final = _post_chat(gateway, _chat(model, (_NO_CONTENT_USER,)))
+        assert response.status_code == expected, response.text
+        assert message in response.text, response.text
+        assert len(wire.drain()) == 1, response.text
+
+
+def test_s09_unknown_model_with_a_content_less_turn_never_reaches_the_peer(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire:
+        response: Final = _post_chat(gateway, _chat(f"integration-missing-{uuid.uuid4().hex}", (_NO_CONTENT_USER,)))
+        assert response.status_code in (400, 404), response.text
+        assert wire.drain() == (), response.text
+
+
+def test_s10_a_deployment_continue_message_without_content_adds_no_converse_block(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire, user_continue_message={"role": "user"})
+        _, received = _converse_cell(gateway, wire, _chat(model, (_NO_CONTENT_USER,)))
+        assert received["messages"] == [], received
+
+
+def test_s11_a_deployment_continue_message_given_as_a_string_errors_in_the_body_and_leaves_the_proxy_serving(
+    gateway: Gateway,
+) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        broken: Final = _converse_deployment(scenario, wire, user_continue_message=_DEPLOYMENT_CONTINUE)
+        healthy: Final = _converse_deployment(scenario, wire)
+        response: Final = _post_chat(gateway, _chat(broken, (_NO_CONTENT_USER,)))
+        assert response.status_code >= 400, response.text
+        assert "error" in response.json(), response.text
+        assert wire.drain() == ()
+        _, received = _converse_cell(gateway, wire, _chat(healthy, (_QUESTION_TURN,)))
+        assert received["messages"] == [_CONVERSE_QUESTION], received
+
+
+def test_e01_the_same_content_less_request_twice_with_no_cache_hits_the_peer_twice(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        first: Final = _chat_answer(_post_chat(gateway, _chat(model, (_NO_CONTENT_USER,))))
+        second: Final = _chat_answer(_post_chat(gateway, _chat(model, (_NO_CONTENT_USER,))))
+        assert first != second
+        assert len(wire.drain()) == 2
+        assert _spend_row(first)["status"] == "success"
+        assert _spend_row(second)["status"] == "success"
+
+
+def test_e02_the_same_content_less_request_twice_is_served_from_the_response_cache(gateway: Gateway) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        body: Final = _chat(model, (_NO_CONTENT_USER,), cached=True)
+        first: Final = _chat_answer(_post_chat(gateway, body))
+        second: Final = _chat_answer(_post_chat(gateway, body))
+        assert first == second
+        assert len(wire.drain()) == 1
+        assert _spend_row(first)["status"] == "success"
+        cache_hits: Final = eventually(
+            lambda: read_rows(
+                'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', (f"{first}_cache_hit%",)
+            ),
+            lambda rows: len(rows) >= 1,
+            seconds=70,
+        )
+        assert len(cache_hits) == 1, cache_hits
+
+
+def _model_id(gateway: Gateway, name: str) -> str:
+    entries: Final = gateway.get("/model/info")["data"]
+    assert isinstance(entries, list), entries
+    (identity,) = (
+        string_value(object_value(object_value(entry)["model_info"])["id"])
+        for entry in entries
+        if object_value(entry)["model_name"] == name
+    )
+    return identity
+
+
+def _content_less_bodies(gateway: Gateway, wire: Wire, model: str, count: int) -> tuple[JsonValue, ...]:
+    identities: Final = tuple(
+        _chat_answer(_post_chat(gateway, _chat(model, (_NO_CONTENT_USER,)))) for _ in range(count)
+    )
+    received: Final = wire.drain()
+    assert len(received) == len(identities), (identities, received)
+    bodies: Final = tuple(_JSON.validate_python(json.loads(request.body))["messages"] for request in received)
+    assert all(body in ([], [_continue_turn(_DEPLOYMENT_CONTINUE)]) for body in bodies), bodies
+    return bodies
+
+
+@pytest.mark.timeout(180)
+def test_e03_updating_the_deployment_continue_message_under_traffic_never_breaks_a_content_less_turn(
+    gateway: Gateway,
+) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        assert _content_less_bodies(gateway, wire, model, 4) == ([],) * 4
+        gateway.post(
+            "/model/update",
+            {
+                "model_info": {"id": _model_id(gateway, model)},
+                "litellm_params": {"user_continue_message": _DEPLOYMENT_CONTINUE_MESSAGE},
+            },
+        )
+        settled: Final = eventually(
+            lambda: _content_less_bodies(gateway, wire, model, 8),
+            lambda bodies: all(body == [_continue_turn(_DEPLOYMENT_CONTINUE)] for body in bodies),
+            seconds=90,
+        )
+        assert len(settled) == 8, settled
+
+
+@pytest.mark.parametrize(
+    ("continue_message", "expected"),
+    ((None, []), ({}, [_continue_turn(_DEFAULT_CONTINUE)])),
+    ids=("e04", "e05"),
+)
+def test_e04_e05_a_null_continue_message_means_absent_and_an_empty_one_means_the_default(
+    gateway: Gateway, continue_message: JsonValue, expected: list[JsonValue]
+) -> None:
+    with wire_server(_converse_peer) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire, user_continue_message=continue_message)
+        _, received = _converse_cell(gateway, wire, _chat(model, (_NO_CONTENT_USER,)))
+        assert received["messages"] == expected, received
+
+
+def _path(endpoint: Endpoint) -> str:
+    match endpoint:
+        case "chat":
+            return "/v1/chat/completions"
+        case "responses":
+            return "/v1/responses"
+
+
+def _calls(marker: str, endpoint: Endpoint, stream: bool, indexes: range) -> tuple[_Call, ...]:
+    return tuple(_Call(endpoint, stream, f"{marker}-{index}", index) for index in indexes)
+
+
+def _burst_body(model: str, call: _Call) -> dict[str, JsonValue]:
+    turns: Final = ({"role": "user", "content": f"{_QUESTION} {call.user}"}, _ANSWERED_TURN, _NO_CONTENT_USER)
+    if call.endpoint == "chat":
+        return {**_chat(model, turns, stream=call.stream), "user": call.user}
+    return {"model": model, "input": [dict(turn) for turn in turns], "stream": call.stream, "user": call.user, **_EXTRA}
+
+
+def _call_index(request: Request) -> int:
+    found: Final = _CALL_INDEX.search(request.body.decode())
+    assert found is not None, request.body
+    return int(found.group(1))
+
+
+async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
+    async with client.stream(
+        "POST", _path(call.endpoint), json=_burst_body(model, call), headers={"Authorization": f"Bearer {key}"}
+    ) as response:
+        raw: Final = await response.aread()
+    return _Served(call=call, status=response.status_code, text=raw.decode())
+
+
+async def _burst(
+    base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False
+) -> tuple[_Served, ...]:
+    async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
+        results: Final = await asyncio.gather(
+            *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors
+        )
+    for result in results:
+        assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
+    return tuple(result for result in results if isinstance(result, _Served))
+
+
+def _inner_response_id(identity: str) -> str:
+    managed: Final = decrypt_if_encrypted_with(identity.removeprefix("resp_"), _SIGNING_KEY)
+    assert managed is not None, identity
+    issued: Final = managed.split(";", 1)[0].rsplit("response_id:", 1)[1]
+    decoded: Final = base64.b64decode(issued.removeprefix("resp_")).decode()
+    return decoded.rsplit("response_id:", 1)[1]
+
+
+def _served_id(served: _Served) -> str:
+    if served.call.stream:
+        (identity,) = {chunk["id"] for chunk in _sse_payloads(served.text.splitlines())}
+        return identity
+    return json.loads(served.text)["id"]
+
+
+def _spend_row_id(served: _Served) -> str:
+    identity: Final = _served_id(served)
+    return identity if served.call.endpoint == "chat" else _inner_response_id(identity)
+
+
+def _answered(served: Iterable[_Served]) -> frozenset[str]:
+    ids: Final = tuple(_spend_row_id(item) for item in served)
+    assert len(set(ids)) == len(ids), ids
+    return frozenset(ids)
+
+
+def _open_peer_connections(pid: int, peer_url: str) -> int:
+    port: Final = urlsplit(peer_url).port
+    return sum(
+        1
+        for connection in psutil.Process(pid).net_connections(kind="tcp")
+        if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
+    )
+
+
+async def test_c01_a_mixed_burst_of_content_less_calls_survives_a_peer_outage_window(gateway: Gateway) -> None:
+    marker: Final = f"call-{uuid.uuid4().hex}"
+    calls: Final = (
+        *_calls(marker, "chat", False, range(0, 10)),
+        *_calls(marker, "chat", True, range(10, 20)),
+        *_calls(marker, "responses", False, range(20, 30)),
+    )
+
+    def respond(request: Request) -> Reply:
+        if _call_index(request) % 3 == 1:
+            return _scripted_error(500, "scripted outage")
+        return _converse_peer(request)
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = _converse_deployment(scenario, wire)
+        served: Final = await _burst(_proxy_url(gateway), gateway.key, model, calls)
+        assert len(served) == 30
+        failed: Final = tuple(item for item in served if item.call.index % 3 == 1)
+        answered: Final = tuple(item for item in served if item.call.index % 3 != 1)
+        assert len(failed) == 10 and len(answered) == 20
+        for item in failed:
+            assert item.status == 503 and "scripted outage" in item.text, (item.call, item.status, item.text)
+        for item in answered:
+            assert item.status == 200 and _ANSWER in item.text, (item.call, item.status, item.text)
+        identities: Final = _answered(answered)
+        assert len(identities) == 20
+        assert {row["request_id"] for row in _success_rows(marker, 20)} == identities
+        assert len(wire.drain()) == 30
+
+
+@pytest.mark.timeout(180)
+async def test_c02_worker_sigkill_mid_burst_leaves_the_sibling_serving_content_less_turns(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    marker: Final = f"call-{uuid.uuid4().hex}"
+    again: Final = f"call-{uuid.uuid4().hex}"
+    calls: Final = _calls(marker, "chat", False, range(20))
+    release: Final = threading.Event()
+    held_indexes: Final[SimpleQueue[int]] = SimpleQueue()
+
+    def held(request: Request) -> Reply:
+        held_indexes.put(_call_index(request))
+        assert release.wait(timeout=60), "The burst was never released"
+        return _converse_peer(request)
+
+    with wire_server(held) as wire:
+        config: Final = _owned_config(wire, tmp_path, modify_params=False)
+        with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
+            candidate: Final = owned.gateway
+            workers: Final = eventually(
+                lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
+                lambda pids: len(pids) == 2,
+                seconds=30,
+            )
+            burst: Final = asyncio.create_task(
+                _burst(_proxy_url(candidate), candidate.key, _PLAIN, calls, tolerate_transport_errors=True)
+            )
+            await asyncio.to_thread(eventually, held_indexes.qsize, lambda size: size == 20, 60)
+            held_by: Final = MappingProxyType({pid: _open_peer_connections(pid, wire.url) for pid in workers})
+            assert sum(held_by.values()) == 20, held_by
+            victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
+            psutil.Process(victim_pid).send_signal(signal.SIGKILL)
+            release.set()
+            served: Final = await burst
+            assert held_by[survivor_pid] >= 10, held_by
+            assert len(served) == held_by[survivor_pid], (held_by, len(served))
+            for item in served:
+                assert item.status == 200 and _ANSWER in item.text, (item.call, item.status, item.text)
+            eventually(
+                lambda: len(_STARTED_WORKER.findall(owned.log.read_text())), lambda count: count == 3, seconds=60
+            )
+            follow_up: Final = await _burst(
+                _proxy_url(candidate), candidate.key, _PLAIN, _calls(again, "chat", False, range(6))
+            )
+            assert len(follow_up) == 6
+            for item in follow_up:
+                assert item.status == 200 and _ANSWER in item.text, (item.call, item.status, item.text)
+            assert len(wire.drain()) == 26
+            assert {row["request_id"] for row in _success_rows(marker, len(served))} == _answered(served)
+            assert {row["request_id"] for row in _success_rows(again, 6)} == _answered(follow_up)
+
+
+@pytest.mark.timeout(180)
+async def test_c03_proxy_terminated_mid_burst_lands_every_answered_content_less_call_at_most_once(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    marker: Final = f"call-{uuid.uuid4().hex}"
+    calls: Final = _calls(marker, "chat", False, range(12))
+    release: Final = threading.Event()
+    held_indexes: Final[SimpleQueue[int]] = SimpleQueue()
+
+    def held(request: Request) -> Reply:
+        held_indexes.put(_call_index(request))
+        assert release.wait(timeout=60), "The burst was never released"
+        return _converse_peer(request)
+
+    with wire_server(held) as wire:
+        config: Final = _owned_config(wire, tmp_path, modify_params=False)
+        with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=1) as owned:
+            candidate: Final = owned.gateway
+            burst: Final = asyncio.create_task(
+                _burst(_proxy_url(candidate), candidate.key, _PLAIN, calls, tolerate_transport_errors=True)
+            )
+            await asyncio.to_thread(eventually, held_indexes.qsize, lambda size: size == 12, 60)
+            owned.process.terminate()
+            release.set()
+            served: Final = await burst
+            eventually(owned.process.poll, lambda code: code is not None, seconds=60)
+            answered: Final = _answered(item for item in served if item.status == 200)
+            assert len(served) <= 12
+            landed: Final = tuple(
+                row["request_id"]
+                for row in read_rows(
+                    'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE end_user LIKE %s AND status=%s',
+                    (f"{marker}%", "success"),
+                )
+            )
+            assert len(landed) == len(set(landed)), landed
+            stray: Final = set(landed) - answered
+            assert len(stray) <= 12 - len(answered), (landed, answered)
+            assert len(wire.drain()) == 12

tests/integration/spend/test_key_metadata_recovery_probe_bounds.py

@@ -1,12 +1,11 @@
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-from pathlib import Path
 from typing import Final
 
-import litellm_proxy_extras
 import psycopg
 import pytest
+from litellm_proxy_extras.request_log_indexes import REQUEST_LOG_INDEXES
 from psycopg.types.json import Jsonb
 from pydantic import JsonValue
 
@@ -29,11 +28,8 @@ _SPEND_LOGS_DDL: Final = """
     )
 """
 
-_API_KEY_START_TIME_INDEX_MIGRATION: Final = (
-    Path(litellm_proxy_extras.__file__).parent
-    / "migrations"
-    / "20260823000000_add_spend_logs_api_key_starttime_index"
-    / "migration.sql"
+_API_KEY_START_TIME_INDEX: Final = next(
+    index for index in REQUEST_LOG_INDEXES if index.name == "LiteLLM_SpendLogs_api_key_startTime_idx"
 )
 
 _STATS_SQL: Final = """
@@ -56,7 +52,12 @@ class _Settle:
 
 def _create_spend_logs_table(database_url: str) -> None:
     write_rows(_SPEND_LOGS_DDL, (), database_url=database_url)
-    write_rows(_API_KEY_START_TIME_INDEX_MIGRATION.read_text(), (), database_url=database_url)
+    write_rows(
+        f'CREATE INDEX "{_API_KEY_START_TIME_INDEX.name}" ON "{_API_KEY_START_TIME_INDEX.table}" '  # pyright: ignore[reportArgumentType]  # DDL from the migration job index list
+        f"{_API_KEY_START_TIME_INDEX.definition}",
+        (),
+        database_url=database_url,
+    )
 
 
 def _spend_log_stats(database_url: str) -> dict[str, int]:

tests/proxy_migration_tests/test_db_schema_migration.py

@@ -5,6 +5,7 @@ import tempfile
 from pathlib import Path
 
 import pytest
+from litellm_proxy_extras.request_log_indexes import filter_request_log_index_diff
 
 
 @pytest.mark.skipif(
@@ -16,7 +17,9 @@ def test_schema_migration_in_sync():
 
     Applies every committed migration to an empty database, then diffs the result
     against schema.prisma. A non-empty diff means the schema was changed without a
-    matching migration being generated.
+    matching migration being generated. The request-log indexes the migration job
+    builds are declared in the schema and deliberately absent from the migrations,
+    so those statements are filtered out before the diff is judged.
     """
     db_url = os.environ["DATABASE_URL"]
     source_migrations_dir = Path(
@@ -60,11 +63,14 @@ def test_schema_migration_in_sync():
         )
 
         if diff.returncode == 2:
-            pytest.fail(
-                "Schema changes detected that no migration captures. Run "
-                "`python litellm/ci_cd/run_migration.py <migration_name>`.\n\n"
-                + diff.stdout
-            )
-        assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}"
+            drift = filter_request_log_index_diff(diff.stdout)
+            if drift.strip():
+                pytest.fail(
+                    "Schema changes detected that no migration captures. Run "
+                    "`python litellm/ci_cd/run_migration.py <migration_name>`.\n\n"
+                    + drift
+                )
+        else:
+            assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}"
     finally:
         shutil.rmtree(temp_base, ignore_errors=True)

tests/proxy_migration_tests/test_invalid_index_repair.py

@@ -1,12 +1,14 @@
 import os
 import threading
+import time
 import uuid
 from collections.abc import Iterator, Mapping
 from types import MappingProxyType
 from typing import Final
 
 import pytest
-from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager
+from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY
+from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager, _InvalidIndex
 
 psycopg = pytest.importorskip("psycopg")
 
@@ -20,6 +22,7 @@ requires_db: Final = pytest.mark.skipif(
 HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable"
 HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx"
 HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC'
+SECOND_HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_name_idx"
 LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable"
 LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx"
 PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable"
@@ -167,7 +170,74 @@ def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: st
 
 
 @requires_db
-def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None:
+def test_repair_yields_to_the_migration_job_building_indexes_under_the_migration_lock(scratch_schema: str) -> None:
+    """A migration job's index build holds the migration lock while its CREATE INDEX CONCURRENTLY
+    is cataloged as invalid; the repair must not rebuild that in-flight index."""
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+
+    with psycopg.connect(_base_url(), autocommit=True) as index_builder:
+        index_builder.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
+        assert ProxyExtrasDBManager.repair_invalid_indexes() is False
+        assert _index_validity(scratch_schema) == {HEALTH_INDEX: False}
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True}
+
+
+def _hold_migration_lock_once_free(release: threading.Event) -> None:
+    with psycopg.connect(_base_url(), autocommit=True) as resolver:
+        resolver.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
+        release.wait(timeout=60)
+
+
+def _wait_until_a_session_queues_for_the_migration_lock() -> None:
+    with psycopg.connect(_base_url(), autocommit=True) as conn:
+        for _ in range(200):
+            queued: Final = conn.execute(
+                "SELECT count(*) FROM pg_locks WHERE locktype = 'advisory' AND NOT granted "
+                "AND classid = %s AND objid = %s",
+                (MIGRATION_LOCK_KEY >> 32, MIGRATION_LOCK_KEY & 0xFFFFFFFF),
+            ).fetchone()
+            if queued is not None and queued[0]:
+                return
+            time.sleep(0.05)
+    pytest.fail("no session queued for the migration lock")
+
+
+@requires_db
+def test_repair_releases_the_migration_lock_between_indexes_so_a_booting_resolver_gets_in(
+    scratch_schema: str,
+) -> None:
+    """A v2 resolver on another replica waits for the migration lock; with two invalid
+    indexes to rebuild it must get the lock after the first REINDEX, not after both."""
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
+    _leave_invalid_index(scratch_schema, HEALTH_TABLE, SECOND_HEALTH_INDEX, '"model_name"')
+    release: Final = threading.Event()
+    resolver: Final = threading.Thread(target=_hold_migration_lock_once_free, args=(release,))
+
+    def repair_then_let_a_resolver_queue_for_the_lock(
+        conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex
+    ) -> None:
+        ProxyExtrasDBManager._repair_index(conn, index)
+        if not resolver.is_alive():
+            resolver.start()
+            _wait_until_a_session_queues_for_the_migration_lock()
+
+    try:
+        assert (
+            ProxyExtrasDBManager.repair_invalid_indexes(repair=repair_then_let_a_resolver_queue_for_the_lock) is False
+        )
+        assert sorted(_index_validity(scratch_schema).values()) == [False, True]
+    finally:
+        release.set()
+        resolver.join()
+
+    assert ProxyExtrasDBManager.repair_invalid_indexes() is True
+    assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, SECOND_HEALTH_INDEX: True}
+
+
+@requires_db
+def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_boot(scratch_schema: str) -> None:
     _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS)
 
     with psycopg.connect(_base_url()) as pin:

tests/proxy_migration_tests/test_request_log_indexes.py

@@ -0,0 +1,763 @@
+import os
+import queue
+import shutil
+import subprocess
+import sys
+import threading
+import time
+import uuid
+from collections.abc import Callable, Iterator, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import psycopg
+import pytest
+from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, migration_lock
+from litellm_proxy_extras.migration_recovery import roll_back_failed_inert_migration
+from litellm_proxy_extras.request_log_indexes import (
+    REQUEST_LOG_INDEXES,
+    RequestLogIndex,
+    build_index_on_partitioned_table,
+    ensure_request_log_indexes,
+)
+from litellm_proxy_extras.utils import ProxyExtrasDBManager
+from psycopg import sql
+from psycopg.abc import Params, QueryNoTemplate
+from psycopg.rows import class_row
+
+pytestmark = pytest.mark.timeout(900)
+
+requires_db: Final = pytest.mark.skipif(
+    "DATABASE_URL" not in os.environ,
+    reason="requires a postgres database (DATABASE_URL)",
+)
+
+REPO: Final = Path(__file__).resolve().parents[2]
+PACKAGE: Final = REPO / "litellm-proxy-extras" / "litellm_proxy_extras"
+PARTITION_SCRIPT: Final = REPO / "db_scripts" / "partition_spend_logs.sql"
+API_KEY_INDEX_MIGRATION: Final = "20260823000000_add_spend_logs_api_key_starttime_index"
+CALL_ID_INDEX_MIGRATION: Final = "20260831120001_spend_logs_litellm_call_id_index"
+API_KEY_INDEX: Final = "LiteLLM_SpendLogs_api_key_startTime_idx"
+CALL_ID_INDEX: Final = "LiteLLM_SpendLogs_litellm_call_id_idx"
+PARTITIONED_PARENT_ERROR: Final = 'cannot create index on partitioned table "LiteLLM_SpendLogs" concurrently'
+ORIGINAL_MIGRATION_SQL: Final = MappingProxyType(
+    {
+        API_KEY_INDEX_MIGRATION: (
+            "-- CreateIndex\n"
+            'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" '
+            'ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
+        ),
+        CALL_ID_INDEX_MIGRATION: (
+            "-- CreateIndex\n"
+            'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" '
+            'ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
+        ),
+    }
+)
+CALL_ID_INDEX_DEFINITION: Final = next(index for index in REQUEST_LOG_INDEXES if index.name == CALL_ID_INDEX)
+RELEASES: Final = pytest.mark.parametrize(
+    "release", (API_KEY_INDEX_MIGRATION, CALL_ID_INDEX_MIGRATION), ids=("v1.102.1", "v1.103.0")
+)
+PARTITIONS: Final = MappingProxyType(
+    {
+        "LiteLLM_SpendLogs_p2026_08": ("2026-08-01", "2026-09-01"),
+        "LiteLLM_SpendLogs_p2026_09": ("2026-09-01", "2026-10-01"),
+    }
+)
+DEFAULT_PARTITION: Final = "LiteLLM_SpendLogs_pdefault"
+ROWS_PER_PARTITION: Final = 200
+RESOLVERS: Final = pytest.mark.parametrize("use_v2_resolver", (True, False), ids=("v2", "v1"))
+
+
+def _base_url() -> str:
+    return os.environ["DATABASE_URL"].split("?")[0]
+
+
+def _migrate_deploy(database_url: str, schema: Path) -> "subprocess.CompletedProcess[str]":
+    return subprocess.run(
+        [sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(schema)],
+        capture_output=True,
+        text=True,
+        env={**os.environ, "DATABASE_URL": database_url},
+    )
+
+
+def _release_layout(prisma_dir: Path, before: str) -> Path:
+    """The shipped migrations older than `before`, with the two index migrations written
+    the way the releases that shipped them did: the Prisma layout of a proxy on that release."""
+    (prisma_dir / "migrations").mkdir(parents=True)
+    shutil.copy(PACKAGE / "schema.prisma", prisma_dir / "schema.prisma")
+    for migration in sorted((PACKAGE / "migrations").iterdir()):
+        if migration.is_dir() and migration.name < before:
+            shutil.copytree(migration, prisma_dir / "migrations" / migration.name)
+    for name, original in ORIGINAL_MIGRATION_SQL.items():
+        if (prisma_dir / "migrations" / name).is_dir():
+            (prisma_dir / "migrations" / name / "migration.sql").write_text(original)
+    return prisma_dir / "schema.prisma"
+
+
+def _deploy_release(database_url: str, prisma_dir: Path, before: str) -> None:
+    deployed: Final = _migrate_deploy(database_url, _release_layout(prisma_dir, before))
+    assert deployed.returncode == 0, deployed.stderr
+
+
+def _insert_spend_log(
+    conn: "psycopg.Connection[tuple[object, ...]]", request_id: str, day: str, table: str = "LiteLLM_SpendLogs"
+) -> None:
+    conn.execute(
+        sql.SQL(
+            'INSERT INTO {} ("request_id", "call_type", "startTime", "endTime", "api_key") VALUES (%s, %s, %s, %s, %s)'
+        ).format(sql.Identifier(table)),
+        (request_id, "acompletion", day, day, f"key-{request_id[-1]}"),
+    )
+
+
+def _partition_spend_logs(database_url: str) -> None:
+    with psycopg.connect(database_url, autocommit=True) as conn:
+        conn.execute(PARTITION_SCRIPT.read_bytes())
+        for partition, (start, stop) in PARTITIONS.items():
+            conn.execute(
+                sql.SQL('CREATE TABLE {} PARTITION OF "LiteLLM_SpendLogs" FOR VALUES FROM ({}) TO ({})').format(
+                    sql.Identifier(partition), sql.Literal(start), sql.Literal(stop)
+                )
+            )
+            for row in range(ROWS_PER_PARTITION):
+                _insert_spend_log(conn, f"{partition}-{row}", start)
+        for row in range(ROWS_PER_PARTITION):
+            _insert_spend_log(conn, f"default-{row}", "2020-01-01")
+
+
+@pytest.fixture
+def release() -> str:
+    """The first migration a database has not applied yet; the v1.103.0 shape unless a test parametrizes it."""
+    return CALL_ID_INDEX_MIGRATION
+
+
+@pytest.fixture
+def scratch_database(release: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Iterator[str]:
+    """A deployment stopped before `release`, with DATABASE_URL pointed at it so
+    ProxyExtrasDBManager upgrades it like a booting proxy."""
+    admin_url: Final = _base_url()
+    name: Final = f"spend_logs_index_{uuid.uuid4().hex[:8]}"
+    with psycopg.connect(admin_url, autocommit=True) as conn:
+        conn.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name)))
+    database_url: Final = f"{admin_url.rsplit('/', 1)[0]}/{name}"
+    try:
+        _deploy_release(database_url, tmp_path / "prisma", release)
+        monkeypatch.delenv("DIRECT_URL", raising=False)
+        monkeypatch.setenv("DATABASE_URL", database_url)
+        yield database_url
+    finally:
+        with psycopg.connect(admin_url, autocommit=True) as conn:
+            conn.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(name)))
+
+
+@pytest.fixture
+def partitioned_database(scratch_database: str) -> str:
+    _partition_spend_logs(scratch_database)
+    return scratch_database
+
+
+def _fail_the_call_id_migration_like_the_shipped_release(database_url: str, tmp_path: Path) -> None:
+    """Boot the original v1.103.0 layout once: its CONCURRENTLY statement fails on the
+    partitioned parent and leaves the call_id ledger row unfinished."""
+    failed: Final = _migrate_deploy(database_url, _release_layout(tmp_path / "v1.103.0", "99999999999999"))
+    assert failed.returncode != 0 and PARTITIONED_PARENT_ERROR in failed.stderr, failed.stderr
+    assert _ledger(database_url)[CALL_ID_INDEX_MIGRATION] == (False, False)
+
+
+@dataclass(frozen=True, slots=True)
+class _IndexRow:
+    name: str
+    valid: bool
+
+
+@dataclass(frozen=True, slots=True)
+class _AttachedRow:
+    table: str
+    index: str
+
+
+@dataclass(frozen=True, slots=True)
+class _LedgerRow:
+    name: str
+    finished: bool
+    rolled_back: bool
+
+
+@dataclass(frozen=True, slots=True)
+class _OidRow:
+    name: str
+    oid: int
+
+
+def _index_validity(database_url: str, suffix: str) -> Mapping[str, bool]:
+    """index name -> indisvalid for every index ending in `suffix` on the SpendLogs parent or one of its partitions."""
+    with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_IndexRow)) as cursor:
+        rows: Final = cursor.execute(
+            "SELECT c.relname AS name, i.indisvalid AS valid FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid "
+            "WHERE c.relname LIKE %s AND (i.indrelid = to_regclass('\"LiteLLM_SpendLogs\"') OR i.indrelid IN "
+            "(SELECT inhrelid FROM pg_inherits WHERE inhparent = to_regclass('\"LiteLLM_SpendLogs\"'))) "
+            "ORDER BY c.relname",
+            (f"%{suffix}",),
+        ).fetchall()
+    return MappingProxyType({row.name: row.valid for row in rows})
+
+
+def _attached_children(database_url: str, parent_index: str) -> frozenset[tuple[str, str]]:
+    """(partition, child index) pairs attached under the parent index."""
+    with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_AttachedRow)) as cursor:
+        rows: Final = cursor.execute(
+            'SELECT t.relname AS "table", c.relname AS index FROM pg_inherits i '
+            "JOIN pg_class c ON c.oid = i.inhrelid JOIN pg_index x ON x.indexrelid = c.oid "
+            "JOIN pg_class t ON t.oid = x.indrelid "
+            "WHERE i.inhparent = to_regclass(%s)",
+            (f'"{parent_index}"',),
+        ).fetchall()
+    return frozenset((row.table, row.index) for row in rows)
+
+
+@dataclass(frozen=True, slots=True)
+class _TableRow:
+    name: str
+
+
+def _indexed_table(database_url: str, index: str) -> "str | None":
+    with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_TableRow)) as cursor:
+        row: Final = cursor.execute(
+            "SELECT t.relname AS name FROM pg_index x JOIN pg_class t ON t.oid = x.indrelid "
+            "WHERE x.indexrelid = to_regclass(%s)",
+            (f'"{index}"',),
+        ).fetchone()
+    return None if row is None else row.name
+
+
+def _ledger(database_url: str) -> Mapping[str, tuple[bool, bool]]:
+    """migration name -> (finished, rolled back) for the newest ledger row of each migration."""
+    with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_LedgerRow)) as cursor:
+        rows: Final = cursor.execute(
+            "SELECT DISTINCT ON (migration_name) migration_name AS name, finished_at IS NOT NULL AS finished, "
+            "rolled_back_at IS NOT NULL AS rolled_back FROM _prisma_migrations ORDER BY migration_name, started_at DESC"
+        ).fetchall()
+    return MappingProxyType({row.name: (row.finished, row.rolled_back) for row in rows})
+
+
+def _index_oids(database_url: str) -> Mapping[str, int]:
+    """index name -> oid for every index on the SpendLogs parent or one of its partitions; a rebuild changes the oid."""
+    with psycopg.connect(database_url) as conn, conn.cursor(row_factory=class_row(_OidRow)) as cursor:
+        rows: Final = cursor.execute(
+            "SELECT c.relname AS name, c.oid::int AS oid FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid "
+            "WHERE i.indrelid = to_regclass('\"LiteLLM_SpendLogs\"') OR i.indrelid IN "
+            "(SELECT inhrelid FROM pg_inherits WHERE inhparent = to_regclass('\"LiteLLM_SpendLogs\"'))"
+        ).fetchall()
+    return MappingProxyType({row.name: row.oid for row in rows})
+
+
+def _migration_job(use_v2_resolver: bool) -> bool:
+    return ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver)
+
+
+def _assert_no_pending_migrations(database_url: str) -> None:
+    status: Final = _migrate_deploy(database_url, PACKAGE / "schema.prisma")
+    assert status.returncode == 0 and "No pending migrations" in status.stdout, status.stdout + status.stderr
+
+
+def _assert_every_ledger_row_is_finished(database_url: str) -> Mapping[str, tuple[bool, bool]]:
+    ledger: Final = _ledger(database_url)
+    assert ledger[API_KEY_INDEX_MIGRATION] == (True, False) and ledger[CALL_ID_INDEX_MIGRATION] == (True, False)
+    assert all(finished and not rolled_back for finished, rolled_back in ledger.values()), ledger
+    return ledger
+
+
+def _expected_children(partitions: tuple[str, ...], suffix: str) -> frozenset[tuple[str, str]]:
+    return frozenset((partition, f"{partition}_{suffix}") for partition in partitions)
+
+
+def _assert_index_covers_every_partition(database_url: str, parent_index: str, suffix: str) -> None:
+    partitions: Final = (*PARTITIONS, DEFAULT_PARTITION)
+    assert _index_validity(database_url, suffix) == {parent_index: True} | {f"{p}_{suffix}": True for p in partitions}
+    assert _attached_children(database_url, parent_index) == _expected_children(partitions, suffix)
+
+
+@requires_db
+@RESOLVERS
+@RELEASES
+def test_a_partitioned_spend_logs_upgrade_builds_both_indexes_per_partition_and_a_rerun_is_idempotent(
+    partitioned_database: str, use_v2_resolver: bool
+) -> None:
+    assert _migration_job(use_v2_resolver) is True
+
+    _assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
+    _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
+    ledger: Final = _assert_every_ledger_row_is_finished(partitioned_database)
+    _assert_no_pending_migrations(partitioned_database)
+    oids: Final = _index_oids(partitioned_database)
+
+    with psycopg.connect(partitioned_database, autocommit=True) as conn:
+        conn.execute(
+            'CREATE TABLE "LiteLLM_SpendLogs_p2026_10" PARTITION OF "LiteLLM_SpendLogs" '
+            "FOR VALUES FROM ('2026-10-01') TO ('2026-11-01')"
+        )
+    inherited: Final = frozenset(
+        ("LiteLLM_SpendLogs_p2026_10", f"LiteLLM_SpendLogs_p2026_10_{suffix}")
+        for suffix in ("api_key_startTime_idx", "litellm_call_id_idx")
+    )
+    attached: Final = _attached_children(partitioned_database, API_KEY_INDEX) | _attached_children(
+        partitioned_database, CALL_ID_INDEX
+    )
+    assert inherited <= attached, attached
+
+    assert _migration_job(use_v2_resolver) is True
+    assert _ledger(partitioned_database) == ledger
+    assert {name: oid for name, oid in _index_oids(partitioned_database).items() if name in oids} == oids
+
+
+@requires_db
+@RESOLVERS
+@RELEASES
+def test_a_plain_spend_logs_upgrade_builds_both_indexes_and_a_second_job_run_rebuilds_nothing(
+    scratch_database: str, use_v2_resolver: bool
+) -> None:
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        for row in range(ROWS_PER_PARTITION):
+            _insert_spend_log(conn, f"flat-{row}", "2026-09-01")
+
+    assert _migration_job(use_v2_resolver) is True
+
+    assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+    _assert_every_ledger_row_is_finished(scratch_database)
+    _assert_no_pending_migrations(scratch_database)
+    oids: Final = _index_oids(scratch_database)
+
+    assert _migration_job(use_v2_resolver) is True
+    assert _index_oids(scratch_database) == oids
+
+
+@requires_db
+@RESOLVERS
+def test_a_database_that_applied_the_original_migration_files_sees_no_pending_migrations_and_no_rebuild(
+    scratch_database: str, use_v2_resolver: bool, tmp_path: Path
+) -> None:
+    """A plain table upgraded on v1.103.0 applied both original files. The inert files in
+    this build must neither re-run nor fail those rows, and the migration job must keep the
+    indexes the migrations built."""
+    deployed: Final = _migrate_deploy(scratch_database, _release_layout(tmp_path / "v1.103.0", "99999999999999"))
+    assert deployed.returncode == 0, deployed.stderr
+    before: Final = _ledger(scratch_database)
+    assert before[API_KEY_INDEX_MIGRATION] == (True, False) and before[CALL_ID_INDEX_MIGRATION] == (True, False)
+    oids: Final = _index_oids(scratch_database)
+    assert {API_KEY_INDEX, CALL_ID_INDEX} <= set(oids)
+
+    _assert_no_pending_migrations(scratch_database)
+    assert _migration_job(use_v2_resolver) is True
+
+    assert _ledger(scratch_database) == before
+    assert _index_oids(scratch_database) == oids
+
+
+@requires_db
+@RESOLVERS
+def test_a_failed_call_id_ledger_row_from_a_v1_103_boot_is_rolled_back_and_the_inert_file_applied(
+    partitioned_database: str, use_v2_resolver: bool, tmp_path: Path
+) -> None:
+    _fail_the_call_id_migration_like_the_shipped_release(partitioned_database, tmp_path)
+
+    assert _migration_job(use_v2_resolver) is True
+
+    _assert_every_ledger_row_is_finished(partitioned_database)
+    _assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
+    _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
+    _assert_no_pending_migrations(partitioned_database)
+    with psycopg.connect(partitioned_database) as conn:
+        rows: Final = conn.execute(
+            "SELECT finished_at IS NOT NULL, rolled_back_at IS NOT NULL FROM _prisma_migrations "
+            "WHERE migration_name = %s ORDER BY started_at",
+            (CALL_ID_INDEX_MIGRATION,),
+        ).fetchall()
+    assert rows == [(False, True), (True, False)], rows
+
+
+@requires_db
+def test_a_failed_row_whose_migration_still_runs_sql_in_this_build_is_left_for_the_operator(
+    partitioned_database: str, tmp_path: Path
+) -> None:
+    _fail_the_call_id_migration_like_the_shipped_release(partitioned_database, tmp_path)
+    still_building: Final = tmp_path / "edited" / CALL_ID_INDEX_MIGRATION / "migration.sql"
+    still_building.parent.mkdir(parents=True)
+    still_building.write_text(ORIGINAL_MIGRATION_SQL[CALL_ID_INDEX_MIGRATION])
+
+    with migration_lock(partitioned_database) as coordinator:
+        assert roll_back_failed_inert_migration(coordinator, "public", still_building) is False
+
+    assert _ledger(partitioned_database)[CALL_ID_INDEX_MIGRATION] == (False, False)
+
+
+@requires_db
+def test_a_migration_without_a_failed_row_is_not_touched(partitioned_database: str) -> None:
+    inert: Final = PACKAGE / "migrations" / CALL_ID_INDEX_MIGRATION / "migration.sql"
+    before: Final = _ledger(partitioned_database)
+
+    with migration_lock(partitioned_database) as coordinator:
+        assert roll_back_failed_inert_migration(coordinator, "public", inert) is False
+
+    assert _ledger(partitioned_database) == before
+
+
+def _pin_a_snapshot_on(database_url: str, table: str) -> "psycopg.Connection[tuple[object, ...]]":
+    pin: Final = psycopg.connect(database_url)
+    pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
+    pin.execute(sql.SQL("SELECT count(*) FROM {}").format(sql.Identifier(table)))
+    return pin
+
+
+def _leave_an_invalid_index(database_url: str, name: str, table: str, column: str) -> None:
+    with _pin_a_snapshot_on(database_url, table):
+        with psycopg.connect(database_url, autocommit=True) as builder:
+            builder.execute("SET statement_timeout = '1s'")
+            with pytest.raises(psycopg.errors.QueryCanceled):
+                builder.execute(
+                    sql.SQL("CREATE INDEX CONCURRENTLY {} ON {} ({})").format(
+                        sql.Identifier(name), sql.Identifier(table), sql.Identifier(column)
+                    )
+                )
+
+
+@requires_db
+def test_an_invalid_index_of_the_managed_name_on_a_plain_table_is_rebuilt(scratch_database: str) -> None:
+    _leave_an_invalid_index(scratch_database, CALL_ID_INDEX, "LiteLLM_SpendLogs", "litellm_call_id")
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: False}
+
+    assert ensure_request_log_indexes(scratch_database, "public") is True
+
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+
+
+def _rebuild_as_another_replica(database_url: str, name: str, table: str, column: str) -> int:
+    """Drop and rebuild the index from a second connection, as a replica that won the
+    race would, and return the oid of the index it built."""
+    with psycopg.connect(database_url, autocommit=True) as other_replica:
+        other_replica.execute(sql.SQL("DROP INDEX {}").format(sql.Identifier(name)))
+        other_replica.execute(
+            sql.SQL("CREATE INDEX {} ON {} ({})").format(
+                sql.Identifier(name), sql.Identifier(table), sql.Identifier(column)
+            )
+        )
+    return _index_oids(database_url)[name]
+
+
+def _connecting_with_another_replica_acting_first(
+    statement: str, other_replica: Callable[[QueryNoTemplate], None]
+) -> Callable[[str], "psycopg.Connection[tuple[object, ...]]"]:
+    """A connect function whose cursors let `other_replica` act, once, right before the
+    first statement containing `statement` runs: the interleaving two replicas booting
+    together can produce, made deterministic."""
+    raced: Final = threading.Event()
+
+    class _RacedCursor(psycopg.Cursor[tuple[object, ...]]):
+        def execute(  # pyright: ignore[reportIncompatibleMethodOverride]  # the builder never runs a Template query
+            self,
+            query: QueryNoTemplate,
+            params: "Params | None" = None,
+            *,
+            prepare: "bool | None" = None,
+            binary: "bool | None" = None,
+        ) -> "_RacedCursor":
+            text: Final = query.as_string(self.connection) if isinstance(query, sql.Composable) else query
+            if isinstance(text, str) and statement in text and not raced.is_set():
+                raced.set()
+                other_replica(query)
+            return super().execute(query, params, prepare=prepare, binary=binary)
+
+    def connect(database_url: str) -> "psycopg.Connection[tuple[object, ...]]":
+        return psycopg.connect(database_url, autocommit=True, cursor_factory=_RacedCursor)
+
+    return connect
+
+
+@requires_db
+def test_an_index_another_replica_made_valid_before_the_lock_was_taken_is_kept(scratch_database: str) -> None:
+    """Two replicas boot against the same invalid index. The one that takes the lock
+    second must read the catalog again under it, or it drops the valid index the first
+    one just finished and starts the whole build over."""
+    _leave_an_invalid_index(scratch_database, CALL_ID_INDEX, "LiteLLM_SpendLogs", "litellm_call_id")
+    theirs: Final[queue.SimpleQueue[int]] = queue.SimpleQueue()
+    connect: Final = _connecting_with_another_replica_acting_first(
+        "pg_try_advisory_lock",
+        lambda _: theirs.put(
+            _rebuild_as_another_replica(scratch_database, CALL_ID_INDEX, "LiteLLM_SpendLogs", "litellm_call_id")
+        ),
+    )
+
+    assert ensure_request_log_indexes(scratch_database, "public", (CALL_ID_INDEX_DEFINITION,), connect) is True
+
+    assert _index_oids(scratch_database)[CALL_ID_INDEX] == theirs.get_nowait()
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+
+
+@requires_db
+def test_a_child_index_another_replica_attached_first_is_not_attached_twice(partitioned_database: str) -> None:
+    """A replica that reaches the attach step after another one attached the same child
+    relies on ATTACH PARTITION being a no-op for an index already under that parent
+    (PostgreSQL 14 ALTER INDEX, ATExecAttachPartitionIdx, checked 2026-10-01); this test
+    is where that would surface if a future version or a code change made it an error."""
+
+    def attach_as_another_replica(statement: QueryNoTemplate) -> None:
+        with psycopg.connect(partitioned_database, autocommit=True) as other_replica:
+            other_replica.execute(statement)
+
+    connect: Final = _connecting_with_another_replica_acting_first("ATTACH PARTITION", attach_as_another_replica)
+
+    assert ensure_request_log_indexes(partitioned_database, "public", (CALL_ID_INDEX_DEFINITION,), connect) is True
+
+    _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
+
+
+@requires_db
+def test_an_invalid_child_index_left_by_an_interrupted_build_is_rebuilt_and_attached(
+    partitioned_database: str,
+) -> None:
+    partition: Final = "LiteLLM_SpendLogs_p2026_08"
+    child: Final = f"{partition}_litellm_call_id_idx"
+    _leave_an_invalid_index(partitioned_database, child, partition, "litellm_call_id")
+    assert _index_validity(partitioned_database, "litellm_call_id_idx") == {child: False}
+
+    assert ensure_request_log_indexes(partitioned_database, "public") is True
+
+    _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
+
+
+@requires_db
+def test_an_index_of_that_name_on_another_table_is_left_alone_and_reported(partitioned_database: str) -> None:
+    with psycopg.connect(partitioned_database, autocommit=True) as conn:
+        conn.execute(f'CREATE INDEX "{CALL_ID_INDEX}" ON "LiteLLM_ErrorLogs" ("request_id")')
+
+    assert ensure_request_log_indexes(partitioned_database, "public") is False
+
+    assert _index_validity(partitioned_database, "litellm_call_id_idx") == {}
+    assert _indexed_table(partitioned_database, CALL_ID_INDEX) == "LiteLLM_ErrorLogs"
+    _assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
+
+
+@requires_db
+def test_an_invalid_index_of_a_child_name_on_another_table_is_not_dropped(partitioned_database: str) -> None:
+    child: Final = "LiteLLM_SpendLogs_p2026_08_litellm_call_id_idx"
+    _leave_an_invalid_index(partitioned_database, child, "LiteLLM_ErrorLogs", "request_id")
+
+    assert ensure_request_log_indexes(partitioned_database, "public") is False
+
+    assert _indexed_table(partitioned_database, child) == "LiteLLM_ErrorLogs"
+    assert _attached_children(partitioned_database, CALL_ID_INDEX) == frozenset()
+
+
+@requires_db
+def test_a_process_holding_the_migration_lock_makes_the_build_wait_for_the_next_job_run(scratch_database: str) -> None:
+    with psycopg.connect(scratch_database, autocommit=True) as other_replica:
+        other_replica.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
+        assert ensure_request_log_indexes(scratch_database, "public") is False
+        assert _index_validity(scratch_database, "litellm_call_id_idx") == {}
+
+    assert ensure_request_log_indexes(scratch_database, "public") is True
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+
+
+@requires_db
+@RESOLVERS
+def test_a_migration_job_that_could_not_build_the_indexes_reports_failure_and_succeeds_when_rerun(
+    scratch_database: str, use_v2_resolver: bool
+) -> None:
+    """The migration job waits for the build and exits by run_migration_job's result; a job
+    that exits 0 with the indexes missing would leave the table unindexed until the next
+    deploy or until a serving proxy's background build gets to them."""
+    with psycopg.connect(scratch_database, autocommit=True) as other_replica:
+        other_replica.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,))
+        assert _migration_job(use_v2_resolver) is False
+        _assert_every_ledger_row_is_finished(scratch_database)
+        assert _index_validity(scratch_database, "litellm_call_id_idx") == {}
+
+    assert _migration_job(use_v2_resolver) is True
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+    assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
+
+
+@requires_db
+@RESOLVERS
+def test_the_serving_proxy_setup_applies_the_inert_migrations_and_builds_no_index(
+    partitioned_database: str, use_v2_resolver: bool
+) -> None:
+    """setup_database alone applies the inert files and builds nothing, so a serving proxy's
+    readiness is never held up by an index build; the build it starts afterwards, or the
+    migration job, is what puts the indexes in place."""
+    api_key_index_before: Final = _index_validity(partitioned_database, "api_key_startTime_idx")
+    assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True
+
+    _assert_every_ledger_row_is_finished(partitioned_database)
+    _assert_no_pending_migrations(partitioned_database)
+    assert _index_validity(partitioned_database, "litellm_call_id_idx") == {}
+    assert _index_validity(partitioned_database, "api_key_startTime_idx") == api_key_index_before
+
+    assert _migration_job(use_v2_resolver) is True
+    _assert_index_covers_every_partition(partitioned_database, API_KEY_INDEX, "api_key_startTime_idx")
+    _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
+
+
+@requires_db
+def test_a_role_that_may_not_create_indexes_is_logged_and_left_for_the_next_job_run(
+    scratch_database: str, caplog: pytest.LogCaptureFixture
+) -> None:
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC")
+        conn.execute("CREATE ROLE spend_logs_reader LOGIN PASSWORD 'reader'")
+        conn.execute("GRANT USAGE ON SCHEMA public TO spend_logs_reader")
+        conn.execute('GRANT SELECT ON "LiteLLM_SpendLogs" TO spend_logs_reader')
+    reader_url: Final = scratch_database.replace("postgres:postgres@", "spend_logs_reader:reader@", 1)
+    try:
+        with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
+            assert ensure_request_log_indexes(reader_url, "public") is False
+    finally:
+        with psycopg.connect(scratch_database, autocommit=True) as conn:
+            conn.execute("DROP OWNED BY spend_logs_reader")
+            conn.execute("DROP ROLE spend_logs_reader")
+    assert "leaving them for the next index build" in caplog.text
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {}
+
+
+@requires_db
+def test_inserts_keep_flowing_while_the_partition_indexes_build(partitioned_database: str) -> None:
+    """With a write open on one partition, the parent index goes on ONLY the parent and
+    the CONCURRENTLY child build waits for that write without blocking new INSERTs. A
+    plain CREATE INDEX on the parent would wait for the same write while holding SHARE
+    on the parent, queueing every new INSERT behind it."""
+    outcome: Final[list[bool]] = []  # mutable-ok: the builder thread hands its result back through it
+    with psycopg.connect(partitioned_database) as writer:
+        _insert_spend_log(writer, "LiteLLM_SpendLogs_p2026_08-open", "2026-08-15", table="LiteLLM_SpendLogs_p2026_08")
+        builder_thread: Final = threading.Thread(
+            target=lambda: outcome.append(_build_in_its_own_session(partitioned_database, CALL_ID_INDEX_DEFINITION))
+        )
+        builder_thread.start()
+        try:
+            _wait_until_the_build_is_waiting(partitioned_database)
+            with psycopg.connect(partitioned_database, autocommit=True) as late_writer:
+                late_writer.execute("SET lock_timeout = '1s'")
+                _insert_spend_log(late_writer, "LiteLLM_SpendLogs_p2026_08-late", "2026-08-16")
+        finally:
+            writer.commit()
+            builder_thread.join()
+    assert outcome == [True]
+    _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
+
+
+def _build_in_its_own_session(database_url: str, index: RequestLogIndex) -> bool:
+    with psycopg.connect(database_url, autocommit=True) as builder:
+        return build_index_on_partitioned_table(builder, "public", index)
+
+
+def _wait_until_the_build_is_waiting(database_url: str) -> None:
+    deadline: Final = time.monotonic() + 30
+    with psycopg.connect(database_url, autocommit=True) as conn:
+        while time.monotonic() < deadline:
+            waiting = conn.execute(
+                "SELECT 1 FROM pg_stat_activity WHERE query LIKE 'CREATE INDEX%' AND wait_event_type IS NOT NULL"
+            ).fetchone()
+            if waiting is not None:
+                return
+            time.sleep(0.05)
+    pytest.fail("the partition index build never started waiting on the open write")
+
+
+def _create_index(database_url: str, name: str, table: str, columns: str) -> int:
+    """Create a plain index by hand, the way an operator's workaround would, and return its oid."""
+    with psycopg.connect(database_url, autocommit=True) as conn:
+        conn.execute(
+            sql.SQL("CREATE INDEX {} ON {} {}").format(sql.Identifier(name), sql.Identifier(table), sql.SQL(columns))
+        )
+    return _index_oids(database_url)[name]
+
+
+@requires_db
+def test_a_valid_index_of_the_same_definition_under_another_name_is_renamed_instead_of_rebuilt(
+    scratch_database: str,
+) -> None:
+    hand_built: Final = _create_index(scratch_database, "call_id_by_hand", "LiteLLM_SpendLogs", '("litellm_call_id")')
+
+    assert ensure_request_log_indexes(scratch_database, "public") is True
+
+    oids: Final = _index_oids(scratch_database)
+    assert "call_id_by_hand" not in oids and oids[CALL_ID_INDEX] == hand_built
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+
+
+@requires_db
+def test_a_hand_built_child_index_under_another_name_is_renamed_and_attached(partitioned_database: str) -> None:
+    partition: Final = "LiteLLM_SpendLogs_p2026_08"
+    hand_built: Final = _create_index(
+        partitioned_database, "p2026_08_call_id_by_hand", partition, '("litellm_call_id")'
+    )
+
+    assert ensure_request_log_indexes(partitioned_database, "public") is True
+
+    _assert_index_covers_every_partition(partitioned_database, CALL_ID_INDEX, "litellm_call_id_idx")
+    oids: Final = _index_oids(partitioned_database)
+    assert "p2026_08_call_id_by_hand" not in oids and oids[f"{partition}_litellm_call_id_idx"] == hand_built
+
+
+@requires_db
+def test_an_index_with_another_definition_is_not_taken_for_the_managed_one(scratch_database: str) -> None:
+    with psycopg.connect(scratch_database, autocommit=True) as conn:
+        conn.execute(sql.SQL("DROP INDEX {}").format(sql.Identifier(API_KEY_INDEX)))
+    others: Final = {
+        "time_then_key": _create_index(
+            scratch_database, "time_then_key", "LiteLLM_SpendLogs", '("startTime", "api_key")'
+        ),
+        "call_id_desc": _create_index(
+            scratch_database, "call_id_desc", "LiteLLM_SpendLogs", '("litellm_call_id" DESC)'
+        ),
+        "call_id_then_key": _create_index(
+            scratch_database, "call_id_then_key", "LiteLLM_SpendLogs", '("litellm_call_id", "api_key")'
+        ),
+        "call_id_pattern": _create_index(
+            scratch_database, "call_id_pattern", "LiteLLM_SpendLogs", '("litellm_call_id" text_pattern_ops)'
+        ),
+    }
+
+    assert ensure_request_log_indexes(scratch_database, "public") is True
+
+    oids: Final = _index_oids(scratch_database)
+    assert {name: oids[name] for name in others} == others
+    assert _index_validity(scratch_database, "litellm_call_id_idx") == {CALL_ID_INDEX: True}
+    assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True}
+
+
+@requires_db
+def test_a_valid_partitioned_parent_index_under_another_name_is_renamed_with_its_children_kept(
+    partitioned_database: str, caplog: pytest.LogCaptureFixture
+) -> None:
+    hand_built: Final = _create_index(
+        partitioned_database, "call_id_parent_by_hand", "LiteLLM_SpendLogs", '("litellm_call_id")'
+    )
+    children_before: Final = _attached_children(partitioned_database, "call_id_parent_by_hand")
+
+    with caplog.at_level("INFO", logger="litellm_proxy_extras"):
+        assert ensure_request_log_indexes(partitioned_database, "public") is True
+
+    assert "Building index" not in caplog.text
+    oids: Final = _index_oids(partitioned_database)
+    assert "call_id_parent_by_hand" not in oids and oids[CALL_ID_INDEX] == hand_built
+    assert _attached_children(partitioned_database, CALL_ID_INDEX) == children_before
+    assert _index_validity(partitioned_database, "litellm_call_id_idx")[CALL_ID_INDEX] is True
+
+
+@requires_db
+def test_a_second_copy_of_a_managed_index_is_reported_with_its_drop_statement_and_left_in_place(
+    scratch_database: str, caplog: pytest.LogCaptureFixture
+) -> None:
+    assert ensure_request_log_indexes(scratch_database, "public") is True
+    copy: Final = _create_index(scratch_database, "call_id_copy", "LiteLLM_SpendLogs", '("litellm_call_id")')
+
+    with caplog.at_level("WARNING", logger="litellm_proxy_extras"):
+        assert ensure_request_log_indexes(scratch_database, "public") is True
+
+    assert 'remove it with: DROP INDEX CONCURRENTLY "public"."call_id_copy"' in caplog.text
+    assert _index_oids(scratch_database)["call_id_copy"] == copy

tests/unit/harness/core_fakes.py

@@ -0,0 +1,246 @@
+"""Fake handler/config, sandbox and endpoint shared by the core runtime tests."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections.abc import AsyncIterator, Callable
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import pytest
+
+from litellm.harness import runtime
+from litellm.harness.context import SessionContext
+from litellm.harness.handlers.base import BaseHarnessHandler
+from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig
+from litellm.harness.options import ClaudeCodeOptions
+from litellm.harness.sandbox.base import CompletedRun
+from litellm.harness.sandbox.snapshot import snapshot_local
+from litellm.harness.types import (
+    Approval,
+    Capabilities,
+    Event,
+    Harness,
+    Text,
+    ToolCall,
+    ToolResult,
+)
+
+ALL_MODES = frozenset({"read-only", "ask", "edit", "full"})
+FULL_CAPS = Capabilities(
+    structured_output=True,
+    tool_approval=True,
+    tool_filtering=True,
+    history=True,
+    custom_tools=True,
+    skills=True,
+    resume=True,
+    permission_modes=ALL_MODES,
+)
+NARROW_CAPS = Capabilities(
+    structured_output=False,
+    tool_approval=False,
+    tool_filtering=False,
+    history=False,
+    custom_tools=False,
+    skills=False,
+    resume=False,
+    permission_modes=frozenset({"read-only", "full"}),
+)
+
+
+class FakeConfig(BaseHarnessConfig):
+    """Declares the fake harness; per-test subclasses override capabilities."""
+
+    harness: ClassVar[Harness] = Harness.CLAUDE_CODE
+    options_type: ClassVar[type] = ClaudeCodeOptions
+    capabilities: ClassVar[Capabilities] = FULL_CAPS
+    uses_model_endpoint: ClassVar[bool] = True
+
+
+Script = Callable[["FakeAdapter", SessionContext, str], AsyncIterator[Event]]
+
+
+class FakeSandbox:
+    """A LocalSandbox-like object over a temp dir; no subprocesses."""
+
+    def __init__(self, workdir: str) -> None:
+        self.workdir = workdir
+        self.closed = False
+
+    def _path(self, path: str) -> str:
+        return path if os.path.isabs(path) else os.path.join(self.workdir, path)
+
+    async def exec(self, cmd: list[str], *, env: Any = None, cwd: Any = None) -> Any:
+        raise NotImplementedError
+
+    async def run(
+        self, cmd: list[str], *, env: Any = None, cwd: Any = None, timeout: Any = None
+    ) -> CompletedRun:
+        return CompletedRun(stdout="", stderr="", exit_code=0)
+
+    async def read(self, path: str) -> bytes:
+        with open(self._path(path), "rb") as fh:
+            return fh.read()
+
+    async def write(self, path: str, data: bytes) -> None:
+        full = self._path(path)
+        os.makedirs(os.path.dirname(full), exist_ok=True)
+        with open(full, "wb") as fh:
+            fh.write(data)
+
+    def host_url(self, port: int) -> str:
+        return f"http://127.0.0.1:{port}"
+
+    async def which(self, binary: str) -> str | None:
+        return None
+
+    async def snapshot(self) -> dict[str, str]:
+        return await snapshot_local(self.workdir)
+
+    async def close(self) -> None:
+        self.closed = True
+
+
+@dataclass
+class FakeUsage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cost: float = 0.0
+    calls: int = 0
+
+    def add(self, input_tokens: int, output_tokens: int, cost: float) -> None:
+        self.input_tokens += input_tokens
+        self.output_tokens += output_tokens
+        self.cost += cost
+        self.calls += 1
+
+
+class FakeEndpoint:
+    """Stands in for ModelEndpoint; records every instance."""
+
+    instances: ClassVar[list[FakeEndpoint]] = []
+
+    def __init__(self, harness: Harness, model: Any, gateway: Any, **kwargs: Any):
+        self.harness = harness
+        self.model = model
+        self.gateway = gateway
+        self.kwargs = kwargs
+        self.usage = FakeUsage()
+        self.url = "http://127.0.0.1:1"
+        self.token = "tok"
+        self.entered = False
+        self.exited = False
+        FakeEndpoint.instances.append(self)
+
+    async def __aenter__(self) -> FakeEndpoint:
+        self.entered = True
+        return self
+
+    async def __aexit__(self, *exc_info: object) -> None:
+        self.exited = True
+
+
+async def script_hello(
+    adapter: FakeAdapter, ctx: SessionContext, prompt: str
+) -> AsyncIterator[Event]:
+    yield Text("hello ")
+    yield ToolCall(id="t1", name="bash", native_name="Bash", input={"cmd": "ls"})
+    yield ToolResult(id="t1", output="a.txt")
+    yield Text("world")
+    if ctx.endpoint is not None:
+        ctx.endpoint.usage.add(10, 5, 0.25)
+    else:
+        ctx.input_tokens += 10
+        ctx.output_tokens += 5
+        ctx.cost += 0.25
+        ctx.calls += 1
+
+
+class FakeAdapter(BaseHarnessHandler):
+    """Configurable adapter; subclass per test and set `script` / `caps`."""
+
+    harness: ClassVar[Harness] = Harness.CLAUDE_CODE
+    options_type: ClassVar[type] = ClaudeCodeOptions
+    capabilities: ClassVar[Capabilities] = FULL_CAPS
+    uses_endpoint: ClassVar[bool] = True
+    script: ClassVar[Script] = script_hello
+    instances: ClassVar[list[FakeAdapter]] = []
+
+    def __init__(self, config: BaseHarnessConfig | None = None) -> None:
+        self.config = config if config is not None else FakeConfig()
+        self.calls: list[str] = []
+        self.prompts: list[str] = []
+        self.resumed_with: str | None = None
+        self.approvals: list[tuple[bool, str]] = []
+        type(self).instances.append(self)
+
+    async def start(self, ctx: SessionContext) -> None:
+        self.calls.append("start")
+
+    async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
+        self.calls.append("turn")
+        self.prompts.append(prompt)
+        async for event in type(self).script(self, ctx, prompt):
+            yield event
+
+    async def stop(self, ctx: SessionContext) -> None:
+        self.calls.append("stop")
+
+    def native_session_id(self) -> str | None:
+        return "native-123"
+
+    async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
+        self.calls.append("resume")
+        self.resumed_with = native_session_id
+
+    async def history(self, ctx: SessionContext) -> list[dict[str, Any]]:
+        return [{"role": "user", "content": p} for p in self.prompts]
+
+
+async def script_approval(
+    adapter: FakeAdapter, ctx: SessionContext, prompt: str
+) -> AsyncIterator[Event]:
+    approval = Approval(tool="bash", input={"cmd": "rm"})
+    yield approval
+    decision = await approval.wait()
+    adapter.approvals.append(decision)
+    yield Text("allowed" if decision[0] else "denied")
+
+
+def install_adapter(
+    monkeypatch: pytest.MonkeyPatch,
+    script: Script = script_hello,
+    caps: Capabilities = FULL_CAPS,
+    uses_endpoint: bool = True,
+) -> type[FakeAdapter]:
+    """Register a FakeAdapter subclass for every harness and fake the endpoint."""
+    adapter_cls = type(
+        "TestAdapter",
+        (FakeAdapter,),
+        {
+            "script": staticmethod(script),
+            "capabilities": caps,
+            "uses_endpoint": uses_endpoint,
+            "instances": [],
+        },
+    )
+    config_cls = type(
+        "TestConfig",
+        (FakeConfig,),
+        {"capabilities": caps, "uses_model_endpoint": uses_endpoint},
+    )
+    monkeypatch.setattr(runtime, "get_harness_config", lambda harness: config_cls())
+    monkeypatch.setattr(
+        runtime, "get_harness_handler", lambda config: adapter_cls(config)
+    )
+    monkeypatch.setattr(runtime, "ModelEndpoint", FakeEndpoint)
+    monkeypatch.delenv("LITELLM_PROXY_API_BASE", raising=False)
+    monkeypatch.delenv("LITELLM_PROXY_API_KEY", raising=False)
+    FakeEndpoint.instances = []
+    return adapter_cls
+
+
+async def wait_forever() -> None:
+    await asyncio.Event().wait()

tests/unit/harness/handlers/test_deepagents_handler.py

@@ -0,0 +1,378 @@
+import asyncio
+import builtins
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import pytest
+from pydantic import BaseModel
+
+from litellm.harness.context import GatewayTarget, SessionContext
+from litellm.harness.errors import HarnessError, HarnessInstallFailed
+from litellm.harness.handlers import deepagents_handler as dh
+from litellm.harness.sandbox.local import LocalSandbox
+from litellm.harness.types import Approval, Harness, Text, ToolCall, ToolResult
+from litellm.llms.deepagents.harness.transformation import DeepAgentsHarnessConfig
+
+pytest.importorskip("deepagents")
+pytest.importorskip("langchain_litellm")
+
+from langchain_core.language_models.fake_chat_models import (  # noqa: E402
+    FakeMessagesListChatModel,
+)
+from langchain_core.messages import AIMessage  # noqa: E402
+
+from litellm.llms.deepagents.harness.sandbox_backend import (  # noqa: E402
+    SandboxBackend,
+    message_cost,
+)
+
+USAGE = {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}
+
+
+class FakeToolModel(FakeMessagesListChatModel):
+    """Canned responses; records the tool names bound on each call."""
+
+    bound: list = []
+
+    def bind_tools(self, tools: Any, **kwargs: Any) -> "FakeToolModel":
+        names = [getattr(t, "name", None) or t.get("name") for t in tools]
+        self.bound.append(sorted(n for n in names if n))
+        return self
+
+
+def tool_call(name: str, args: dict, call_id: str) -> AIMessage:
+    return AIMessage(
+        content="",
+        tool_calls=[{"name": name, "args": args, "id": call_id}],
+        usage_metadata=USAGE,
+    )
+
+
+def final(text: str) -> AIMessage:
+    return AIMessage(content=text, usage_metadata=USAGE)
+
+
+@pytest.fixture
+def fake_model(monkeypatch: pytest.MonkeyPatch):
+    def install(responses: list) -> FakeToolModel:
+        model = FakeToolModel(responses=responses, bound=[])
+        monkeypatch.setattr(dh, "build_chat_model", lambda ctx, deps: model)
+        return model
+
+    return install
+
+
+def make_ctx(tmp_path: Path, **kwargs: Any) -> SessionContext:
+    base: dict[str, Any] = {
+        "harness": Harness.DEEPAGENTS,
+        "sandbox": LocalSandbox(tmp_path),
+        "session_id": f"s-{os.urandom(4).hex()}",
+        "model": "gpt-4o-mini",
+    }
+    return SessionContext(**{**base, **kwargs})
+
+
+def make_handler() -> dh.DeepAgentsHandler:
+    return dh.DeepAgentsHandler(DeepAgentsHarnessConfig())
+
+
+async def started(ctx: SessionContext) -> dh.DeepAgentsHandler:
+    handler = make_handler()
+    await handler.start(ctx)
+    return handler
+
+
+async def run_turn(
+    handler: dh.DeepAgentsHandler,
+    ctx: SessionContext,
+    prompt: str,
+    approve: bool = True,
+) -> list:
+    events = []
+    async for event in handler.turn(ctx, prompt):
+        events.append(event)
+        if isinstance(event, Approval):
+            event.allow() if approve else event.deny("no")
+    return events
+
+
+async def test_write_then_read_events_and_file(tmp_path: Path, fake_model) -> None:
+    fake_model(
+        [
+            tool_call("write_file", {"file_path": "/hello.txt", "content": "hi"}, "c1"),
+            tool_call("read_file", {"file_path": "/hello.txt"}, "c2"),
+            final("done"),
+        ]
+    )
+    ctx = make_ctx(tmp_path)
+    handler = await started(ctx)
+    events = await run_turn(handler, ctx, "write hello.txt with hi")
+
+    calls = [e for e in events if isinstance(e, ToolCall)]
+    results = [e for e in events if isinstance(e, ToolResult)]
+    assert [(c.name, c.native_name, c.builtin) for c in calls] == [
+        ("write", "write_file", True),
+        ("read", "read_file", True),
+    ]
+    assert [r.id for r in results] == ["c1", "c2"]
+    assert "hi" in results[1].output
+    assert not any(r.is_error for r in results)
+    assert "done" in "".join(e.delta for e in events if isinstance(e, Text))
+    assert (tmp_path / "hello.txt").read_text() == "hi"
+    assert ctx.final_text == "done"
+    assert (ctx.input_tokens, ctx.output_tokens, ctx.calls) == (30, 15, 3)
+    assert ctx.cost > 0
+    history = await handler.history(ctx)
+    assert history[0] == {"role": "user", "content": "write hello.txt with hi"}
+    assert history[-1]["content"] == "done"
+
+
+async def test_read_only_hides_write_tools(tmp_path: Path, fake_model) -> None:
+    model = fake_model(
+        [
+            tool_call("write_file", {"file_path": "/x.txt", "content": "no"}, "c1"),
+            final("ok"),
+        ]
+    )
+    ctx = make_ctx(tmp_path, permissions="read-only")
+    handler = await started(ctx)
+    events = await run_turn(handler, ctx, "try to write")
+
+    first = model.bound[0]
+    assert "read_file" in first and "ls" in first
+    assert not {"write_file", "edit_file", "execute", "delete"} & set(first)
+    result = next(e for e in events if isinstance(e, ToolResult))
+    assert result.is_error
+    assert not (tmp_path / "x.txt").exists()
+
+
+async def test_disable_tools_uses_normalized_names(tmp_path: Path, fake_model) -> None:
+    model = fake_model([final("ok")])
+    ctx = make_ctx(tmp_path, disable_tools=["bash", "grep"])
+    handler = await started(ctx)
+    await run_turn(handler, ctx, "hi")
+    assert "execute" not in model.bound[0] and "grep" not in model.bound[0]
+    assert "write_file" in model.bound[0]
+
+
+class Answer(BaseModel):
+    city: str
+
+
+async def test_structured_output(tmp_path: Path, fake_model) -> None:
+    fake_model([tool_call("Answer", {"city": "Paris"}, "c1")])
+    ctx = make_ctx(tmp_path, output=Answer)
+    handler = await started(ctx)
+    events = await run_turn(handler, ctx, "capital of France?")
+    assert Answer.model_validate_json(ctx.output_json or "") == Answer(city="Paris")
+    assert not any(isinstance(e, ToolCall) for e in events)
+
+
+async def test_custom_tool(tmp_path: Path, fake_model) -> None:
+    def add(a: int, b: int) -> int:
+        """Add two numbers."""
+        return a + b
+
+    fake_model([tool_call("add", {"a": 2, "b": 3}, "c1"), final("5")])
+    ctx = make_ctx(tmp_path, tools=[add])
+    handler = await started(ctx)
+    events = await run_turn(handler, ctx, "2+3")
+    call = next(e for e in events if isinstance(e, ToolCall))
+    assert (call.name, call.builtin) == ("add", False)
+    assert next(e for e in events if isinstance(e, ToolResult)).output == "5"
+
+
+@pytest.mark.parametrize("approve", [True, False])
+async def test_ask_permissions_emit_approval(
+    tmp_path: Path, fake_model, approve: bool
+) -> None:
+    fake_model(
+        [
+            tool_call("write_file", {"file_path": "/a.txt", "content": "x"}, "c1"),
+            final("end"),
+        ]
+    )
+    ctx = make_ctx(tmp_path, permissions="ask")
+    handler = await started(ctx)
+    events = await run_turn(handler, ctx, "write a", approve=approve)
+    approval = next(e for e in events if isinstance(e, Approval))
+    assert approval.tool == "write"
+    assert approval.input["file_path"] == "/a.txt"
+    assert (tmp_path / "a.txt").exists() is approve
+    assert ctx.final_text == "end"
+
+
+async def test_edit_and_execute_through_sandbox(tmp_path: Path, fake_model) -> None:
+    (tmp_path / "f.txt").write_text("one two\n")
+    fake_model(
+        [
+            tool_call(
+                "edit_file",
+                {"file_path": "/f.txt", "old_string": "two", "new_string": "three"},
+                "c1",
+            ),
+            tool_call("execute", {"command": "cat f.txt"}, "c2"),
+            final("ok"),
+        ]
+    )
+    ctx = make_ctx(tmp_path)
+    handler = await started(ctx)
+    events = await run_turn(handler, ctx, "edit")
+    calls = [e.name for e in events if isinstance(e, ToolCall)]
+    assert calls == ["edit", "bash"]
+    results = [e for e in events if isinstance(e, ToolResult)]
+    assert "one three" in results[1].output
+    assert (tmp_path / "f.txt").read_text() == "one three\n"
+
+
+async def test_resume_keeps_thread(tmp_path: Path, fake_model) -> None:
+    fake_model([final("first"), final("second")])
+    ctx = make_ctx(tmp_path)
+    handler = await started(ctx)
+    await run_turn(handler, ctx, "one")
+    native = handler.native_session_id()
+    assert native == ctx.session_id
+
+    other = make_handler()
+    ctx2 = make_ctx(tmp_path)
+    await other.start(ctx2)
+    await other.resume(ctx2, native or "")
+    await run_turn(other, ctx2, "two")
+    history = await other.history(ctx2)
+    assert [m["content"] for m in history if m["role"] == "user"] == ["one", "two"]
+
+
+async def test_skills_copied_and_loaded(tmp_path: Path, fake_model) -> None:
+    skill = tmp_path / "src-skills" / "greeter"
+    skill.mkdir(parents=True)
+    (skill / "SKILL.md").write_text(
+        "---\nname: greeter\ndescription: Says hi\n---\nSay hi.\n"
+    )
+    work = tmp_path / "work"
+    work.mkdir()
+    fake_model([final("ok")])
+    ctx = make_ctx(work, skills=[str(skill)])
+    handler = await started(ctx)
+    await run_turn(handler, ctx, "hi")
+    assert (work / ".deepagents" / "skills" / "greeter" / "SKILL.md").exists()
+
+
+async def test_turn_and_history_before_start_and_after_stop(
+    tmp_path: Path, fake_model
+) -> None:
+    fake_model([final("ok")])
+    ctx = make_ctx(tmp_path)
+    handler = make_handler()
+    with pytest.raises(HarnessError, match="not started"):
+        await run_turn(handler, ctx, "hi")
+    await handler.start(ctx)
+    await handler.stop(ctx)
+    with pytest.raises(HarnessError, match="not started"):
+        await handler.history(ctx)
+
+
+async def test_start_validates_model(tmp_path: Path, fake_model) -> None:
+    fake_model([final("ok")])
+    with pytest.raises(ValueError, match="needs model="):
+        await started(make_ctx(tmp_path, model=None))
+
+
+def test_build_chat_model_uses_chat_model_kwargs(tmp_path: Path) -> None:
+    deps = dh.load_deps()
+    gw = GatewayTarget(api_base="https://gw.example.com", api_key="sk-virtual")
+    model = dh.build_chat_model(make_ctx(tmp_path, gateway=gw), deps)
+    assert isinstance(model, deps.chat_litellm)
+    assert model.model == "litellm_proxy/gpt-4o-mini"
+
+
+def test_shared_checkpointer_is_process_wide() -> None:
+    deps = dh.load_deps()
+    assert dh.shared_checkpointer(deps) is dh.shared_checkpointer(deps)
+
+
+async def test_sandbox_backend_fs_ops(tmp_path: Path) -> None:
+    (tmp_path / "src").mkdir()
+    (tmp_path / "src" / "a.py").write_text("print('hello')\n")
+    (tmp_path / "b.txt").write_text("hello world\n")
+    backend = SandboxBackend(LocalSandbox(tmp_path), loop=asyncio.get_running_loop())
+
+    ls = await backend.als("/")
+    assert {e["path"] for e in ls.entries or []} == {"/b.txt", "/src/"}
+    assert (await backend.als("/missing")).error
+    globbed = await backend.aglob("*.py")
+    assert [m["path"] for m in globbed.matches or []] == ["/src/a.py"]
+    grep = await backend.agrep("hello", glob="*.txt")
+    assert [(m["path"], m["line"]) for m in grep.matches or []] == [("/b.txt", 1)]
+    read = await backend.aread("/b.txt")
+    assert read.file_data and read.file_data["content"] == "hello world\n"
+    assert (await backend.aread("/nope.txt")).error
+    assert (await backend.aread("/../etc/passwd")).error
+    edit = await backend.aedit("/b.txt", "hello", "bye")
+    assert edit.occurrences == 1
+    assert (await backend.aedit("/b.txt", "zzz", "q")).error
+    assert (await backend.adelete("/src")).path == "/src"
+    assert not (tmp_path / "src").exists()
+    assert (
+        backend.to_real(str(tmp_path / "b.txt"))
+        == str(LocalSandbox(tmp_path).workdir) + "/b.txt"
+    )
+    sync_ls = await asyncio.to_thread(backend.ls, "/")
+    assert [e["path"] for e in sync_ls.entries or []] == ["/b.txt"]
+
+    read_only = SandboxBackend(
+        LocalSandbox(tmp_path),
+        loop=asyncio.get_running_loop(),
+        writable=False,
+        allow_execute=False,
+    )
+    assert (await read_only.awrite("/c.txt", "x")).error
+    assert (await read_only.aexecute("ls")).exit_code == 1
+    assert not (tmp_path / "c.txt").exists()
+
+
+def test_message_cost_prefers_reported_and_never_raises() -> None:
+    reported = AIMessage(content="", response_metadata={"response_cost": 0.5})
+    assert message_cost(reported, "gpt-4o-mini", 1, 1) == 0.5
+    assert message_cost(AIMessage(content=""), "not-a-real-model-xyz", 10, 10) == 0.0
+    assert message_cost(AIMessage(content=""), "gpt-4o-mini", 1000, 1000) > 0
+
+
+def test_missing_deps_raise_install_hint(monkeypatch: pytest.MonkeyPatch) -> None:
+    real_import = builtins.__import__
+
+    def fake_import(name: str, *args: Any, **kwargs: Any) -> Any:
+        if name.startswith("deepagents"):
+            raise ImportError("No module named 'deepagents'")
+        return real_import(name, *args, **kwargs)
+
+    for mod in [m for m in sys.modules if m.startswith("deepagents")]:
+        monkeypatch.delitem(sys.modules, mod)
+    monkeypatch.setattr(builtins, "__import__", fake_import)
+    with pytest.raises(
+        HarnessInstallFailed, match="pip install deepagents langchain-litellm"
+    ):
+        dh.load_deps()
+
+
+LIVE_BASE = os.environ.get("LITELLM_PROXY_API_BASE", "")
+LIVE_KEY = os.environ.get("LITELLM_PROXY_API_KEY", "")
+
+
+@pytest.mark.skipif(
+    not (LIVE_BASE and LIVE_KEY), reason="LITELLM_PROXY_API_BASE / KEY not set"
+)
+async def test_live_gateway_write_file(tmp_path: Path) -> None:
+    model = os.environ.get("HARNESS_DEEPAGENTS_LIVE_MODEL", "claude-haiku-4-5-20251001")
+    ctx = make_ctx(
+        tmp_path,
+        model=model,
+        gateway=GatewayTarget(api_base=LIVE_BASE, api_key=LIVE_KEY),
+        max_turns=6,
+    )
+    handler = await started(ctx)
+    events = await run_turn(handler, ctx, "write hello.txt with hi")
+    assert any(isinstance(e, ToolCall) and e.name == "write" for e in events)
+    assert (tmp_path / "hello.txt").read_text().strip() == "hi"
+    assert ctx.calls >= 1 and ctx.input_tokens > 0

tests/unit/harness/sandbox/test_docker.py

@@ -0,0 +1,241 @@
+import asyncio
+import hashlib
+import shutil
+import subprocess
+from typing import Optional
+
+import pytest
+
+from litellm import sandbox
+from litellm.harness.errors import SandboxError
+from litellm.harness.sandbox import DockerSandbox, Sandbox
+from litellm.harness.sandbox.docker import parse_sha256sum
+
+DOCKER_IMAGE = "alpine:3.20"
+CID = "cid123"
+
+
+class FakeStdin:
+    def __init__(self) -> None:
+        self.data = b""
+        self.closed = False
+
+    def write(self, data: bytes) -> None:
+        self.data += data
+
+    async def drain(self) -> None:
+        return None
+
+    def close(self) -> None:
+        self.closed = True
+
+
+class FakeHandle:
+    def __init__(self, stdout: bytes = b"", stderr: bytes = b"", code: int = 0):
+        self.stdin = FakeStdin()
+        self.stdout = asyncio.StreamReader()
+        self.stdout.feed_data(stdout)
+        self.stdout.feed_eof()
+        self.stderr = asyncio.StreamReader()
+        self.stderr.feed_data(stderr)
+        self.stderr.feed_eof()
+        self.returncode: Optional[int] = code
+        self._code = code
+
+    async def wait(self) -> int:
+        return self._code
+
+    async def kill(self) -> None:
+        return None
+
+
+class Recorder:
+    """Stands in for DockerSandbox._spawn; scripted responses by docker subcommand."""
+
+    def __init__(self) -> None:
+        self.calls: list[list[str]] = []
+        self.handles: list[FakeHandle] = []
+        self.responses: dict[str, FakeHandle] = {}
+
+    async def __call__(self, args: list[str]) -> FakeHandle:
+        self.calls.append(args)
+        handle = self.responses.pop(args[0], None)
+        if handle is None:
+            handle = FakeHandle(stdout=f"{CID}\n".encode() if args[0] == "run" else b"")
+        self.handles.append(handle)
+        return handle
+
+
+@pytest.fixture
+def fake(monkeypatch):
+    rec = Recorder()
+    monkeypatch.setattr(DockerSandbox, "_spawn", lambda self, args: rec(args))
+    return rec
+
+
+def test_run_args():
+    box = sandbox.docker(
+        "img:1",
+        mounts={"/host/src": "/workspace"},
+        env={"A": "1"},
+        name="h1",
+    )
+    assert isinstance(box, Sandbox)
+    assert box.run_args() == [
+        "run",
+        "-d",
+        "--rm",
+        "--add-host=host.docker.internal:host-gateway",
+        "--name",
+        "h1",
+        "-v",
+        "/host/src:/workspace",
+        "-e",
+        "A=1",
+        "-w",
+        "/workspace",
+        "img:1",
+        "sleep",
+        "infinity",
+    ]
+    assert box.host_url(8080) == "http://host.docker.internal:8080"
+
+
+def test_relative_workdir_rejected():
+    with pytest.raises(SandboxError):
+        sandbox.docker("img", workdir="rel")
+
+
+async def test_lazy_start_and_exec_args(fake):
+    box = sandbox.docker("img")
+    assert fake.calls == []
+    await box.exec(["echo", "hi"], env={"K": "V"}, cwd="sub")
+    await box.exec(["true"])
+    assert fake.calls[0][0] == "run"
+    assert [c for c in fake.calls if c[0] == "run"] == [fake.calls[0]]
+    assert fake.calls[1] == [
+        "exec",
+        "-i",
+        "-w",
+        "/workspace/sub",
+        "-e",
+        "K=V",
+        CID,
+        "echo",
+        "hi",
+    ]
+    assert fake.calls[2] == ["exec", "-i", "-w", "/workspace", CID, "true"]
+
+
+async def test_run_start_failure(fake):
+    fake.responses["run"] = FakeHandle(stderr=b"no such image", code=125)
+    box = sandbox.docker("img")
+    with pytest.raises(SandboxError, match="no such image"):
+        await box.run(["echo"])
+
+
+async def test_read_write_which_tempdir(fake):
+    box = sandbox.docker("img")
+    await box.start()
+
+    fake.responses["exec"] = FakeHandle(stdout=b"content")
+    assert await box.read("a.txt") == b"content"
+    assert fake.calls[-1][-2:] == ["cat", "/workspace/a.txt"]
+
+    await box.write("d/b.txt", b"payload")
+    assert fake.calls[-1][-5:-1] == ["sh", "-c", fake.calls[-1][-3], "sh"]
+    assert fake.calls[-1][-1] == "/workspace/d/b.txt"
+    assert fake.handles[-1].stdin.data == b"payload"
+    assert fake.handles[-1].stdin.closed
+
+    fake.responses["exec"] = FakeHandle(stdout=b"/usr/bin/codex\n")
+    assert await box.which("codex") == "/usr/bin/codex"
+    assert fake.calls[-1][-5:] == ["sh", "-lc", 'command -v "$1"', "sh", "codex"]
+
+    fake.responses["exec"] = FakeHandle(code=1)
+    assert await box.which("nope") is None
+
+    fake.responses["exec"] = FakeHandle(stdout=b"/tmp/tmp.abc\n")
+    assert await box.tempdir() == "/tmp/tmp.abc"
+
+    fake.responses["exec"] = FakeHandle(stderr=b"No such file", code=1)
+    with pytest.raises(SandboxError, match="No such file"):
+        await box.read("missing")
+
+
+async def test_snapshot_parses_output(fake):
+    box = sandbox.docker("img")
+    digest = "a" * 64
+    fake.responses["exec"] = FakeHandle(
+        stdout=f"{digest}  ./x.txt\n{digest}  ./dir/with space.txt\n".encode()
+    )
+    snap = await box.snapshot()
+    assert snap == {"x.txt": digest, "dir/with space.txt": digest}
+    script = fake.calls[-1][-3]
+    assert "-name '.git'" in script and "-prune" in script
+    assert fake.calls[-1][-1] == "/workspace"
+
+
+async def test_close_removes_container(fake):
+    box = sandbox.docker("img")
+    await box.start()
+    await box.close()
+    assert fake.calls[-1] == ["rm", "-f", CID]
+    with pytest.raises(SandboxError):
+        await box.start()
+
+
+async def test_close_without_start_is_noop(fake):
+    await sandbox.docker("img").close()
+    assert fake.calls == []
+
+
+async def test_missing_docker_binary(monkeypatch):
+    monkeypatch.setattr(shutil, "which", lambda name, *a, **k: None)
+    with pytest.raises(SandboxError, match="docker"):
+        await sandbox.docker("img").start()
+
+
+def test_parse_sha256sum_ignores_junk():
+    assert parse_sha256sum("garbage\n\n") == {}
+
+
+def _docker_usable() -> bool:
+    if shutil.which("docker") is None:
+        return False
+    try:
+        return (
+            subprocess.run(
+                ["docker", "info"], capture_output=True, timeout=20
+            ).returncode
+            == 0
+        )
+    except (OSError, subprocess.SubprocessError):
+        return False
+
+
+@pytest.mark.skipif(not _docker_usable(), reason="docker daemon not available")
+async def test_real_docker_roundtrip():
+    box = sandbox.docker(DOCKER_IMAGE, workdir="/workspace")
+    try:
+        result = await box.run(["echo", "hello"], timeout=120)
+        assert result.stdout.strip() == "hello"
+        assert result.exit_code == 0
+
+        await box.write("seed.txt", b"seed")
+        await box.write("sub/out.txt", b"from host")
+        assert await box.read("sub/out.txt") == b"from host"
+        await box.write("node_modules/skip.js", b"x")
+
+        assert await box.which("sh") is not None
+        assert await box.which("definitely-not-a-binary-xyz") is None
+        tmp = await box.tempdir()
+        assert tmp.startswith("/")
+
+        snap = await box.snapshot()
+        assert snap == {
+            "seed.txt": hashlib.sha256(b"seed").hexdigest(),
+            "sub/out.txt": hashlib.sha256(b"from host").hexdigest(),
+        }
+    finally:
+        await box.close()

tests/unit/harness/sandbox/test_local.py

@@ -0,0 +1,180 @@
+import os
+import sys
+
+import pytest
+
+from litellm import sandbox
+from litellm.harness.errors import SandboxError
+from litellm.harness.sandbox import LocalSandbox, Process, Sandbox
+from litellm.harness.sandbox.local import filtered_environ, is_secret_env_name
+
+PY = sys.executable
+
+
+@pytest.fixture
+async def sbx(tmp_path):
+    box = sandbox.local(tmp_path)
+    yield box
+    await box.close()
+
+
+def test_local_requires_existing_dir(tmp_path):
+    with pytest.raises(SandboxError):
+        sandbox.local(tmp_path / "missing")
+
+
+def test_local_resolves_absolute_and_satisfies_protocol(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+    (tmp_path / "ws").mkdir()
+    box = sandbox.local("ws")
+    assert isinstance(box, LocalSandbox)
+    assert isinstance(box, Sandbox)
+    assert box.workdir == os.path.realpath(tmp_path / "ws")
+    assert box.host_url(4321) == "http://127.0.0.1:4321"
+
+
+async def test_run_collects_output(sbx):
+    result = await sbx.run(
+        [
+            PY,
+            "-c",
+            "import os,sys;print(os.getcwd());print('err',file=sys.stderr);sys.exit(3)",
+        ]
+    )
+    assert result.stdout.strip() == sbx.workdir
+    assert result.stderr.strip() == "err"
+    assert result.exit_code == 3
+
+
+async def test_run_cwd_inside_workdir(sbx):
+    os.mkdir(os.path.join(sbx.workdir, "sub"))
+    result = await sbx.run([PY, "-c", "import os;print(os.getcwd())"], cwd="sub")
+    assert result.stdout.strip() == os.path.join(sbx.workdir, "sub")
+    with pytest.raises(SandboxError):
+        await sbx.run([PY, "-c", "pass"], cwd="/")
+
+
+async def test_exec_streams_stdin(sbx):
+    proc = await sbx.exec([PY, "-c", "import sys;print(sys.stdin.read().upper())"])
+    assert isinstance(proc, Process)
+    assert proc.stdin is not None
+    proc.stdin.write(b"hello")
+    await proc.stdin.drain()
+    proc.stdin.close()
+    assert (await proc.stdout.read()).strip() == b"HELLO"
+    assert await proc.wait() == 0
+
+
+async def test_run_timeout_kills(sbx):
+    with pytest.raises(SandboxError, match="timed out"):
+        await sbx.run([PY, "-c", "import time;time.sleep(30)"], timeout=0.5)
+
+
+async def test_missing_binary_raises(sbx):
+    with pytest.raises(SandboxError):
+        await sbx.run(["definitely-not-a-binary-xyz"])
+
+
+async def test_close_kills_live_processes(tmp_path):
+    box = sandbox.local(tmp_path)
+    proc = await box.exec([PY, "-c", "import time;time.sleep(30)"])
+    await box.close()
+    assert proc.returncode is not None
+    with pytest.raises(SandboxError):
+        await box.run([PY, "-c", "pass"])
+
+
+async def test_read_write_roundtrip(sbx):
+    await sbx.write("a/b/c.txt", b"data")
+    assert await sbx.read("a/b/c.txt") == b"data"
+    abs_path = os.path.join(sbx.workdir, "a", "b", "c.txt")
+    assert await sbx.read(abs_path) == b"data"
+
+
+@pytest.mark.parametrize("bad", ["../escape.txt", "a/../../escape.txt", "/etc/passwd"])
+async def test_path_escape_rejected(sbx, bad):
+    with pytest.raises(SandboxError, match="escapes"):
+        await sbx.read(bad)
+    with pytest.raises(SandboxError, match="escapes"):
+        await sbx.write(bad, b"x")
+
+
+async def test_symlink_escape_rejected(sbx, tmp_path_factory):
+    outside = tmp_path_factory.mktemp("outside")
+    os.symlink(outside, os.path.join(sbx.workdir, "link"))
+    with pytest.raises(SandboxError, match="escapes"):
+        await sbx.write("link/x.txt", b"x")
+
+
+async def test_tempdir_is_allowed_and_cleaned(tmp_path):
+    box = sandbox.local(tmp_path)
+    tmp = await box.tempdir()
+    assert os.path.isdir(tmp)
+    target = os.path.join(tmp, "config.toml")
+    await box.write(target, b"k = 1")
+    assert await box.read(target) == b"k = 1"
+    await box.close()
+    assert not os.path.exists(tmp)
+
+
+async def test_env_filters_provider_secrets(sbx, monkeypatch):
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-fake")
+    monkeypatch.setenv("OPENAI_BASE_URL", "http://x")
+    monkeypatch.setenv("GITHUB_TOKEN", "ghp_fake")
+    monkeypatch.setenv("HARNESS_TEST_PLAIN", "visible")
+    script = (
+        "import os;"
+        "print(os.environ.get('ANTHROPIC_API_KEY','<none>'));"
+        "print(os.environ.get('OPENAI_BASE_URL','<none>'));"
+        "print(os.environ.get('GITHUB_TOKEN','<none>'));"
+        "print(os.environ.get('HARNESS_TEST_PLAIN','<none>'));"
+        "print(os.environ.get('ANTHROPIC_BASE_URL','<none>'))"
+    )
+    result = await sbx.run(
+        [PY, "-c", script], env={"ANTHROPIC_BASE_URL": "http://127.0.0.1:1"}
+    )
+    assert result.stdout.split() == [
+        "<none>",
+        "<none>",
+        "<none>",
+        "visible",
+        "http://127.0.0.1:1",
+    ]
+
+
+@pytest.mark.parametrize(
+    "name,secret",
+    [
+        ("ANTHROPIC_API_KEY", True),
+        ("AWS_REGION", True),
+        ("VERTEXAI_PROJECT", True),
+        ("GOOGLE_APPLICATION_CREDENTIALS", True),
+        ("MY_API_KEY", True),
+        ("SLACK_BOT_TOKEN", True),
+        ("CLIENT_SECRET", True),
+        ("PATH", False),
+        ("HOME", False),
+    ],
+)
+def test_is_secret_env_name(name, secret):
+    assert is_secret_env_name(name) is secret
+
+
+def test_filtered_environ_overlay_wins():
+    env = filtered_environ(
+        {"PATH": "/bin", "OPENAI_API_KEY": "x"}, {"PATH": "/usr/bin"}
+    )
+    assert env == {"PATH": "/usr/bin"}
+
+
+async def test_which_uses_filtered_path(sbx):
+    assert await sbx.which("sh") is not None
+    assert await sbx.which("definitely-not-a-binary-xyz") is None
+
+
+async def test_snapshot_skips_dirs(sbx):
+    await sbx.write("keep.txt", b"k")
+    await sbx.write(".git/HEAD", b"ref")
+    await sbx.write("node_modules/x/index.js", b"x")
+    snap = await sbx.snapshot()
+    assert list(snap) == ["keep.txt"]

tests/unit/harness/sandbox/test_snapshot.py

@@ -0,0 +1,129 @@
+import hashlib
+import os
+
+import pytest
+
+from litellm import sandbox
+from litellm.constants import HARNESS_MAX_DIFF_BYTES
+from litellm.harness.sandbox.snapshot import (
+    build_file_changes,
+    capture_text_contents,
+    diff_snapshots,
+    snapshot_local,
+    unified_diff,
+)
+from litellm.harness.types import FileChange
+
+
+def test_diff_snapshots_kinds():
+    before = {"a": "1", "b": "2", "c": "3"}
+    after = {"a": "1", "b": "9", "d": "4"}
+    assert diff_snapshots(before, after) == [
+        ("b", "modified"),
+        ("c", "deleted"),
+        ("d", "created"),
+    ]
+
+
+async def test_snapshot_local_hashes_and_skips(tmp_path):
+    (tmp_path / "x.txt").write_bytes(b"hello")
+    (tmp_path / "sub").mkdir()
+    (tmp_path / "sub" / "y.txt").write_bytes(b"y")
+    (tmp_path / "__pycache__").mkdir()
+    (tmp_path / "__pycache__" / "z.pyc").write_bytes(b"z")
+    os.symlink(tmp_path / "x.txt", tmp_path / "link.txt")
+    snap = await snapshot_local(str(tmp_path))
+    assert snap == {
+        "x.txt": hashlib.sha256(b"hello").hexdigest(),
+        "sub/y.txt": hashlib.sha256(b"y").hexdigest(),
+    }
+
+
+def test_unified_diff_created():
+    diff = unified_diff("f.txt", None, "one\n")
+    assert diff.startswith("--- /dev/null\n+++ b/f.txt\n")
+    assert "+one\n" in diff
+
+
+async def test_created_modified_deleted_end_to_end(tmp_path):
+    box = sandbox.local(tmp_path)
+    try:
+        await box.write("mod.txt", b"line1\nline2\n")
+        await box.write("gone.txt", b"bye\n")
+        await box.write("bin.dat", b"\x00\x01\x02")
+        before = await box.snapshot()
+        contents = await capture_text_contents(box, before)
+        assert set(contents) == {"mod.txt", "gone.txt"}
+
+        await box.write("mod.txt", b"line1\nchanged\n")
+        await box.write("new.txt", b"fresh\n")
+        await box.write("bin.dat", b"\x00\x09")
+        os.remove(os.path.join(box.workdir, "gone.txt"))
+        after = await box.snapshot()
+
+        changes = await build_file_changes(box, before, after, contents)
+        by_path = {c.path: c for c in changes}
+        assert all(isinstance(c, FileChange) for c in changes)
+        assert [(c.path, c.kind) for c in changes] == [
+            ("bin.dat", "modified"),
+            ("gone.txt", "deleted"),
+            ("mod.txt", "modified"),
+            ("new.txt", "created"),
+        ]
+        assert by_path["bin.dat"].diff is None
+        assert "-line2\n" in by_path["mod.txt"].diff
+        assert "+changed\n" in by_path["mod.txt"].diff
+        assert "+fresh\n" in by_path["new.txt"].diff
+        assert "-bye\n" in by_path["gone.txt"].diff
+    finally:
+        await box.close()
+
+
+async def test_modified_without_before_contents_has_no_diff(tmp_path):
+    box = sandbox.local(tmp_path)
+    try:
+        await box.write("f.txt", b"a\n")
+        before = await box.snapshot()
+        await box.write("f.txt", b"b\n")
+        after = await box.snapshot()
+        changes = await build_file_changes(box, before, after, None)
+        assert changes == [FileChange(path="f.txt", kind="modified", diff=None)]
+    finally:
+        await box.close()
+
+
+async def test_large_text_file_has_no_diff(tmp_path):
+    box = sandbox.local(tmp_path)
+    try:
+        before = await box.snapshot()
+        await box.write("big.txt", b"a" * (HARNESS_MAX_DIFF_BYTES + 1))
+        after = await box.snapshot()
+        changes = await build_file_changes(box, before, after, {})
+        assert changes == [FileChange(path="big.txt", kind="created", diff=None)]
+    finally:
+        await box.close()
+
+
+async def test_capture_respects_total_cap(tmp_path, monkeypatch):
+    monkeypatch.setattr(
+        "litellm.harness.sandbox.snapshot.HARNESS_SNAPSHOT_MAX_TOTAL_BYTES", 10
+    )
+    box = sandbox.local(tmp_path)
+    try:
+        await box.write("a.txt", b"x" * 8)
+        await box.write("b.txt", b"x" * 8)
+        await box.write("c.txt", b"x" * 8)
+        captured = await capture_text_contents(box, await box.snapshot())
+        assert set(captured) == {"a.txt", "b.txt"}
+    finally:
+        await box.close()
+
+
+@pytest.mark.parametrize("data", [b"\xff\xfe bad utf8", b"has\x00nul"])
+async def test_capture_skips_binary(tmp_path, data):
+    box = sandbox.local(tmp_path)
+    try:
+        await box.write("f", data)
+        assert await capture_text_contents(box, {"f": "h"}) == {}
+    finally:
+        await box.close()

tests/unit/harness/test_endpoint.py

@@ -0,0 +1,359 @@
+import json
+import sys
+from collections.abc import AsyncIterator
+from typing import Any
+
+import httpx
+import pytest
+
+import litellm
+from litellm.harness import endpoint as endpoint_module
+from litellm.harness.endpoint import (
+    ModelEndpoint,
+    SSEUsageParser,
+    UsageTracker,
+    compute_cost,
+    usage_from_body,
+)
+from litellm.harness.errors import HarnessInstallFailed
+from litellm.harness.context import GatewayTarget
+from litellm.harness.types import Harness, Usage
+from litellm.types.utils import ModelResponse, ModelResponseStream
+
+GATEWAY = GatewayTarget(api_base="https://gw.example.com", api_key="sk-gateway-secret")
+
+ANTHROPIC_SSE = (
+    b"event: message_start\n"
+    b'data: {"type":"message_start","message":{"usage":{"input_tokens":11,"output_tokens":1}}}\n\n'
+    b"event: content_block_delta\n"
+    b'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"hi"}}\n\n'
+    b"event: message_delta\n"
+    b'data: {"type":"message_delta","usage":{"output_tokens":7}}\n\n'
+    b"event: message_stop\n"
+    b'data: {"type":"message_stop"}\n\n'
+)
+CHAT_SSE = (
+    b'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n'
+    b'data: {"choices":[],"usage":{"prompt_tokens":5,"completion_tokens":3}}\n\n'
+    b"data: [DONE]\n\n"
+)
+RESPONSES_SSE = (
+    b"event: response.output_text.delta\n"
+    b'data: {"type":"response.output_text.delta","delta":"hi"}\n\n'
+    b"event: response.completed\n"
+    b'data: {"type":"response.completed","response":{"usage":{"input_tokens":20,"output_tokens":4}}}\n\n'
+)
+
+
+class Recorder:
+    def __init__(self, response: httpx.Response) -> None:
+        self.response = response
+        self.requests: list[httpx.Request] = []
+
+    def __call__(self, request: httpx.Request) -> httpx.Response:
+        self.requests.append(request)
+        return self.response
+
+
+def sse_response(body: bytes, headers: dict[str, str] | None = None) -> httpx.Response:
+    return httpx.Response(
+        200,
+        content=body,
+        headers={"content-type": "text/event-stream", **(headers or {})},
+    )
+
+
+def gateway_endpoint(recorder: Recorder, **kwargs: Any) -> ModelEndpoint:
+    return ModelEndpoint(
+        Harness.CLAUDE_CODE,
+        kwargs.pop("model", "claude-sonnet"),
+        GATEWAY,
+        client=httpx.AsyncClient(transport=httpx.MockTransport(recorder)),
+        **kwargs,
+    )
+
+
+def auth(ep: ModelEndpoint) -> dict[str, str]:
+    return {"authorization": f"Bearer {ep.token}"}
+
+
+@pytest.fixture(autouse=True)
+def no_real_cost(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(
+        litellm,
+        "cost_per_token",
+        lambda model, prompt_tokens, completion_tokens: (
+            prompt_tokens * 0.001,
+            completion_tokens * 0.002,
+        ),
+    )
+
+
+async def test_rejects_bad_token_and_accepts_both_header_styles() -> None:
+    recorder = Recorder(httpx.Response(200, json={"usage": {}}))
+    async with gateway_endpoint(recorder) as ep:
+        assert ep.url == f"http://127.0.0.1:{ep.port}" and ep.port > 0
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            missing = await client.post("/v1/messages", json={})
+            wrong = await client.post(
+                "/v1/messages", json={}, headers={"x-api-key": "nope"}
+            )
+            bearer = await client.post("/v1/messages", json={}, headers=auth(ep))
+            api_key = await client.post(
+                "/messages", json={}, headers={"x-api-key": ep.token}
+            )
+    assert missing.status_code == 401
+    assert wrong.status_code == 401
+    assert "error" in wrong.json()
+    assert bearer.status_code == 200
+    assert api_key.status_code == 200
+    assert len(recorder.requests) == 2
+
+
+async def test_gateway_rewrites_headers_and_model() -> None:
+    recorder = Recorder(
+        httpx.Response(
+            200,
+            json={"id": "m", "usage": {"input_tokens": 3, "output_tokens": 2}},
+        )
+    )
+    async with gateway_endpoint(recorder, metadata={"run": "abc"}) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            resp = await client.post(
+                "/v1/messages",
+                json={"model": "whatever", "max_tokens": 5},
+                headers={
+                    "x-api-key": ep.token,
+                    "anthropic-version": "2023-06-01",
+                    "anthropic-beta": "tools-2024",
+                },
+            )
+    assert resp.status_code == 200
+    sent = recorder.requests[0]
+    assert str(sent.url) == "https://gw.example.com/v1/messages"
+    assert sent.headers["authorization"] == "Bearer sk-gateway-secret"
+    assert "x-api-key" not in sent.headers
+    assert sent.headers["x-litellm-tags"] == "harness,claude_code"
+    assert json.loads(sent.headers["x-litellm-spend-logs-metadata"]) == {"run": "abc"}
+    assert sent.headers["anthropic-version"] == "2023-06-01"
+    assert sent.headers["anthropic-beta"] == "tools-2024"
+    assert json.loads(sent.content)["model"] == "claude-sonnet"
+    assert ep.usage.input_tokens == 3 and ep.usage.output_tokens == 2
+    assert ep.usage.calls == 1
+
+
+@pytest.mark.parametrize(
+    "path,body,expected",
+    [
+        ("/v1/messages", ANTHROPIC_SSE, (11, 7)),
+        ("/v1/chat/completions", CHAT_SSE, (5, 3)),
+        ("/responses", RESPONSES_SSE, (20, 4)),
+    ],
+)
+async def test_gateway_sse_passthrough_and_usage(
+    path: str, body: bytes, expected: tuple[int, int]
+) -> None:
+    recorder = Recorder(sse_response(body))
+    async with gateway_endpoint(recorder) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            resp = await client.post(path, json={"stream": True}, headers=auth(ep))
+    assert resp.status_code == 200
+    assert resp.headers["content-type"].startswith("text/event-stream")
+    assert resp.content == body
+    assert (ep.usage.input_tokens, ep.usage.output_tokens) == expected
+    expected_cost = expected[0] * 0.001 + expected[1] * 0.002
+    assert ep.usage.cost == pytest.approx(expected_cost)
+
+
+async def test_cost_header_preferred_over_computed() -> None:
+    recorder = Recorder(
+        httpx.Response(
+            200,
+            json={"usage": {"prompt_tokens": 100, "completion_tokens": 100}},
+            headers={"x-litellm-response-cost": "0.42"},
+        )
+    )
+    async with gateway_endpoint(recorder) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            await client.post("/v1/chat/completions", json={}, headers=auth(ep))
+    assert ep.usage.cost == pytest.approx(0.42)
+    assert ep.usage.snapshot() == Usage(input_tokens=100, output_tokens=100, calls=1)
+
+
+async def test_gateway_error_status_preserved_and_not_counted() -> None:
+    recorder = Recorder(httpx.Response(429, json={"error": "rate limited"}))
+    async with gateway_endpoint(recorder) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            resp = await client.post("/v1/chat/completions", json={}, headers=auth(ep))
+    assert resp.status_code == 429
+    assert ep.usage.calls == 0
+
+
+async def test_models_route() -> None:
+    recorder = Recorder(httpx.Response(200))
+    async with gateway_endpoint(recorder) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            with_model = await client.get("/v1/models", headers=auth(ep))
+            unauth = await client.get("/models")
+    assert unauth.status_code == 401
+    assert with_model.json()["object"] == "list"
+    assert [m["id"] for m in with_model.json()["data"]] == ["claude-sonnet"]
+
+    async with ModelEndpoint(Harness.CODEX, None, None) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            empty = await client.get("/models", headers=auth(ep))
+    assert empty.json() == {"object": "list", "data": []}
+
+
+async def test_sdk_chat_non_stream(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls: list[dict[str, Any]] = []
+
+    async def fake_acompletion(**kwargs: Any) -> ModelResponse:
+        calls.append(kwargs)
+        response = ModelResponse(
+            model="gpt-x",
+            choices=[{"message": {"role": "assistant", "content": "hello"}}],
+            usage={"prompt_tokens": 9, "completion_tokens": 4, "total_tokens": 13},
+        )
+        response._hidden_params["response_cost"] = 0.5
+        return response
+
+    monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
+    async with ModelEndpoint(
+        Harness.OPENCODE, "openai/gpt-x", None, api_key="sk-real", api_base="https://x"
+    ) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            resp = await client.post(
+                "/v1/chat/completions",
+                json={
+                    "model": "ignored",
+                    "messages": [{"role": "user", "content": "hi"}],
+                },
+                headers=auth(ep),
+            )
+    assert resp.status_code == 200
+    assert resp.json()["choices"][0]["message"]["content"] == "hello"
+    assert calls[0]["model"] == "openai/gpt-x"
+    assert calls[0]["api_key"] == "sk-real"
+    assert calls[0]["api_base"] == "https://x"
+    assert (ep.usage.input_tokens, ep.usage.output_tokens) == (9, 4)
+    assert ep.usage.cost == pytest.approx(0.5)
+
+
+async def fake_chat_stream() -> AsyncIterator[ModelResponseStream]:
+    yield ModelResponseStream(choices=[{"delta": {"content": "he"}}])
+    yield ModelResponseStream(choices=[{"delta": {"content": "llo"}}])
+    final = ModelResponseStream(choices=[])
+    final.usage = litellm.Usage(prompt_tokens=6, completion_tokens=2, total_tokens=8)
+    yield final
+
+
+async def test_sdk_chat_stream(monkeypatch: pytest.MonkeyPatch) -> None:
+    calls: list[dict[str, Any]] = []
+
+    async def fake_acompletion(**kwargs: Any) -> AsyncIterator[ModelResponseStream]:
+        calls.append(kwargs)
+        return fake_chat_stream()
+
+    monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
+    async with ModelEndpoint(Harness.OPENCODE, "openai/gpt-x", None) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            resp = await client.post(
+                "/chat/completions",
+                json={"messages": [], "stream": True},
+                headers=auth(ep),
+            )
+    assert resp.headers["content-type"].startswith("text/event-stream")
+    lines = [line for line in resp.text.split("\n") if line.startswith("data: ")]
+    assert lines[-1] == "data: [DONE]"
+    assert json.loads(lines[0][6:])["choices"][0]["delta"]["content"] == "he"
+    assert calls[0]["stream_options"] == {"include_usage": True}
+    assert (ep.usage.input_tokens, ep.usage.output_tokens) == (6, 2)
+    assert ep.usage.cost == pytest.approx(6 * 0.001 + 2 * 0.002)
+
+
+async def fake_anthropic_stream() -> AsyncIterator[Any]:
+    yield {"type": "message_start", "message": {"usage": {"input_tokens": 4}}}
+    yield b'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":9}}\n\n'
+
+
+async def test_sdk_messages_stream_handles_dicts_and_bytes(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    async def fake_acreate(**kwargs: Any) -> AsyncIterator[Any]:
+        return fake_anthropic_stream()
+
+    monkeypatch.setattr(litellm.anthropic.messages, "acreate", fake_acreate)
+    async with ModelEndpoint(Harness.CLAUDE_CODE, "anthropic/claude", None) as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            resp = await client.post(
+                "/v1/messages", json={"stream": True}, headers=auth(ep)
+            )
+    assert "event: message_start" in resp.text
+    assert "event: message_delta" in resp.text
+    assert (ep.usage.input_tokens, ep.usage.output_tokens) == (4, 9)
+
+
+async def test_sdk_error_is_sanitized(monkeypatch: pytest.MonkeyPatch) -> None:
+    async def failing(**kwargs: Any) -> Any:
+        raise litellm.RateLimitError(
+            message="too many requests for key sk-real",
+            llm_provider="openai",
+            model="gpt-x",
+        )
+
+    monkeypatch.setattr(litellm, "aresponses", failing)
+    async with ModelEndpoint(Harness.CODEX, "gpt-x", None, api_key="sk-real") as ep:
+        async with httpx.AsyncClient(base_url=ep.url) as client:
+            resp = await client.post("/v1/responses", json={}, headers=auth(ep))
+    assert resp.status_code == 429
+    assert "sk-real" not in resp.text
+    assert resp.json()["error"]["type"] == "RateLimitError"
+
+
+async def test_missing_server_deps_raises_install_failed(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    def missing() -> Any:
+        raise HarnessInstallFailed(endpoint_module.MISSING_DEPS_MESSAGE)
+
+    monkeypatch.setattr(endpoint_module, "_load_server_deps", missing)
+    with pytest.raises(HarnessInstallFailed, match="pip install starlette uvicorn"):
+        async with ModelEndpoint(Harness.CODEX, None, None):
+            pass
+
+
+def test_load_server_deps_maps_import_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setitem(sys.modules, "uvicorn", None)
+    with pytest.raises(HarnessInstallFailed, match="starlette and uvicorn"):
+        endpoint_module._load_server_deps()
+
+
+def test_usage_helpers() -> None:
+    assert usage_from_body({"usage": {"prompt_tokens": 1, "completion_tokens": 2}}) == (
+        1,
+        2,
+    )
+    assert usage_from_body({"response": {"usage": {"input_tokens": 3}}}) == (3, 0)
+    assert usage_from_body("nope") == (0, 0)
+
+    parser = SSEUsageParser()
+    for i in range(0, len(ANTHROPIC_SSE), 7):  # split across arbitrary chunk borders
+        parser.feed(ANTHROPIC_SSE[i : i + 7])
+    parser.close()
+    assert (parser.input_tokens, parser.output_tokens) == (11, 7)
+
+    tracker = UsageTracker()
+    tracker.add(1, 2, 0.1)
+    tracker.add(3, 4, 0.2)
+    assert tracker.snapshot() == Usage(input_tokens=4, output_tokens=6, calls=2)
+    assert tracker.cost == pytest.approx(0.3)
+
+
+def test_compute_cost_never_raises(monkeypatch: pytest.MonkeyPatch) -> None:
+    def boom(**kwargs: Any) -> Any:
+        raise ValueError("unknown model")
+
+    monkeypatch.setattr(litellm, "cost_per_token", boom)
+    assert compute_cost("mystery", 10, 10) == 0.0
+    assert compute_cost(None, 10, 10) == 0.0

tests/unit/harness/test_init.py

@@ -0,0 +1,95 @@
+"""Tests for litellm/harness/__init__.py: the public API surface."""
+
+from __future__ import annotations
+
+
+from litellm import harness
+from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter
+from litellm.utils import ProviderConfigManager
+
+PUBLIC_NAMES = [
+    "Harness",
+    "agent",
+    "aagent",
+    "agent_session",
+    "aagent_session",
+    "agent_resume",
+    "aagent_resume",
+    "agent_capabilities",
+    "Result",
+    "Usage",
+    "State",
+    "Capabilities",
+    "Session",
+    "EventStream",
+    "Text",
+    "Reasoning",
+    "ToolCall",
+    "ToolResult",
+    "FileChange",
+    "Compaction",
+    "Approval",
+    "Done",
+    "Event",
+    "ClaudeCodeOptions",
+    "CodexOptions",
+    "OpenCodeOptions",
+    "DeepAgentsOptions",
+    "HarnessError",
+    "CapabilityUnsupported",
+    "OptionsMismatch",
+    "HarnessInstallFailed",
+    "SandboxError",
+    "SessionClosed",
+    "StateIncompatible",
+    "OutputInvalid",
+]
+ERROR_NAMES = [
+    "CapabilityUnsupported",
+    "OptionsMismatch",
+    "HarnessInstallFailed",
+    "SandboxError",
+    "SessionClosed",
+    "StateIncompatible",
+    "OutputInvalid",
+]
+LAZY_IMPORT_CHECK = (
+    "import sys, litellm\n"
+    "assert 'litellm.harness' not in sys.modules\n"
+    "h = litellm.harness\n"
+    "assert h.Harness.CODEX.value == 'codex'\n"
+    "assert 'starlette' not in sys.modules and 'uvicorn' not in sys.modules\n"
+    "print('ok')\n"
+)
+
+
+def test_public_api_names_exported():
+    missing = [name for name in PUBLIC_NAMES if not hasattr(harness, name)]
+    assert missing == []
+    assert set(PUBLIC_NAMES) <= set(harness.__all__)
+
+
+def test_errors_share_base_class():
+    for name in ERROR_NAMES:
+        assert issubclass(getattr(harness, name), harness.HarnessError)
+
+
+def test_litellm_harness_attribute_is_lazy():
+    out = run_child_interpreter(LAZY_IMPORT_CHECK, timeout=120)
+    assert out.returncode == 0, out.stderr
+    assert out.stdout.strip() == "ok"
+
+
+def test_adapter_registry_paths_cover_every_harness():
+    for member in harness.Harness:
+        config = ProviderConfigManager.get_provider_harness_config(member)
+        assert config is not None and config.harness is member
+
+
+def test_litellm_agent_is_top_level_and_lazy():
+    code = (
+        "import sys, litellm; assert 'litellm.harness' not in sys.modules; "
+        "assert litellm.agent is litellm.harness.agent; assert litellm.Harness.CODEX.value == 'codex'"
+    )
+    out = run_child_interpreter(code, timeout=120)
+    assert out.returncode == 0, out.stderr

tests/unit/harness/test_runtime.py

@@ -0,0 +1,621 @@
+"""Tests for litellm/harness/runtime.py using a fake adapter, sandbox and endpoint."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections.abc import AsyncIterator
+
+import pytest
+from pydantic import BaseModel
+
+from litellm.harness import runtime
+from litellm.harness.context import SessionContext
+from litellm.harness.errors import (
+    CapabilityUnsupported,
+    HarnessInstallFailed,
+    OptionsMismatch,
+    OutputInvalid,
+    SessionClosed,
+    StateIncompatible,
+)
+from litellm.harness.options import CodexOptions
+from litellm.harness.types import (
+    Approval,
+    Done,
+    Event,
+    FileChange,
+    Harness,
+    State,
+    Text,
+    ToolCall,
+)
+from tests.unit.harness.core_fakes import (
+    NARROW_CAPS,
+    FakeAdapter,
+    FakeEndpoint,
+    FakeSandbox,
+    install_adapter,
+    script_approval,
+    wait_forever,
+)
+
+
+class Answer(BaseModel):
+    value: int
+
+
+@pytest.fixture
+def sandbox(tmp_path) -> FakeSandbox:
+    return FakeSandbox(str(tmp_path))
+
+
+async def _collect(stream) -> list[Event]:
+    return [event async for event in stream]
+
+
+# -- validation ---------------------------------------------------------------
+
+
+async def test_string_harness_raises_type_error_with_hint(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    with pytest.raises(TypeError, match=r"Harness\.CODEX"):
+        await runtime.aagent("codex", "hi", sandbox=sandbox)  # type: ignore[arg-type]
+
+
+async def test_options_mismatch(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch)
+    with pytest.raises(OptionsMismatch, match="CodexOptions"):
+        await runtime.aagent(
+            Harness.CLAUDE_CODE, "hi", sandbox=sandbox, options=CodexOptions()
+        )
+    assert adapter_cls.instances == []
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"permissions": "edit"},
+        {"output": Answer},
+        {"tools": [print]},
+        {"disable_tools": ["bash"]},
+        {"permissions": "ask", "on_approval": lambda a: True},
+    ],
+)
+async def test_capability_errors_before_start(monkeypatch, sandbox, kwargs):
+    adapter_cls = install_adapter(monkeypatch, caps=NARROW_CAPS)
+    with pytest.raises(CapabilityUnsupported):
+        await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, **kwargs)
+    assert all("start" not in a.calls for a in adapter_cls.instances)
+    assert FakeEndpoint.instances == []
+
+
+async def test_skills_capability_error_before_start(monkeypatch, sandbox, tmp_path):
+    skill = tmp_path / "skill"
+    skill.mkdir()
+    (skill / "SKILL.md").write_text("# s")
+    adapter_cls = install_adapter(monkeypatch, caps=NARROW_CAPS)
+    with pytest.raises(CapabilityUnsupported, match="skills"):
+        await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, skills=[skill])
+    assert adapter_cls.instances == []
+
+
+async def test_skill_folder_without_skill_md_rejected(monkeypatch, sandbox, tmp_path):
+    install_adapter(monkeypatch)
+    with pytest.raises(ValueError, match=r"SKILL\.md"):
+        await runtime.aagent(
+            Harness.CLAUDE_CODE, "hi", sandbox=sandbox, skills=[tmp_path]
+        )
+
+
+async def test_ask_without_handler_only_allowed_for_stream(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    with pytest.raises(ValueError, match="on_approval"):
+        await runtime.aagent(
+            Harness.CLAUDE_CODE, "hi", sandbox=sandbox, permissions="ask"
+        )
+    stream = runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, permissions="ask", stream=True
+    )
+    events = await _collect(stream)
+    assert isinstance(events[-1], Done)
+
+
+async def test_invalid_permissions_value(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    with pytest.raises(ValueError, match="permissions"):
+        await runtime.aagent(
+            Harness.CLAUDE_CODE, "hi", sandbox=sandbox, permissions="yolo"  # type: ignore[arg-type]
+        )
+
+
+# -- gateway routing (litellm_proxy/ prefix) ---------------------------------
+
+
+def test_litellm_proxy_prefix_routes_through_gateway_env(monkeypatch):
+    monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://gw.example.com/")
+    monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-test")
+    model, gateway = runtime.resolve_model_route("litellm_proxy/coder", None, None)
+    assert model == "coder"
+    assert gateway == runtime.GatewayTarget(
+        api_base="https://gw.example.com", api_key="sk-test"
+    )
+
+
+def test_litellm_proxy_call_args_win_over_env(monkeypatch):
+    monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://env.example.com")
+    monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-env")
+    _, gateway = runtime.resolve_model_route(
+        "litellm_proxy/coder", "sk-arg", "https://arg.example.com"
+    )
+    assert gateway == runtime.GatewayTarget(
+        api_base="https://arg.example.com", api_key="sk-arg"
+    )
+
+
+def test_litellm_proxy_without_base_raises(monkeypatch):
+    monkeypatch.delenv("LITELLM_PROXY_API_BASE", raising=False)
+    monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-test")
+    with pytest.raises(ValueError, match="LITELLM_PROXY_API_BASE"):
+        runtime.resolve_model_route("litellm_proxy/coder", None, None)
+
+
+def test_litellm_proxy_without_key_raises(monkeypatch):
+    monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://gw.example.com")
+    monkeypatch.setenv("LITELLM_PROXY_API_KEY", "  ")
+    with pytest.raises(ValueError, match="LITELLM_PROXY_API_KEY"):
+        runtime.resolve_model_route("litellm_proxy/coder", None, None)
+
+
+def test_plain_model_is_sdk_mode_even_with_gateway_env(monkeypatch):
+    monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://gw.example.com")
+    monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-test")
+    assert runtime.resolve_model_route("anthropic/claude-sonnet-4-5", None, None) == (
+        "anthropic/claude-sonnet-4-5",
+        None,
+    )
+
+
+def test_use_litellm_proxy_flag_routes_unprefixed_model(monkeypatch):
+    monkeypatch.setattr(runtime.litellm, "use_litellm_proxy", True)
+    monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://gw.example.com")
+    monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-test")
+    model, gateway = runtime.resolve_model_route("coder", None, None)
+    assert model == "coder" and gateway is not None
+
+
+async def test_gateway_passed_to_endpoint(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    monkeypatch.setenv("LITELLM_PROXY_API_BASE", "https://gw.example.com")
+    monkeypatch.setenv("LITELLM_PROXY_API_KEY", "sk-test")
+    await runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, model="litellm_proxy/m"
+    )
+    endpoint = FakeEndpoint.instances[0]
+    assert endpoint.gateway.api_key == "sk-test"
+    assert endpoint.model == "m"
+    assert endpoint.entered and endpoint.exited
+
+
+# -- event flow ---------------------------------------------------------------
+
+
+async def test_text_and_tool_events_flow_and_done_last(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch)
+    events = await _collect(
+        runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, stream=True)
+    )
+    kinds = [type(e).__name__ for e in events]
+    assert kinds == ["Text", "ToolCall", "ToolResult", "Text", "Done"]
+    assert sum(isinstance(e, Done) for e in events) == 1
+    result = events[-1].result
+    assert result.text == "hello world"
+    assert result.stop_reason == "done"
+    assert result.usage.input_tokens == 10 and result.usage.output_tokens == 5
+    assert result.cost == pytest.approx(0.25)
+    assert adapter_cls.instances[0].calls == ["start", "turn", "stop"]
+
+
+async def test_arun_returns_result(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    result = await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    assert result.text == "hello world"
+    assert len(result.events) == 4
+
+
+async def test_final_text_from_ctx_preferred(monkeypatch, sandbox):
+    async def script(adapter, ctx: SessionContext, prompt) -> AsyncIterator[Event]:
+        yield Text("partial")
+        ctx.final_text = "final answer"
+
+    install_adapter(monkeypatch, script=script)
+    result = await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    assert result.text == "final answer"
+
+
+async def test_endpointless_adapter_usage(monkeypatch, sandbox):
+    install_adapter(monkeypatch, uses_endpoint=False)
+    result = await runtime.aagent(Harness.DEEPAGENTS, "hi", sandbox=sandbox)
+    assert FakeEndpoint.instances == []
+    assert result.usage.calls == 1
+    assert result.cost == pytest.approx(0.25)
+
+
+async def test_stream_result_property(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    stream = runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, stream=True)
+    assert stream.result is None
+    await _collect(stream)
+    assert stream.result is not None and stream.result.text == "hello world"
+
+
+# -- stop reasons -------------------------------------------------------------
+
+
+async def _tool_loop(adapter, ctx, prompt) -> AsyncIterator[Event]:
+    for i in range(10):
+        yield ToolCall(id=str(i), name="bash", native_name="Bash", input={})
+
+
+async def test_max_turns_stop(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch, script=_tool_loop)
+    result = await runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, max_turns=3
+    )
+    assert result.stop_reason == "max_turns"
+    assert sum(isinstance(e, ToolCall) for e in result.events) == 3
+    assert "stop" in adapter_cls.instances[0].calls
+
+
+async def _slow(adapter, ctx, prompt) -> AsyncIterator[Event]:
+    yield Text("thinking")
+    await wait_forever()
+    yield Text("never")
+
+
+async def test_timeout_stop(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_slow)
+    result = await runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, timeout=0.2
+    )
+    assert result.stop_reason == "timeout"
+    assert result.text == "thinking"
+
+
+async def test_cancel_stop(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_slow)
+    stream = runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, stream=True)
+    events = []
+    async for event in stream:
+        events.append(event)
+        if isinstance(event, Text):
+            stream.cancel()
+    assert isinstance(events[-1], Done)
+    assert events[-1].stop_reason == "cancelled"
+
+
+async def _crash(adapter, ctx, prompt) -> AsyncIterator[Event]:
+    yield Text("partial")
+    raise RuntimeError("process exited with code 1")
+
+
+async def test_runtime_error_stop_reason(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_crash)
+    result = await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    assert result.stop_reason == "runtime_error"
+    assert "process exited with code 1" in result.text
+
+
+async def _missing_binary(adapter, ctx, prompt) -> AsyncIterator[Event]:
+    raise HarnessInstallFailed("claude not found on PATH")
+    yield Text("unreachable")  # pragma: no cover
+
+
+async def test_install_failed_propagates(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_missing_binary)
+    with pytest.raises(HarnessInstallFailed, match="claude"):
+        await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    assert FakeEndpoint.instances[0].exited
+
+
+# -- approvals ----------------------------------------------------------------
+
+
+async def test_approval_on_approval_allow(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch, script=script_approval)
+    result = await runtime.aagent(
+        Harness.CLAUDE_CODE,
+        "hi",
+        sandbox=sandbox,
+        permissions="ask",
+        on_approval=lambda approval: approval.tool == "bash",
+    )
+    assert adapter_cls.instances[0].approvals[0][0] is True
+    assert result.text == "allowed"
+
+
+async def test_approval_async_handler_deny(monkeypatch, sandbox):
+    async def handler(approval: Approval) -> bool:
+        await asyncio.sleep(0)
+        return False
+
+    adapter_cls = install_adapter(monkeypatch, script=script_approval)
+    result = await runtime.aagent(
+        Harness.CLAUDE_CODE,
+        "hi",
+        sandbox=sandbox,
+        permissions="ask",
+        on_approval=handler,
+    )
+    assert adapter_cls.instances[0].approvals[0][0] is False
+    assert result.text == "denied"
+
+
+async def test_approval_handler_raises_denies(monkeypatch, sandbox):
+    def handler(approval: Approval) -> bool:
+        raise RuntimeError("boom")
+
+    adapter_cls = install_adapter(monkeypatch, script=script_approval)
+    result = await runtime.aagent(
+        Harness.CLAUDE_CODE,
+        "hi",
+        sandbox=sandbox,
+        permissions="ask",
+        on_approval=handler,
+    )
+    allowed, reason = adapter_cls.instances[0].approvals[0]
+    assert allowed is False and "boom" in reason
+    assert result.stop_reason == "done"
+
+
+async def test_stream_consumer_answers_approval(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch, script=script_approval)
+    stream = runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, permissions="ask", stream=True
+    )
+    async for event in stream:
+        if isinstance(event, Approval):
+            event.allow()
+    assert adapter_cls.instances[0].approvals[0][0] is True
+
+
+async def test_unanswered_approval_denied(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch, script=script_approval)
+    events = await _collect(
+        runtime.aagent(
+            Harness.CLAUDE_CODE, "hi", sandbox=sandbox, permissions="ask", stream=True
+        )
+    )
+    allowed, reason = adapter_cls.instances[0].approvals[0]
+    assert allowed is False and "not answered" in reason
+    assert isinstance(events[-1], Done)
+
+
+async def test_approval_without_ask_denied_in_run(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch, script=script_approval)
+    await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    assert adapter_cls.instances[0].approvals[0][0] is False
+
+
+# -- structured output --------------------------------------------------------
+
+
+def _answer_script(text: str, output_json: str | None = None):
+    async def script(adapter, ctx, prompt) -> AsyncIterator[Event]:
+        yield Text(text)
+        ctx.output_json = output_json
+
+    return script
+
+
+async def test_structured_output_from_text(monkeypatch, sandbox):
+    install_adapter(
+        monkeypatch, script=_answer_script('Sure. {"x": 1} then {"value": 42} done')
+    )
+    result = await runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, output=Answer
+    )
+    assert result.output == Answer(value=42)
+
+
+async def test_structured_output_from_ctx_output_json(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_answer_script("ok", '{"value": 7}'))
+    result = await runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, output=Answer
+    )
+    assert result.output == Answer(value=7)
+
+
+async def test_structured_output_invalid_carries_result(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_answer_script('{"value": "nope"}'))
+    with pytest.raises(OutputInvalid) as info:
+        await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, output=Answer)
+    assert info.value.raw == '{"value": "nope"}'
+    assert info.value.result is not None
+    assert info.value.result.text == '{"value": "nope"}'
+
+
+async def test_structured_output_missing_json(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_answer_script("no json here"))
+    with pytest.raises(OutputInvalid, match="no JSON"):
+        await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, output=Answer)
+
+
+async def test_stream_yields_done_before_output_invalid(monkeypatch, sandbox):
+    install_adapter(monkeypatch, script=_answer_script("nothing"))
+    stream = runtime.aagent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, output=Answer, stream=True
+    )
+    seen: list[object] = []
+    with pytest.raises(OutputInvalid):
+        await _drain_into(stream, seen)
+    assert isinstance(seen[-1], Done)
+
+
+async def _drain_into(stream, seen: list[object]) -> None:
+    async for event in stream:
+        seen.append(event)
+
+
+def test_last_json_object():
+    assert runtime.last_json_object('a {"a": {"b": 1}} b {"c": 2}') == '{"c": 2}'
+    assert runtime.last_json_object("{broken") is None
+
+
+# -- files --------------------------------------------------------------------
+
+
+async def _edit_files(adapter, ctx, prompt) -> AsyncIterator[Event]:
+    root = ctx.sandbox.workdir
+    with open(os.path.join(root, "new.txt"), "w") as fh:
+        fh.write("new\n")
+    with open(os.path.join(root, "keep.txt"), "w") as fh:
+        fh.write("changed\n")
+    os.remove(os.path.join(root, "gone.txt"))
+    yield FileChange(path="new.txt", kind="created", diff=None)
+    yield Text("edited")
+
+
+async def test_file_changes_emitted_once(monkeypatch, sandbox, tmp_path):
+    (tmp_path / "keep.txt").write_text("original\n")
+    (tmp_path / "gone.txt").write_text("bye\n")
+    install_adapter(monkeypatch, script=_edit_files)
+    events = await _collect(
+        runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, stream=True)
+    )
+    file_events = [e for e in events if isinstance(e, FileChange)]
+    assert sorted((e.path, e.kind) for e in file_events) == [
+        ("gone.txt", "deleted"),
+        ("keep.txt", "modified"),
+        ("new.txt", "created"),
+    ]
+    result = events[-1].result
+    by_path = {f.path: f for f in result.files}
+    assert set(by_path) == {"gone.txt", "keep.txt", "new.txt"}
+    assert "+changed" in by_path["keep.txt"].diff
+    assert "-bye" in by_path["gone.txt"].diff
+    assert "+new" in by_path["new.txt"].diff
+    assert isinstance(events[-1], Done)
+
+
+# -- sessions -----------------------------------------------------------------
+
+
+async def test_session_multi_turn_cost(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch)
+    async with runtime.aagent_session(Harness.CLAUDE_CODE, sandbox=sandbox) as session:
+        first = await session.arun("one")
+        second = await session.arun("two")
+        assert first.cost == pytest.approx(0.25)
+        assert second.cost == pytest.approx(0.25)
+        assert second.usage.input_tokens == 10
+        assert session.cost == pytest.approx(0.5)
+        assert session.usage.calls == 2
+        assert await session.history() == [
+            {"role": "user", "content": "one"},
+            {"role": "user", "content": "two"},
+        ]
+    adapter = adapter_cls.instances[0]
+    assert adapter.calls == ["start", "turn", "turn", "stop"]
+    assert len(FakeEndpoint.instances) == 1
+    with pytest.raises(SessionClosed):
+        await session.arun("three")
+
+
+async def test_await_asession(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    session = await runtime.aagent_session(Harness.CLAUDE_CODE, sandbox=sandbox)
+    result = await session.arun("hi")
+    await session.aclose()
+    assert result.text == "hello world"
+
+
+async def test_session_restarts_after_timeout(monkeypatch, sandbox):
+    calls = {"n": 0}
+
+    async def script(adapter, ctx, prompt) -> AsyncIterator[Event]:
+        calls["n"] += 1
+        if calls["n"] == 1:
+            await wait_forever()
+        yield Text("ok")
+
+    adapter_cls = install_adapter(monkeypatch, script=script)
+    async with runtime.aagent_session(
+        Harness.CLAUDE_CODE, sandbox=sandbox, timeout=0.2
+    ) as session:
+        assert (await session.arun("one")).stop_reason == "timeout"
+        assert (await session.arun("two")).text == "ok"
+    adapter = adapter_cls.instances[0]
+    assert adapter.calls[:4] == ["start", "turn", "stop", "start"]
+    assert adapter.resumed_with == "native-123"
+
+
+async def test_detach_state_round_trip_resume(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch)
+    async with runtime.aagent_session(
+        Harness.CLAUDE_CODE, sandbox=sandbox, model="m1"
+    ) as session:
+        await session.arun("one")
+        state = await session.adetach()
+    data = state.dumps()
+    assert b"sk-" not in data
+    restored = State.loads(data)
+    assert restored == state and restored.native_session_id == "native-123"
+
+    async with runtime.aagent_resume(data, sandbox=sandbox) as resumed:
+        await resumed.arun("two")
+    new_adapter = adapter_cls.instances[-1]
+    assert new_adapter.calls[:2] == ["start", "resume"]
+    assert new_adapter.resumed_with == "native-123"
+    assert resumed.config.model == "m1"
+
+
+async def test_resume_requires_capability(monkeypatch, sandbox):
+    install_adapter(
+        monkeypatch,
+        caps=NARROW_CAPS.__class__(**{**NARROW_CAPS.__dict__, "resume": False}),
+    )
+    state = State(harness=Harness.CODEX, native_session_id="x", workdir="/tmp")
+    with pytest.raises(CapabilityUnsupported, match="resume"):
+        runtime.aagent_resume(state, sandbox=sandbox)
+
+
+async def test_resume_state_without_native_id(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    state = State(harness=Harness.CODEX, native_session_id=None, workdir="/tmp")
+    with pytest.raises(StateIncompatible):
+        runtime.aagent_resume(state, sandbox=sandbox)
+
+
+async def test_history_requires_capability(monkeypatch, sandbox):
+    install_adapter(monkeypatch, caps=NARROW_CAPS)
+    async with runtime.aagent_session(Harness.CLAUDE_CODE, sandbox=sandbox) as session:
+        with pytest.raises(CapabilityUnsupported):
+            await session.history()
+
+
+def test_capabilities_uses_registry(monkeypatch):
+    install_adapter(monkeypatch, caps=NARROW_CAPS)
+    assert runtime.agent_capabilities(Harness.CODEX) is NARROW_CAPS
+    with pytest.raises(TypeError):
+        runtime.agent_capabilities("codex")  # type: ignore[arg-type]
+
+
+def test_fake_adapter_is_a_harness_adapter():
+    assert issubclass(FakeAdapter, runtime.BaseHarnessHandler)
+
+
+async def test_turn_keeps_every_event_when_queue_overflows(monkeypatch, sandbox):
+    """A turn that emits more events than the queue holds must not drop any of them."""
+    total = 40
+    monkeypatch.setattr(runtime, "HARNESS_EVENT_QUEUE_MAX_SIZE", 4)
+
+    async def burst(adapter, ctx, prompt):
+        for i in range(total):
+            yield Text(f"{i},")
+
+    install_adapter(monkeypatch, script=burst)
+    result = await runtime.aagent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    texts = [e.delta for e in result.events if isinstance(e, Text)]
+    assert texts == [f"{i}," for i in range(total)]
+    assert result.stop_reason == "done"

tests/unit/harness/test_sync.py

@@ -0,0 +1,114 @@
+"""Tests for litellm/harness/sync.py: the sync bridge over the async runtime."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+
+import pytest
+
+from litellm.harness import sync
+from litellm.harness.types import Approval, Done, Harness, State, Text
+from tests.unit.harness.core_fakes import (
+    FakeSandbox,
+    install_adapter,
+    script_approval,
+)
+
+
+@pytest.fixture
+def sandbox(tmp_path) -> FakeSandbox:
+    return FakeSandbox(str(tmp_path))
+
+
+async def _call_run_in_loop(sandbox: FakeSandbox) -> None:
+    sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+
+
+def test_sync_run_from_plain_code(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    result = sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    assert result.text == "hello world"
+    assert result.stop_reason == "done"
+
+
+def test_sync_stream_from_plain_code(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    stream = sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, stream=True)
+    events = list(stream)
+    assert isinstance(events[-1], Done)
+    assert [e.delta for e in events if isinstance(e, Text)] == ["hello ", "world"]
+    assert stream.result is not None and stream.result.text == "hello world"
+    assert list(stream) == []
+
+
+def test_sync_stream_answers_approval(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch, script=script_approval)
+    for event in sync.agent(
+        Harness.CLAUDE_CODE, "hi", sandbox=sandbox, permissions="ask", stream=True
+    ):
+        if isinstance(event, Approval):
+            event.allow()
+    assert adapter_cls.instances[0].approvals[0][0] is True
+
+
+def test_sync_stream_close_early(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch)
+    with sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, stream=True) as stream:
+        next(stream)
+    assert adapter_cls.instances[0].calls[-1] == "stop"
+
+
+def test_sync_validation_errors_raise_eagerly(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    with pytest.raises(TypeError, match=r"Harness\.OPENCODE"):
+        sync.agent("opencode", "hi", sandbox=sandbox, stream=True)  # type: ignore[arg-type]
+
+
+def test_sync_session_multi_turn_and_detach(monkeypatch, sandbox):
+    adapter_cls = install_adapter(monkeypatch)
+    with sync.agent_session(Harness.CLAUDE_CODE, sandbox=sandbox) as session:
+        session.run("one")
+        events = list(session.stream("two"))
+        assert isinstance(events[-1], Done)
+        assert session.cost == pytest.approx(0.5)
+        assert len(session.history()) == 2
+        state = session.detach()
+    assert isinstance(state, State)
+    with sync.agent_resume(state.dumps(), sandbox=sandbox) as resumed:
+        assert resumed.run("three").text == "hello world"
+    assert adapter_cls.instances[-1].resumed_with == "native-123"
+
+
+def test_sync_session_stop_returns_state(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    session = sync.agent_session(Harness.CLAUDE_CODE, sandbox=sandbox).start()
+    session.run("one")
+    state = session.stop()
+    assert state.native_session_id == "native-123"
+
+
+def test_single_background_loop_thread(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    first = sync._LOOP.loop()
+    sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    assert sync._LOOP.loop() is first
+    names = [t.name for t in threading.enumerate()]
+    assert names.count("litellm-harness-loop") == 1
+
+
+async def test_run_inside_event_loop_raises(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    with pytest.raises(RuntimeError, match="aagent"):
+        sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox)
+    with pytest.raises(RuntimeError, match="aagent"):
+        sync.agent(Harness.CLAUDE_CODE, "hi", sandbox=sandbox, stream=True)
+    with pytest.raises(RuntimeError, match="aagent_session"):
+        sync.agent_session(Harness.CLAUDE_CODE, sandbox=sandbox)
+
+
+def test_run_inside_asyncio_run_raises(monkeypatch, sandbox):
+    install_adapter(monkeypatch)
+    with pytest.raises(RuntimeError, match=r"await litellm\.aagent"):
+        asyncio.run(_call_run_in_loop(sandbox))

tests/unit/harness/test_types.py

@@ -0,0 +1,90 @@
+"""Tests for litellm/harness/types.py."""
+
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from litellm.harness.errors import StateIncompatible
+from litellm.harness.types import (
+    Approval,
+    Done,
+    Harness,
+    Result,
+    State,
+    Usage,
+    require_harness,
+)
+
+
+def test_harness_is_plain_enum():
+    assert Harness.CODEX.value == "codex"
+    assert not isinstance(Harness.CODEX, str)
+
+
+@pytest.mark.parametrize(
+    "given,hint",
+    [
+        ("codex", "Harness.CODEX"),
+        ("claude-code", "Harness.CLAUDE_CODE"),
+        ("OPENCODE", "Harness.OPENCODE"),
+    ],
+)
+def test_require_harness_hint(given, hint):
+    with pytest.raises(TypeError, match=hint):
+        require_harness(given)
+
+
+def test_require_harness_no_hint_for_unknown():
+    with pytest.raises(TypeError) as info:
+        require_harness(42)
+    assert "Did you mean" not in str(info.value)
+    assert require_harness(Harness.DEEPAGENTS) is Harness.DEEPAGENTS
+
+
+def test_usage_total_tokens():
+    assert Usage(input_tokens=3, output_tokens=4, calls=1).total_tokens == 7
+
+
+def test_done_exposes_result_fields():
+    result = Result(
+        text="t",
+        output=None,
+        files=[],
+        events=[],
+        usage=Usage(1, 2, 1),
+        cost=0.5,
+        stop_reason="done",
+        session_id="s",
+    )
+    done = Done(result)
+    assert done.usage.total_tokens == 3
+    assert done.cost == 0.5
+    assert done.stop_reason == "done"
+
+
+def test_state_round_trip_and_errors():
+    state = State(Harness.CODEX, "thread-1", "/work", model="gpt")
+    assert State.loads(state.dumps()) == state
+    with pytest.raises(StateIncompatible):
+        State.loads(b"not json")
+    with pytest.raises(StateIncompatible):
+        State.loads(b'{"harness": "nope", "version": 1, "workdir": "/"}')
+    with pytest.raises(StateIncompatible, match="version"):
+        State.loads(b'{"harness": "codex", "version": 99, "workdir": "/"}')
+
+
+async def test_approval_allow_deny_once():
+    approval = Approval(tool="bash", input={})
+    assert not approval.answered
+    approval.allow()
+    approval.deny("late")
+    assert await approval.wait() == (True, "")
+    assert approval.answered
+
+
+async def test_approval_resolved_from_other_thread():
+    approval = Approval(tool="bash", input={})
+    await asyncio.to_thread(approval.deny, "nope")
+    assert await approval.wait() == (False, "nope")

tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py

@@ -28,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     sanitize_messages_for_tool_calling,
 )
 from litellm.types.llms.openai import ChatCompletionToolMessage
+from litellm.utils import validate_and_fix_openai_messages
 
 
 def _get_gemini_function_response_inline_data_parts(result):
@@ -4095,3 +4096,168 @@ def test_is_unsignable_thinking_block_treats_whitespace_only_as_empty():
     }
 
     assert is_unsignable_thinking_block(whitespace_only_block) is True
+
+
+_CONTENT_LESS_USER_MESSAGES: Final = ({"role": "user"}, {"role": "user", "content": None})
+_CONTENT_LESS_TOOL_MESSAGES: Final = (
+    {"role": "tool", "tool_call_id": "call_1"},
+    {"role": "tool", "tool_call_id": "call_1", "content": None},
+)
+_BOSTON_WEATHER_TOOL_CALL_TURN: Final = (
+    {"role": "user", "content": "What is the weather in Boston?"},
+    {
+        "role": "assistant",
+        "tool_calls": [
+            {
+                "id": "call_1",
+                "type": "function",
+                "function": {"name": "get_weather", "arguments": '{"city": "Boston"}'},
+            }
+        ],
+    },
+)
+
+
+def _conversation_around(
+    content_less_user_message: dict[str, object],
+) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
+    with_message: Final = [
+        {"role": "user", "content": "What is the capital of France?"},
+        content_less_user_message,
+        {"role": "assistant", "content": "Paris."},
+        {"role": "user", "content": "And of Spain?"},
+    ]
+    without_message: Final = [message for message in with_message if message is not content_less_user_message]
+    return validate_and_fix_openai_messages(with_message), validate_and_fix_openai_messages(without_message)
+
+
+@pytest.mark.parametrize("content_less_user_message", _CONTENT_LESS_USER_MESSAGES)
+def test_bedrock_converse_messages_pt_user_message_without_content_adds_no_block(
+    content_less_user_message: dict[str, object],
+):
+    with_message, without_message = _conversation_around(content_less_user_message)
+
+    assert _bedrock_converse_messages_pt(
+        messages=with_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock"
+    ) == _bedrock_converse_messages_pt(messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("content_less_user_message", _CONTENT_LESS_USER_MESSAGES)
+async def test_bedrock_converse_messages_pt_async_user_message_without_content_adds_no_block(
+    content_less_user_message: dict[str, object],
+):
+    with_message, without_message = _conversation_around(content_less_user_message)
+
+    assert await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
+        messages=with_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock"
+    ) == await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
+        messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock"
+    )
+
+
+@pytest.mark.parametrize("content_less_tool_message", _CONTENT_LESS_TOOL_MESSAGES)
+def test_bedrock_converse_messages_pt_tool_message_without_content_yields_empty_tool_result(
+    content_less_tool_message: dict[str, object],
+):
+    result: Final = _bedrock_converse_messages_pt(
+        messages=validate_and_fix_openai_messages([*_BOSTON_WEATHER_TOOL_CALL_TURN, content_less_tool_message]),
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+    )
+
+    tool_result: Final = result[-1]["content"][0]["toolResult"]
+    assert result[-1]["role"] == "user"
+    assert tool_result["toolUseId"] == "call_1"
+    assert tool_result["content"] == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("content_less_tool_message", _CONTENT_LESS_TOOL_MESSAGES)
+async def test_bedrock_converse_messages_pt_async_tool_message_without_content_yields_empty_tool_result(
+    content_less_tool_message: dict[str, object],
+):
+    result: Final = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
+        messages=validate_and_fix_openai_messages([*_BOSTON_WEATHER_TOOL_CALL_TURN, content_less_tool_message]),
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+    )
+
+    tool_result: Final = result[-1]["content"][0]["toolResult"]
+    assert tool_result["toolUseId"] == "call_1"
+    assert tool_result["content"] == []
+
+
+def test_bedrock_converse_messages_pt_blank_user_text_sends_the_continue_message_text():
+    continue_message: Final = {"role": "user", "content": "Please continue."}
+    blank_last_turn: Final = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi."},
+        {"role": "user", "content": "   "},
+    ]
+    explicit_last_turn: Final = [*blank_last_turn[:2], continue_message]
+
+    assert _bedrock_converse_messages_pt(
+        messages=blank_last_turn,
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+        user_continue_message=continue_message,
+    ) == _bedrock_converse_messages_pt(
+        messages=explicit_last_turn,
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+        user_continue_message=continue_message,
+    )
+
+
+@pytest.mark.parametrize("content_less_user_message", _CONTENT_LESS_USER_MESSAGES)
+def test_bedrock_converse_messages_pt_lone_content_less_user_turn_sends_the_continue_message(
+    content_less_user_message: dict[str, object],
+):
+    continue_message: Final = {"role": "user", "content": "Please continue."}
+
+    assert _bedrock_converse_messages_pt(
+        messages=validate_and_fix_openai_messages([content_less_user_message]),
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+        user_continue_message=continue_message,
+    ) == _bedrock_converse_messages_pt(
+        messages=[continue_message],
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+        user_continue_message=continue_message,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("content_less_user_message", _CONTENT_LESS_USER_MESSAGES)
+async def test_bedrock_converse_messages_pt_async_lone_content_less_user_turn_continues_under_modify_params(
+    content_less_user_message: dict[str, object], monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setattr(litellm, "modify_params", True)
+
+    assert await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
+        messages=validate_and_fix_openai_messages([content_less_user_message]),
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+    ) == await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
+        messages=[{"role": "user", "content": ""}],
+        model="anthropic.claude-haiku-4-5",
+        llm_provider="bedrock",
+    )
+
+
+@pytest.mark.parametrize("content_less_user_message", _CONTENT_LESS_USER_MESSAGES)
+def test_bedrock_converse_messages_pt_lone_content_less_user_turn_adds_no_block_without_a_continue_message(
+    content_less_user_message: dict[str, object], monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setattr(litellm, "modify_params", False)
+
+    assert (
+        _bedrock_converse_messages_pt(
+            messages=validate_and_fix_openai_messages([content_less_user_message]),
+            model="anthropic.claude-haiku-4-5",
+            llm_provider="bedrock",
+        )
+        == []
+    )

tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py

@@ -2,7 +2,9 @@ import glob
 import os
 import re
 import sys
+import threading
 from pathlib import Path
+from typing import Final
 
 import pytest
 
@@ -1024,3 +1026,143 @@ class TestJWTKeyMappingCascade:
                 f"{path} must declare onDelete: Cascade on the JWT key mapping "
                 "relation (issue #33702)"
             )
+
+
+
+class TestBuildRequestLogIndexes:
+    """The migration job hands the index build the direct database URL and the schema
+    the migrations target, waits for it, and reports its result."""
+
+    @pytest.fixture
+    def builds(self):
+        return []
+
+    @pytest.fixture
+    def build(self, builds):
+        def record(database_url: str, schema: str) -> bool:
+            builds.append((database_url, schema))
+            return True
+
+        return record
+
+    def test_the_build_gets_the_direct_url_without_prisma_params_and_the_prisma_schema(self, monkeypatch, builds, build):
+        monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@pooler:6543/db?schema=tenant&pgbouncer=true")
+        monkeypatch.setenv("DIRECT_URL", "postgresql://u:p@primary:5432/db?connection_limit=1")
+
+        assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
+
+        assert builds == [("postgresql://u:p@primary:5432/db", "tenant")]
+
+    def test_the_build_defaults_to_the_database_url_and_the_public_schema(self, monkeypatch, builds, build):
+        monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
+        monkeypatch.delenv("DIRECT_URL", raising=False)
+
+        assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
+
+        assert builds == [("postgresql://u:p@primary:5432/db", "public")]
+
+    def test_a_build_that_leaves_indexes_missing_is_reported_so_the_job_reruns(self, monkeypatch):
+        monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@primary:5432/db")
+
+        assert ProxyExtrasDBManager.build_request_log_indexes(build=lambda url, schema: False) is False
+
+    def test_without_a_database_url_nothing_is_built(self, monkeypatch, builds, build):
+        monkeypatch.delenv("DATABASE_URL", raising=False)
+
+        assert ProxyExtrasDBManager.build_request_log_indexes(build=build) is True
+
+        assert builds == []
+
+
+class TestStartRequestLogIndexBuild:
+    """A serving proxy that ran the migrations starts the index build on a daemon thread
+    and goes on to serve while it runs."""
+
+    def test_the_build_runs_on_a_daemon_thread_that_does_not_hold_up_the_caller(self):
+        release: Final = threading.Event()
+        builds: Final[list[str]] = []  # mutable-ok: the builder thread hands back the thread it ran on
+
+        def build() -> bool:
+            assert release.wait(5), "the caller never came back from start_request_log_index_build"
+            builds.append(threading.current_thread().name)
+            return True
+
+        thread: Final = ProxyExtrasDBManager.start_request_log_index_build(build=build)
+
+        assert builds == [], "the build ran before start_request_log_index_build returned"
+        assert thread.daemon is True
+        release.set()
+        thread.join(5)
+        assert builds == ["litellm-request-log-indexes"]
+
+
+class TestRunMigrationJob:
+    """`run_migration_job` is `setup_database` followed by the index build, each step's
+    result deciding whether the job reports success."""
+
+    @pytest.fixture
+    def calls(self):
+        return []
+
+    @pytest.fixture
+    def setup(self, calls):
+        def record(result: bool):
+            def setup_database(use_migrate: bool, use_v2_resolver: bool) -> bool:
+                calls.append(("setup", use_migrate, use_v2_resolver))
+                return result
+
+            return setup_database
+
+        return record
+
+    @pytest.fixture
+    def build(self, calls):
+        def record(result: bool):
+            def build_request_log_indexes() -> bool:
+                calls.append(("build",))
+                return result
+
+            return build_request_log_indexes
+
+        return record
+
+    def test_the_job_builds_the_indexes_after_the_migrations_succeed(self, calls, setup, build):
+        assert ProxyExtrasDBManager.run_migration_job(True, False, setup=setup(True), build=build(True)) is True
+
+        assert calls == [("setup", True, False), ("build",)]
+
+    def test_the_job_fails_without_building_when_the_migrations_fail(self, calls, setup, build):
+        assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(False), build=build(True)) is False
+
+        assert calls == [("setup", True, True)]
+
+    def test_the_job_fails_when_an_index_could_not_be_built(self, calls, setup, build):
+        assert ProxyExtrasDBManager.run_migration_job(True, True, setup=setup(True), build=build(False)) is False
+
+        assert calls == [("setup", True, True), ("build",)]
+
+
+class TestMigrationJobOwnedDrift:
+    JOB_INDEXES = (
+        "-- CreateIndex\n"
+        'CREATE INDEX "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
+        "\n-- CreateIndex\n"
+        'CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
+    )
+
+    def test_a_plain_spend_logs_table_only_loses_the_migration_job_indexes(self):
+        filtered = ProxyExtrasDBManager._filter_migration_job_owned_drift(
+            _PARTITIONED_DRIFT_SQL + self.JOB_INDEXES, partitioned=False
+        )
+        assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
+        assert "LiteLLM_SpendLogs_api_key_startTime_idx" not in filtered
+        assert 'PRIMARY KEY ("request_id")' in filtered
+
+    def test_a_partitioned_spend_logs_table_also_loses_its_partitioning_artifacts(self):
+        filtered = ProxyExtrasDBManager._filter_migration_job_owned_drift(
+            _PARTITIONED_DRIFT_SQL + self.JOB_INDEXES, partitioned=True
+        )
+        assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
+        assert 'PRIMARY KEY ("request_id")' not in filtered
+        assert "LiteLLM_SpendLogs_legacy" not in filtered
+        assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN     "updated_by" TEXT;' in filtered

tests/unit/litellm_proxy_extras/test_request_log_indexes.py

@@ -0,0 +1,149 @@
+import re
+from pathlib import Path
+from typing import Final
+
+import pytest
+from litellm_proxy_extras.migration_recovery import is_inert_migration
+from litellm_proxy_extras.request_log_indexes import (
+    REQUEST_LOG_INDEXES,
+    RequestLogIndex,
+    filter_request_log_index_diff,
+)
+
+PACKAGE: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras"
+SCHEMA: Final = PACKAGE / "schema.prisma"
+INERT_MIGRATIONS: Final = (
+    "20260823000000_add_spend_logs_api_key_starttime_index",
+    "20260831120001_spend_logs_litellm_call_id_index",
+)
+CALL_ID_INDEX: Final = RequestLogIndex(
+    "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'
+)
+
+
+def _prisma_indexes_of(schema: str, model: str) -> frozenset[str]:
+    """The index names Prisma derives for a model's @@index declarations: <model>_<columns joined by _>_idx."""
+    body: Final = re.search(rf"model {model} \{{(.*?)\n\}}", schema, re.DOTALL)
+    assert body is not None, model
+    declarations: Final[tuple[str, ...]] = tuple(
+        match.group(1) for match in re.finditer(r"@@index\(\[([^\]]+)\]\)", body.group(1))
+    )
+    return frozenset(
+        f"{model}_{'_'.join(column.strip() for column in columns.split(','))}_idx" for columns in declarations
+    )
+
+
+class TestTheIndexList:
+    def test_every_migration_job_index_is_declared_in_the_prisma_schema_under_the_same_name(self):
+        schema: Final = SCHEMA.read_text()
+        for index in REQUEST_LOG_INDEXES:
+            assert index.name in _prisma_indexes_of(schema, index.table), index
+
+    @pytest.mark.parametrize("name", INERT_MIGRATIONS)
+    def test_the_migrations_that_used_to_build_these_indexes_run_no_sql(self, name: str):
+        assert is_inert_migration((PACKAGE / "migrations" / name / "migration.sql").read_text())
+
+
+class TestIsInertMigration:
+    @pytest.mark.parametrize(
+        "script",
+        (
+            "",
+            "-- only a comment\n",
+            "/* block */\n-- line\n",
+            "-- a semicolon; in a comment\n",
+            ";\n;",
+            "-- why\nSELECT 1;\n",
+            "select 1",
+        ),
+        ids=(
+            "empty",
+            "line-comment",
+            "both-comments",
+            "semicolon-in-comment",
+            "bare-separators",
+            "select-1",
+            "lowercase",
+        ),
+    )
+    def test_comments_and_a_select_1_alone_are_inert(self, script: str):
+        assert is_inert_migration(script) is True
+
+    @pytest.mark.parametrize(
+        "script",
+        (
+            "SELECT 2;",
+            'SELECT 1 FROM "LiteLLM_SpendLogs";',
+            '-- comment\nCREATE INDEX "ix" ON "t" ("a");',
+            "/* c */ ALTER TABLE t ADD COLUMN a TEXT",
+            'SELECT 1; DROP INDEX "ix";',
+        ),
+        ids=("select-2", "select-from", "index-after-comment", "alter-after-block-comment", "drop-after-select-1"),
+    )
+    def test_any_statement_is_not_inert(self, script: str):
+        assert is_inert_migration(script) is False
+
+
+class TestPartitionIndexName:
+    def test_a_partition_gets_the_name_postgres_would_give_an_inherited_index(self):
+        assert CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p2026_09") == (
+            "LiteLLM_SpendLogs_p2026_09_litellm_call_id_idx"
+        )
+
+    def test_an_index_not_prefixed_by_its_table_keeps_its_whole_name(self):
+        index = RequestLogIndex("LiteLLM_SpendLogs", "call_id_lookup", '("litellm_call_id")')
+        assert index.partition_index_name("LiteLLM_SpendLogs_pdefault") == "LiteLLM_SpendLogs_pdefault_call_id_lookup"
+
+    def test_a_long_name_is_cut_to_63_bytes_with_a_digest_that_keeps_partitions_apart(self):
+        first = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_09")
+        second = CALL_ID_INDEX.partition_index_name("LiteLLM_SpendLogs_p" + "x" * 50 + "_2026_10")
+        assert len(first.encode()) == 63 and len(second.encode()) == 63
+        assert first != second
+        assert first.startswith("LiteLLM_SpendLogs_p") and first[-9] == "_"
+
+    def test_the_byte_limit_counts_multibyte_characters(self):
+        name = CALL_ID_INDEX.partition_index_name("é" * 40)
+        assert len(name.encode()) <= 63 and len(name) < 63
+
+
+class TestColumns:
+    def test_the_columns_are_the_quoted_names_of_the_definition_in_order(self):
+        index = RequestLogIndex(
+            "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'
+        )
+        assert index.columns == ("api_key", "startTime")
+
+    def test_every_migration_job_index_names_at_least_one_column(self):
+        assert all(index.columns for index in REQUEST_LOG_INDEXES)
+
+
+DRIFT_WITH_BOTH_INDEXES: Final = (
+    "-- CreateIndex\n"
+    'CREATE INDEX "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");\n'
+    "\n"
+    "-- CreateIndex\n"
+    'CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");\n'
+)
+
+
+class TestFilterRequestLogIndexDiff:
+    def test_a_drift_script_that_only_creates_the_migration_job_indexes_becomes_empty(self):
+        assert filter_request_log_index_diff(DRIFT_WITH_BOTH_INDEXES) == ""
+
+    def test_other_statements_survive_with_the_migration_job_indexes_removed(self):
+        other: Final = '-- AlterTable\nALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN     "updated_by" TEXT;\n'
+        filtered = filter_request_log_index_diff(other + DRIFT_WITH_BOTH_INDEXES)
+        assert 'ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN     "updated_by" TEXT;' in filtered
+        assert "LiteLLM_SpendLogs_litellm_call_id_idx" not in filtered
+        assert "LiteLLM_SpendLogs_api_key_startTime_idx" not in filtered
+
+    def test_an_index_of_another_name_on_spend_logs_is_kept(self):
+        sql: Final = 'CREATE INDEX "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user");\n'
+        assert filter_request_log_index_diff(sql) == sql
+
+    def test_a_drop_of_a_migration_job_index_is_kept_for_the_operator_to_see(self):
+        sql: Final = 'DROP INDEX "LiteLLM_SpendLogs_litellm_call_id_idx";\n'
+        assert filter_request_log_index_diff(sql) == sql
+
+    def test_an_empty_script_stays_empty(self):
+        assert filter_request_log_index_diff("") == ""

tests/unit/llms/bedrock/chat/test_converse_transformation.py

@@ -7769,6 +7769,17 @@ def test_mid_conversation_system_entry_without_text_is_dropped(empty_content):
     assert out_messages == [{"role": "user", "content": "hi"}, {"role": "user", "content": "done"}]
 
 
+def test_system_entry_without_content_key_transforms_like_an_empty_one():
+    config = AmazonConverseConfig()
+    leading_without_key = [{"role": "system"}, {"role": "user", "content": "hi"}]
+    leading_empty = [{"role": "system", "content": ""}, {"role": "user", "content": "hi"}]
+    assert config._transform_system_message(leading_without_key) == config._transform_system_message(leading_empty)
+    assert config._transform_system_message(leading_without_key) == ([{"role": "user", "content": "hi"}], [])
+    mid_without_key = [{"role": "user", "content": "hi"}, {"role": "system"}, {"role": "user", "content": "done"}]
+    mid_empty = [{"role": "user", "content": "hi"}, {"role": "system", "content": ""}, {"role": "user", "content": "done"}]
+    assert config._transform_system_message(mid_without_key) == config._transform_system_message(mid_empty)
+
+
 def _thinking_reply(text: str) -> dict:
     return {
         "role": "assistant",

tests/unit/llms/claude_code/harness/fixtures/api_error.jsonl

@@ -0,0 +1,3 @@
+{"type": "system", "subtype": "init", "cwd": "/workspace", "session_id": "53af83ee-c3e1-4b96-a70a-f15b6cb6c794", "tools": ["Task", "Bash", "CronCreate", "CronDelete", "CronList", "Edit", "EnterWorktree", "ExitWorktree", "ListAgents", "NotebookEdit", "Read", "ReportFindings", "ScheduleWakeup", "SendMessage", "Skill", "TaskStop", "WebFetch", "WebSearch", "Workflow", "Write"], "mcp_servers": [], "model": "does-not-exist-model-xyz", "permissionMode": "bypassPermissions", "slash_commands": ["deep-research", "dataviz", "update-config", "verify", "debug", "code-review", "simplify", "batch", "fewer-permission-prompts", "doctor", "loop", "claude-api", "workflow-authoring", "run", "run-skill-generator", "agents", "auto-mode-setup", "autocompact", "clear", "color", "compact", "config", "output-style", "context", "effort", "fast", "focus", "heapdump", "init", "mcp", "model", "__remote-workflow", "workflow-launch-exec", "reload-plugins", "reload-skills", "rename", "security-review", "usage", "insights", "recap", "goal", "list-agents", "team-onboarding"], "terminal_slash_commands": ["doctor", "color", "focus", "reload-plugins"], "apiKeySource": "none", "claude_code_version": "2.1.285", "output_style": "default", "agents": ["claude", "Explore", "general-purpose", "Plan", "statusline-setup"], "skills": ["deep-research", "dataviz", "update-config", "verify", "debug", "code-review", "simplify", "batch", "fewer-permission-prompts", "doctor", "loop", "claude-api", "workflow-authoring", "run", "run-skill-generator"], "plugins": [{"name": "cc-plugin-agents-md", "path": "builtin", "source": "cc-plugin-agents-md@builtin"}], "capabilities": ["interrupt_receipt_v1", "interrupt_cancel_queued_v1", "msg_lifecycle_v1", "mcp_read_resource_v1", "mcp_tool_ui_meta_v1"], "analytics_disabled": true, "product_feedback_disabled": true, "uuid": "a5308e12-af9c-41a0-97dc-91fc574b03dc", "fast_mode_state": "off", "fast_mode_disabled_reason": "sdk_opt_in_required", "per_turn_effort_active": false, "view_mode": "default"}
+{"type": "assistant", "message": {"diagnostics": null, "id": "4a8ebe84-f673-472b-9f28-b38722e84b33", "container": null, "model": "<synthetic>", "role": "assistant", "stop_details": null, "stop_reason": "stop_sequence", "stop_sequence": "", "type": "message", "usage": {"output_tokens_details": null, "input_tokens": 0, "output_tokens": 0, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "server_tool_use": {"web_search_requests": 0, "web_fetch_requests": 0}, "service_tier": null, "cache_creation": {"ephemeral_1h_input_tokens": 0, "ephemeral_5m_input_tokens": 0}, "inference_geo": null, "iterations": null, "speed": null, "fallback_credit": null}, "content": [{"type": "text", "text": "API Error: 400 litellm.BadRequestError: You passed in model=does-not-exist-model-xyz. There are no healthy deployments for this model\n\nLiteLLM: model group 'does-not-exist-model-xyz' failed with the error above and no fallback model group was found for it, so the request was not retried on another model. Fallbacks are configured for: anthropic/*, anthropic/claude-opus-4-8, claude-mixed-router, anthropic/claude-fable-5, claude-opus-5, claude-sonnet-5, claude-fable-5, claude-fable-5-1, claude-haiku-4-5-20251001. Add a fallbacks entry for that model group (Router fallbacks or proxy router_settings.fallbacks) to retry on another model."}], "context_management": null}, "parent_tool_use_id": null, "session_id": "53af83ee-c3e1-4b96-a70a-f15b6cb6c794", "uuid": "f8c605c9-c1d6-43b0-89f8-aa64015d7895", "timestamp": "2026-09-30T17:00:14.185Z", "error": "unknown", "is_api_error_message": true}
+{"duration_api_ms": 0, "stop_reason": "stop_sequence", "session_id": "53af83ee-c3e1-4b96-a70a-f15b6cb6c794", "total_cost_usd": 0, "usage": {"output_tokens_details": {"thinking_tokens": 0}, "input_tokens": 0, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0, "output_tokens": 0, "server_tool_use": {"web_search_requests": 0, "web_fetch_requests": 0}, "service_tier": "standard", "cache_creation": {"ephemeral_1h_input_tokens": 0, "ephemeral_5m_input_tokens": 0}, "inference_geo": "", "iterations": [], "speed": "standard", "fallback_credit": null}, "modelUsage": {}, "permission_denials": [], "terminal_reason": "api_error", "fast_mode_state": "off", "fast_mode_disabled_reason": "sdk_opt_in_required", "subagent_stats": {"spawned": 0, "requested": {"background": 0, "foreground": 0, "unset": 0}, "started_in_background": 0, "max_depth": 0, "spawned_by_subagents": 0, "completed": 0, "failed": 0, "killed": {"parent": 0, "user": 0, "system": 0}, "refused": {"depth_limit": 0, "concurrency_limit": 0, "budget": 0}, "by_type": {}}, "is_error": true, "num_turns": 1, "subtype": "success", "api_error_status": 400, "result": "API Error: 400 litellm.BadRequestError: You passed in model=does-not-exist-model-xyz. There are no healthy deployments for this model\n\nLiteLLM: model group 'does-not-exist-model-xyz' failed with the error above and no fallback model group was found for it, so the request was not retried on another model. Fallbacks are configured for: anthropic/*, anthropic/claude-opus-4-8, claude-mixed-router, anthropic/claude-fable-5, claude-opus-5, claude-sonnet-5, claude-fable-5, claude-fable-5-1, claude-haiku-4-5-20251001. Add a fallbacks entry for that model group (Router fallbacks or proxy router_settings.fallbacks) to retry on another model.", "type": "result", "duration_ms": 6781, "uuid": "c38e7f6d-8920-44f4-bab5-a599535509a0", "queued_turn_count": 0, "result_index": 0}