diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index bfaa27c3ef3..542984dd2e0 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.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 ;; diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 445a8519436..eea25e8e285 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.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 diff --git a/README.md b/README.md index 98c5343daee..4004e6474ee 100644 --- a/README.md +++ b/README.md @@ -268,6 +268,31 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse +
+Agents - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK) + +### 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) + +
+ ### 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` | diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py index e4ccbe585a9..bea5e36fd18 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py +++ b/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,)) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py index 9202317c776..5a55b35b255 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py +++ b/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 diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql index 9a061aaed43..a2bec81ca00 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql +++ b/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; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql index 62ad5c42ba7..7eba7fc9b97 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql +++ b/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; diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py new file mode 100644 index 00000000000..a344684a4cd --- /dev/null +++ b/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: 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[^"]+)"\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) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2f74df63367..3acc19d397d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/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: diff --git a/litellm/__init__.py b/litellm/__init__.py index d79086a90f3..9a4f4605519 100644 --- a/litellm/__init__.py +++ b/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 diff --git a/litellm/constants.py b/litellm/constants.py index 76ab419272f..23962400f2a 100644 --- a/litellm/constants.py +++ b/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", + } +) diff --git a/litellm/harness/__init__.py b/litellm/harness/__init__.py new file mode 100644 index 00000000000..79322cdc3e5 --- /dev/null +++ b/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", +) diff --git a/litellm/harness/context.py b/litellm/harness/context.py new file mode 100644 index 00000000000..eaae9aafe1f --- /dev/null +++ b/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 diff --git a/litellm/harness/endpoint.py b/litellm/harness/endpoint.py new file mode 100644 index 00000000000..c492e4794ba --- /dev/null +++ b/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) diff --git a/litellm/harness/errors.py b/litellm/harness/errors.py new file mode 100644 index 00000000000..efa155fdffc --- /dev/null +++ b/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 diff --git a/litellm/harness/handlers/__init__.py b/litellm/harness/handlers/__init__.py new file mode 100644 index 00000000000..4b32d7ca649 --- /dev/null +++ b/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") diff --git a/litellm/harness/handlers/base.py b/litellm/harness/handlers/base.py new file mode 100644 index 00000000000..18b7391e988 --- /dev/null +++ b/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") diff --git a/litellm/harness/handlers/cli_handler.py b/litellm/harness/handlers/cli_handler.py new file mode 100644 index 00000000000..0fda16c34fe --- /dev/null +++ b/litellm/harness/handlers/cli_handler.py @@ -0,0 +1,161 @@ +""" +Generic handler for CLI harnesses (Claude Code, Codex, OpenCode). + +The config (`litellm/llms//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 / 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 diff --git a/litellm/harness/handlers/deepagents_handler.py b/litellm/harness/handlers/deepagents_handler.py new file mode 100644 index 00000000000..3bdc158a0d1 --- /dev/null +++ b/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 diff --git a/litellm/harness/options.py b/litellm/harness/options.py new file mode 100644 index 00000000000..2e3ec5d0fd5 --- /dev/null +++ b/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 diff --git a/litellm/harness/runtime.py b/litellm/harness/runtime.py new file mode 100644 index 00000000000..53a06e7d50d --- /dev/null +++ b/litellm/harness/runtime.py @@ -0,0 +1,1070 @@ +"""The harness engine: validation, sessions, turns, approvals, files, usage and results. + +Adapters only translate a runtime's native protocol into events. Everything that must behave +the same across harnesses (timeouts, max_turns, approvals, FileChange, structured output, +usage and cost) lives here. +""" + +from __future__ import annotations + +import asyncio +import inspect +import logging +import os +import uuid +from collections.abc import AsyncIterator, Callable, Coroutine, Generator, Mapping, Sequence +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import ( + Any, + Final, + get_args, +) + +from pydantic import BaseModel, ValidationError + +import litellm +from litellm.constants import HARNESS_EVENT_QUEUE_MAX_SIZE +from litellm.harness.context import ApprovalHandler, GatewayTarget, SessionContext +from litellm.harness.endpoint import ModelEndpoint +from litellm.harness.errors import ( + CapabilityUnsupported, + HarnessError, + HarnessInstallFailed, + OptionsMismatch, + OutputInvalid, + SessionClosed, + StateIncompatible, +) +from litellm.harness.handlers import get_harness_config, get_harness_handler +from litellm.harness.handlers.base import BaseHarnessHandler +from litellm.harness.options import HarnessOptions +from litellm.harness.sandbox.base import Sandbox +from litellm.harness.sandbox.snapshot import build_file_changes, capture_text_contents +from litellm.harness.types import ( + Approval, + Capabilities, + Done, + Event, + FileChange, + Harness, + PermissionMode, + Result, + State, + StopReason, + Text, + ToolCall, + Usage, + require_harness, +) +from litellm.llms.base_llm.harness.transformation import BaseHarnessConfig +from litellm.llms.base_llm.harness.utils import last_json_object + +PERMISSION_MODES: Final = frozenset(get_args(PermissionMode)) +SKILL_FILE: Final = "SKILL.md" +# Adapter errors that mean "misconfigured", not "the runtime crashed": re-raised to the caller. +verbose_logger: Final = logging.getLogger("LiteLLM") + +PROPAGATED_ERRORS: Final = (HarnessInstallFailed, CapabilityUnsupported) + + +# --------------------------------------------------------------------------- +# Configuration + validation +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class SessionConfig: + """Every per-session parameter a caller can pass, already normalized.""" + + harness: Harness + sandbox: Sandbox + model: str | None = None + gateway: GatewayTarget | None = None + api_key: str | None = None + api_base: str | 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 + install: bool = False + + +LITELLM_PROXY_PREFIX: Final = "litellm_proxy/" + + +def resolve_model_route( + model: str | None, api_key: str | None, api_base: str | None +) -> tuple[str | None, GatewayTarget | None]: + """(model sent to the runtime, gateway or None). + + `litellm_proxy/` (or `litellm.use_litellm_proxy = True`) routes every model call + through the LiteLLM AI Gateway, using api_base/api_key or LITELLM_PROXY_API_BASE / + LITELLM_PROXY_API_KEY. Anything else is called directly through the LiteLLM SDK. + """ + prefixed = model is not None and model.startswith(LITELLM_PROXY_PREFIX) + if not prefixed and not litellm.use_litellm_proxy: + return model, None + group = model[len(LITELLM_PROXY_PREFIX) :] if prefixed and model is not None else model + base = (api_base or os.environ.get("LITELLM_PROXY_API_BASE") or "").strip() + key = (api_key or os.environ.get("LITELLM_PROXY_API_KEY") or "").strip() + if not base: + raise ValueError("litellm_proxy/ models need the gateway URL: pass api_base= or set LITELLM_PROXY_API_BASE") + if not key: + raise ValueError("litellm_proxy/ models need a gateway virtual key: pass api_key= or set LITELLM_PROXY_API_KEY") + return group, GatewayTarget(api_base=base.rstrip("/"), api_key=key) + + +def _normalize_skill(skill: str | os.PathLike[str]) -> str: + path = os.path.abspath(os.fspath(skill)) + if not os.path.isfile(os.path.join(path, SKILL_FILE)): + raise ValueError(f"Skill folder {path!r} has no {SKILL_FILE}") + return path + + +def _normalize_skills(skills: Sequence[str | os.PathLike[str]]) -> tuple[str, ...]: + return tuple(_normalize_skill(skill) for skill in skills) + + +def _check_basic(config: SessionConfig) -> None: + if config.permissions not in PERMISSION_MODES: + raise ValueError(f"permissions must be one of {sorted(PERMISSION_MODES)}, got {config.permissions!r}") + if config.max_turns is not None and config.max_turns < 1: + raise ValueError("max_turns must be >= 1") + if config.timeout is not None and config.timeout <= 0: + raise ValueError("timeout must be > 0") + if config.install: + raise CapabilityUnsupported("install=True is not supported yet; put the runtime binary on PATH in the sandbox") + + +def _check_options(config: SessionConfig, harness_config: BaseHarnessConfig) -> None: + if config.options is None or isinstance(config.options, harness_config.options_type): + return + raise OptionsMismatch( + f"{type(config.options).__name__} cannot be used with Harness.{config.harness.name}; " + f"use {harness_config.options_type.__name__}" + ) + + +def _check_capabilities(config: SessionConfig, caps: Capabilities, interactive: bool) -> None: + name = f"Harness.{config.harness.name}" + if config.permissions not in caps.permission_modes: + raise CapabilityUnsupported( + f"{name} does not support permissions={config.permissions!r}; supported: {sorted(caps.permission_modes)}" + ) + if config.permissions == "ask": + if not caps.tool_approval: + raise CapabilityUnsupported(f"{name} does not support tool approvals") + if config.on_approval is None and not interactive: + raise ValueError("permissions='ask' needs on_approval=, or use stream() and answer Approval events") + if config.output is not None and not caps.structured_output: + raise CapabilityUnsupported(f"{name} does not support output=") + if config.tools and not caps.custom_tools: + raise CapabilityUnsupported(f"{name} does not support custom tools=") + if config.skills and not caps.skills: + raise CapabilityUnsupported(f"{name} does not support skills=") + if config.disable_tools and not caps.tool_filtering: + raise CapabilityUnsupported(f"{name} does not support disable_tools=") + + +def validate(config: SessionConfig, harness_config: BaseHarnessConfig, interactive: bool) -> None: + """Raise before anything starts if the request cannot be served.""" + _check_basic(config) + _check_options(config, harness_config) + _check_capabilities(config, harness_config.capabilities, interactive) + + +def build_config( + 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, +) -> SessionConfig: + """Normalize public keyword arguments into a SessionConfig.""" + resolved_harness = require_harness(harness) + routed_model, gateway = resolve_model_route(model, api_key, api_base) + return SessionConfig( + harness=resolved_harness, + sandbox=sandbox, + model=routed_model, + gateway=gateway, + api_key=api_key, + api_base=api_base, + instructions=instructions, + tools=tuple(tools), + skills=tuple(_normalize_skills(skills)), + disable_tools=tuple(disable_tools), + permissions=permissions, + on_approval=on_approval, + output=output, + max_turns=max_turns, + timeout=timeout, + metadata=MappingProxyType(dict(metadata or ())), + options=options, + install=install, + ) + + +def _context_for(config: SessionConfig) -> SessionContext: + return SessionContext( + harness=config.harness, + sandbox=config.sandbox, + session_id=uuid.uuid4().hex, + model=config.model, + gateway=config.gateway, + api_key=config.api_key, + api_base=config.api_base, + instructions=config.instructions, + tools=config.tools, + skills=config.skills, + disable_tools=config.disable_tools, + permissions=config.permissions, + on_approval=config.on_approval, + output=config.output, + max_turns=config.max_turns, + timeout=config.timeout, + metadata=config.metadata, + options=config.options, + ) + + +# --------------------------------------------------------------------------- +# Structured output +# --------------------------------------------------------------------------- + + +def parse_output(output: type[BaseModel], output_json: str | None, text: str) -> tuple[BaseModel | None, str | None]: + """Return (model, None) on success or (None, error message) on failure.""" + raw = output_json or last_json_object(text) + if raw is None: + return None, "no JSON object found in the final answer" + try: + return output.model_validate_json(raw), None + except ValidationError as e: + return None, str(e) + + +# --------------------------------------------------------------------------- +# Turn machinery +# --------------------------------------------------------------------------- + + +@dataclass +class _End: + """Sentinel the producer puts on the queue when the handler turn is over.""" + + reason: StopReason | None = None + error: BaseException | None = None + + +class TurnControl: + """Lets a stream consumer cancel the running turn.""" + + def __init__(self) -> None: + self.cancelled = False + self.producer: asyncio.Task[None] | None = None + + def cancel(self) -> None: + self.cancelled = True + if self.producer is not None and not self.producer.done(): + self.producer.cancel() + + +async def _aclose(events: AsyncIterator[Event]) -> None: + closer = getattr(events, "aclose", None) + if closer is None: + return + try: + await closer() + except Exception: # closing must not mask the turn's own outcome + verbose_logger.debug("harness: error closing handler turn", exc_info=True) + + +async def pump_events( + events: AsyncIterator[Event], + queue: asyncio.Queue[Event | _End], + max_turns: int | None, +) -> None: + """Drive the handler turn in one task, enforcing max_turns on ToolCall events.""" + end = _End() + tool_calls = 0 + try: + async for event in events: + if isinstance(event, ToolCall): + tool_calls += 1 + if max_turns is not None and tool_calls > max_turns: + end = _End(reason="max_turns") + break + # Backpressure: a runtime that streams faster than the consumer waits here. + await queue.put(event) + except asyncio.CancelledError: + end = _End(reason="cancelled") + raise + except Exception as e: # any runtime failure becomes stop_reason="runtime_error" (see _Turn._finish) + verbose_logger.debug("harness: handler turn raised", exc_info=True) + end = _End(error=e) + finally: + await _aclose(events) + await _put_end(queue, end) + + +async def _put_end(queue: asyncio.Queue[Event | _End], end: _End) -> None: + """Queue the end marker behind every event, waiting for room so no event is dropped. + + A cancelled turn has no consumer left to drain the queue, so only then is space made + by discarding queued events. + """ + if end.reason != "cancelled": + try: + await queue.put(end) + return + except asyncio.CancelledError: + pass + while queue.full(): + queue.get_nowait() + queue.put_nowait(end) + + +async def call_approval_handler(handler: ApprovalHandler, approval: Approval) -> None: + """Run on_approval (sync in a worker thread, or async) and resolve approval.""" + try: + if inspect.iscoroutinefunction(handler): + decision: Any = await handler(approval) + else: + decision = await asyncio.to_thread(handler, approval) + if inspect.isawaitable(decision): + decision = await decision + except Exception as e: # a failing user callback denies the tool instead of crashing the turn + verbose_logger.warning("harness: on_approval raised for tool %s; denying", approval.tool, exc_info=True) + approval.deny(f"on_approval raised: {e}") + return + if decision: + approval.allow() + else: + approval.deny("denied by on_approval") + + +class _Turn: + """One prompt -> events -> Done cycle on a started session.""" + + def __init__( + self, + session: AsyncSession, + prompt: str, + control: TurnControl, + interactive: bool, + ) -> None: + self.session = session + self.ctx = session.ctx + self.prompt = prompt + self.control = control + self.interactive = interactive + self.queue: asyncio.Queue[Event | _End] = asyncio.Queue(maxsize=HARNESS_EVENT_QUEUE_MAX_SIZE) + self.events: list[Event] = [] # mutable-ok: per-turn accumulator the runtime appends events to + self.text_parts: list[str] = [] # mutable-ok: per-turn accumulator of streamed text deltas + self.emitted_files: set[tuple[str, str]] = set() # mutable-ok: per-turn record of emitted FileChanges + self.approval_tasks: list[asyncio.Future[None]] = [] # mutable-ok: per-turn in-flight approval tasks + self.stop_reason: StopReason = "done" + self.error_text: str | None = None + self.before: Mapping[str, str] = MappingProxyType({}) + self.before_contents: Mapping[str, bytes] = MappingProxyType({}) + self.usage_before: tuple[int, int, int, float] = (0, 0, 0, 0.0) + self.deadline: float | None = None + + # -- setup / teardown --------------------------------------------------- + + async def _begin(self) -> None: + sandbox = self.ctx.sandbox + self.before = await sandbox.snapshot() + self.before_contents = await capture_text_contents(sandbox, self.before) + self.usage_before = self.session.usage_counters() + self.ctx.final_text = "" + self.ctx.output_json = None + if self.ctx.timeout is not None: + self.deadline = asyncio.get_running_loop().time() + self.ctx.timeout + + def _start_producer(self) -> None: + events = self.session.handler.turn(self.ctx, self.prompt) + self.control.producer = asyncio.ensure_future(pump_events(events, self.queue, self.ctx.max_turns)) + if self.control.cancelled: + self.control.producer.cancel() + + async def _stop_producer(self) -> None: + producer = self.control.producer + live_producer = (producer,) if producer is not None and not producer.done() else () + pending = (*(task for task in self.approval_tasks if not task.done()), *live_producer) + for task in pending: + task.cancel() + if pending: + await asyncio.wait(pending) + + # -- event loop --------------------------------------------------------- + + async def _next_item(self) -> Event | _End: + if self.deadline is None: + return await self.queue.get() + remaining = self.deadline - asyncio.get_running_loop().time() + try: + if remaining <= 0: + raise asyncio.TimeoutError + return await asyncio.wait_for(self.queue.get(), remaining) + except asyncio.TimeoutError: + await self._stop_producer() + return _End(reason="timeout") + + def _finish(self, end: _End) -> None: + if end.error is not None: + if isinstance(end.error, PROPAGATED_ERRORS): + raise end.error + self.stop_reason = "runtime_error" + self.error_text = f"{type(end.error).__name__}: {end.error}" + verbose_logger.warning("harness %s runtime error: %s", self.ctx.harness.value, self.error_text) + return + if self.control.cancelled: + self.stop_reason = "cancelled" + elif end.reason is not None: + self.stop_reason = end.reason + + async def _on_approval(self, approval: Approval) -> None: + handler = self.ctx.on_approval + if handler is not None: + self.approval_tasks.append(asyncio.ensure_future(call_approval_handler(handler, approval))) + elif not self.interactive: + approval.deny("no approval handler") + + async def _record(self, event: Event) -> None: + if isinstance(event, Text): + self.text_parts.append(event.delta) + elif isinstance(event, FileChange): + self.emitted_files.add((event.path, event.kind)) + elif isinstance(event, Approval): + await self._on_approval(event) + self.events.append(event) + + async def _drain(self) -> AsyncIterator[Event]: + while True: + item = await self._next_item() + if isinstance(item, _End): + self._finish(item) + return + if isinstance(item, Done): + continue + await self._record(item) + yield item + if isinstance(item, Approval) and self.ctx.on_approval is None: + # The consumer asked for the next event without answering. + item.deny("approval not answered") + + # -- results ------------------------------------------------------------ + + async def _file_changes(self) -> list[FileChange]: # mutable-ok: becomes the public Result.files list + sandbox = self.ctx.sandbox + after = await sandbox.snapshot() + files = await build_file_changes(sandbox, self.before, after, self.before_contents) + seen = { # mutable-ok: dedupe set grown while merging streamed FileChange events + change.path for change in files + } + for event in self.events: + if isinstance(event, FileChange) and event.path not in seen: + files.append(event) + seen.add(event.path) + return files + + def _text(self) -> str: + text = self.ctx.final_text or "".join(self.text_parts) + if self.error_text is None: + return text + return f"{text}\n\n{self.error_text}" if text else self.error_text + + def _usage(self) -> tuple[Usage, float]: + now = self.session.usage_counters() + before = self.usage_before + usage = Usage( + input_tokens=now[0] - before[0], + output_tokens=now[1] - before[1], + calls=now[2] - before[2], + ) + return usage, max(now[3] - before[3], 0.0) + + def _result( + self, + files: list[FileChange], # mutable-ok: Result.files is a public list field + output: BaseModel | None, + ) -> Result: + usage, cost = self._usage() + return Result( + text=self._text(), + output=output, + files=files, + events=list( # mutable-ok: Result.events is a public list field; copy detaches it from the accumulator + self.events + ), + usage=usage, + cost=cost, + stop_reason=self.stop_reason, + session_id=self.ctx.session_id, + ) + + def _output(self) -> tuple[BaseModel | None, str | None, str | None]: + """(parsed output, raw text, error) for the structured-output check.""" + output_type = self.ctx.output + if output_type is None or self.stop_reason != "done": + return None, None, None + text = self._text() + parsed, error = parse_output(output_type, self.ctx.output_json, text) + return parsed, self.ctx.output_json or text, error + + # -- entry -------------------------------------------------------------- + + async def run(self) -> AsyncIterator[Event]: + await self._begin() + self._start_producer() + try: + async for event in self._drain(): + yield event + finally: + await self._stop_producer() + if self.stop_reason != "done": + await self.session.interrupt() + files = await self._file_changes() + for change in files: + if (change.path, change.kind) not in self.emitted_files: + self.events.append(change) + yield change + parsed, raw, error = self._output() + result = self._result(files, parsed) + self.session.record(result) + yield Done(result) + if error is not None: + raise OutputInvalid( + f"Final answer did not match {self.ctx.output.__name__ if self.ctx.output else 'output'}: {error}", + raw=raw or "", + result=result, + ) + + +# --------------------------------------------------------------------------- +# Streams +# --------------------------------------------------------------------------- + + +class AsyncEventStream: + """Async iterator of events for one turn. `.result` is set once Done is seen.""" + + def __init__(self, source: AsyncIterator[Event], control: TurnControl) -> None: + self._source = source + self._control = control + self._result: Result | None = None + + def __aiter__(self) -> AsyncEventStream: + return self + + async def __anext__(self) -> Event: + event = await self._source.__anext__() + if isinstance(event, Done): + self._result = event.result + return event + + @property + def result(self) -> Result | None: + return self._result + + def cancel(self) -> None: + """Stop the turn. The stream still ends with Done(stop_reason='cancelled').""" + self._control.cancel() + + async def aclose(self) -> None: + await _aclose(self._source) + + +async def _one_shot(session: AsyncSession, prompt: str, control: TurnControl) -> AsyncIterator[Event]: + """Stream one turn on a fresh session and close it before Done is handed out.""" + try: + async for event in session.turn_events(prompt, control, interactive=True): + if isinstance(event, Done): + await session.aclose() + yield event + finally: + await session.aclose() + + +# --------------------------------------------------------------------------- +# Sessions +# --------------------------------------------------------------------------- + + +class AsyncSession: + """A multi-turn conversation with one harness. Use `async with` or `await`.""" + + def __init__( + self, + config: SessionConfig, + *, + resume_from: str | None = None, + interactive: bool = True, + ) -> None: + self.config = config + self.harness_config = get_harness_config(config.harness) + validate(config, self.harness_config, interactive=interactive) + if resume_from is not None and not self.harness_config.capabilities.resume: + raise CapabilityUnsupported(f"Harness.{config.harness.name} does not support resume") + self.ctx = _context_for(config) + # Config-specific static checks (managed option keys, required model) before any I/O. + self.harness_config.validate_environment(self.ctx) + self.handler: BaseHarnessHandler = get_harness_handler(self.harness_config) + self.results: list[Result] = [] # mutable-ok: session accumulator; each turn's Result is appended + self._resume_from = resume_from + self._native_id: str | None = resume_from + self._started = False + self._closed = False + self._busy = False + self._restart_needed = False + + # -- lifecycle ---------------------------------------------------------- + + def __await__(self) -> Generator[object, None, AsyncSession]: + return self.start().__await__() + + async def __aenter__(self) -> AsyncSession: + return await self.start() + + async def __aexit__(self, *exc_info: object) -> None: + await self.aclose() + + async def _open_endpoint(self) -> None: + if not self.harness_config.uses_model_endpoint or self.ctx.endpoint is not None: + return + endpoint = ModelEndpoint( + self.config.harness, + self.config.model, + self.config.gateway, + api_key=self.config.api_key, + api_base=self.config.api_base, + metadata=self.config.metadata, + ) + await endpoint.__aenter__() + self.ctx.endpoint = endpoint + + async def _launch(self) -> None: + await self.handler.start(self.ctx) + if self._native_id is not None and (self._resume_from is not None or self._restart_needed): + await self.handler.resume(self.ctx, self._native_id) + + async def start(self) -> AsyncSession: + if self._closed: + raise SessionClosed("session is closed") + if self._started: + return self + await self._open_endpoint() + try: + await self._launch() + except BaseException: + await self._close_endpoint() + raise + self._started = True + return self + + async def interrupt(self) -> None: + """Stop the runtime after a timeout / max_turns / cancel; next turn restarts it.""" + self._native_id = self.handler.native_session_id() or self._native_id + try: + await self.handler.stop(self.ctx) + except Exception: # the next turn restarts the runtime regardless + verbose_logger.warning("harness: handler stop failed", exc_info=True) + self._restart_needed = True + + async def _ensure_ready(self) -> None: + if self._closed: + raise SessionClosed("session is closed") + if not self._started: + await self.start() + elif self._restart_needed: + await self._launch() + self._restart_needed = False + + async def _close_endpoint(self) -> None: + endpoint = self.ctx.endpoint + self.ctx.endpoint = None + if endpoint is None: + return + try: + await endpoint.__aexit__(None, None, None) + except Exception: # shutdown is best-effort cleanup + verbose_logger.warning("harness: endpoint shutdown failed", exc_info=True) + + async def aclose(self) -> None: + """Stop the runtime and the endpoint. Safe to call twice.""" + if self._closed: + return + self._closed = True + if self._started: + self._native_id = self.handler.native_session_id() or self._native_id + try: + await self.handler.stop(self.ctx) + except Exception: # still close the endpoint below + verbose_logger.warning("harness: handler stop failed", exc_info=True) + await self._close_endpoint() + + close = aclose + + def state(self) -> State: + native = self._native_id + if self._started and not self._closed: + native = self.handler.native_session_id() or native + return State( + harness=self.config.harness, + native_session_id=native, + workdir=self.config.sandbox.workdir, + model=self.config.model, + ) + + async def adetach(self) -> State: + """Release local resources and return State to resume() later.""" + await self.aclose() + return self.state() + + async def astop(self) -> State: + """Stop the session for good and return its final State.""" + await self.aclose() + return self.state() + + detach = adetach + stop = astop + + # -- turns -------------------------------------------------------------- + + def usage_counters(self) -> tuple[int, int, int, float]: + """(input_tokens, output_tokens, calls, cost) so far, from endpoint or handler.""" + endpoint = self.ctx.endpoint + if endpoint is not None: + usage = endpoint.usage + return (usage.input_tokens, usage.output_tokens, usage.calls, usage.cost) + ctx = self.ctx + return (ctx.input_tokens, ctx.output_tokens, ctx.calls, ctx.cost) + + def record(self, result: Result) -> None: + self.results.append(result) + + async def turn_events(self, prompt: str, control: TurnControl, interactive: bool) -> AsyncIterator[Event]: + if self._busy: + raise HarnessError("a turn is already running on this session") + self._busy = True + try: + await self._ensure_ready() + async for event in _Turn(self, prompt, control, interactive).run(): + yield event + finally: + self._busy = False + + def astream(self, prompt: str) -> AsyncEventStream: + control = TurnControl() + return AsyncEventStream(self.turn_events(prompt, control, interactive=True), control) + + async def arun(self, prompt: str) -> Result: + return await _collect(self.turn_events(prompt, TurnControl(), False)) + + async def history( + self, + ) -> list[dict[str, Any]]: # mutable-ok: public API returns OpenAI-format message dicts from the handler + if not self.harness_config.capabilities.history: + raise CapabilityUnsupported(f"Harness.{self.config.harness.name} does not expose history") + await self._ensure_ready() + return await self.handler.history(self.ctx) + + @property + def cost(self) -> float: + return sum(result.cost for result in self.results) + + @property + def usage(self) -> Usage: + return Usage( + input_tokens=sum(r.usage.input_tokens for r in self.results), + output_tokens=sum(r.usage.output_tokens for r in self.results), + calls=sum(r.usage.calls for r in self.results), + ) + + @property + def session_id(self) -> str: + return self.ctx.session_id + + @property + def closed(self) -> bool: + return self._closed + + +async def _collect(events: AsyncIterator[Event]) -> Result: + result: Result | None = None + async for event in events: + if isinstance(event, Done): + result = event.result + if result is None: + raise HarnessError("turn ended without a result") + return result + + +# --------------------------------------------------------------------------- +# Public async API +# --------------------------------------------------------------------------- + + +def aagent_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, +) -> AsyncSession: + """A multi-turn agent session: `async with litellm.aagent_session(...) as s:`.""" + config = build_config( + 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, + ) + return AsyncSession(config) + + +async def arun_agent( + 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 and return the Result.""" + config = build_config( + 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, + ) + async with AsyncSession(config, interactive=False) as session: + return await session.arun(prompt) + + +def astream_agent( + 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, +) -> AsyncEventStream: + """Stream events for one prompt. Validation errors raise here, before iteration.""" + 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, + ) + control = TurnControl() + return AsyncEventStream(_one_shot(session, prompt, control), control) + + +def _coerce_state(state: State | bytes) -> State: + if isinstance(state, (bytes, bytearray)): + return State.loads(bytes(state)) + if not isinstance(state, State): + raise TypeError(f"state must be a State or bytes, got {type(state).__name__}") + return state + + +def aagent_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, +) -> AsyncSession: + """Continue a detached/stopped session from its State.""" + resolved = _coerce_state(state) + if not resolved.native_session_id: + raise StateIncompatible("State has no native session id to resume") + config = build_config( + resolved.harness, + sandbox=sandbox, + model=model if model is not None else resolved.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 AsyncSession(config, resume_from=resolved.native_session_id) + + +def agent_capabilities(harness: Harness) -> Capabilities: + """What a harness supports (permission modes, structured output, tools...).""" + return get_harness_config(require_harness(harness)).capabilities + + +def aagent( + 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, +) -> Coroutine[Any, Any, Result] | AsyncEventStream: + """Run an agent harness on one prompt. + + `await litellm.aagent(...)` returns a Result. With stream=True it returns an async + iterator of events instead: `async for event in litellm.aagent(..., stream=True)`. + """ + kwargs: dict[str, Any] = { # mutable-ok: forwarded as **kwargs to arun_agent/astream_agent + "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 astream_agent(harness, prompt, **kwargs) + return arun_agent(harness, prompt, **kwargs) diff --git a/litellm/harness/sandbox/__init__.py b/litellm/harness/sandbox/__init__.py new file mode 100644 index 00000000000..434917d625a --- /dev/null +++ b/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", +) diff --git a/litellm/harness/sandbox/base.py b/litellm/harness/sandbox/base.py new file mode 100644 index 00000000000..18bdeac58a1 --- /dev/null +++ b/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: ... diff --git a/litellm/harness/sandbox/docker.py b/litellm/harness/sandbox/docker.py new file mode 100644 index 00000000000..32dbfc27ddc --- /dev/null +++ b/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 (" ./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 ` 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 ` 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) diff --git a/litellm/harness/sandbox/local.py b/litellm/harness/sandbox/local.py new file mode 100644 index 00000000000..10c7303ef27 --- /dev/null +++ b/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) diff --git a/litellm/harness/sandbox/snapshot.py b/litellm/harness/sandbox/snapshot.py new file mode 100644 index 00000000000..61407c6c063 --- /dev/null +++ b/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) + ] diff --git a/litellm/harness/sync.py b/litellm/harness/sync.py new file mode 100644 index 00000000000..4583a98bf54 --- /dev/null +++ b/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) diff --git a/litellm/harness/types.py b/litellm/harness/types.py new file mode 100644 index 00000000000..88ad7f711ea --- /dev/null +++ b/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, + ) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ae4ac29de9e..0c48b7c1c2a 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/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 ) diff --git a/litellm/llms/base_llm/harness/__init__.py b/litellm/llms/base_llm/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/base_llm/harness/transformation.py b/litellm/llms/base_llm/harness/transformation.py new file mode 100644 index 00000000000..643f808d4d8 --- /dev/null +++ b/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/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.""" diff --git a/litellm/llms/base_llm/harness/utils.py b/litellm/llms/base_llm/harness/utils.py new file mode 100644 index 00000000000..7c0278d5e36 --- /dev/null +++ b/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)) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index c9fec2db1d7..48a8b1b44bb 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/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) diff --git a/litellm/llms/claude_code/__init__.py b/litellm/llms/claude_code/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/claude_code/harness/__init__.py b/litellm/llms/claude_code/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/claude_code/harness/transformation.py b/litellm/llms/claude_code/harness/transformation.py new file mode 100644 index 00000000000..a85968f78be --- /dev/null +++ b/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 = "" + +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) diff --git a/litellm/llms/codex/__init__.py b/litellm/llms/codex/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/codex/harness/__init__.py b/litellm/llms/codex/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/codex/harness/transformation.py b/litellm/llms/codex/harness/transformation.py new file mode 100644 index 00000000000..ab17cf3d869 --- /dev/null +++ b/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) diff --git a/litellm/llms/deepagents/__init__.py b/litellm/llms/deepagents/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/deepagents/harness/__init__.py b/litellm/llms/deepagents/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/deepagents/harness/sandbox_backend.py b/litellm/llms/deepagents/harness/sandbox_backend.py new file mode 100644 index 00000000000..49610014e21 --- /dev/null +++ b/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 +`/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/" or "f/" 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) diff --git a/litellm/llms/deepagents/harness/transformation.py b/litellm/llms/deepagents/harness/transformation.py new file mode 100644 index 00000000000..339a6274d33 --- /dev/null +++ b/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/` 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=") diff --git a/litellm/llms/opencode/__init__.py b/litellm/llms/opencode/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/opencode/harness/__init__.py b/litellm/llms/opencode/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/opencode/harness/transformation.py b/litellm/llms/opencode/harness/transformation.py new file mode 100644 index 00000000000..af5fa1ae71d --- /dev/null +++ b/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:/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) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 2c63d957e75..9c7778e2b77 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/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, diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 44523f5cf45..d17b849750b 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/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) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 44979d12a00..2090a0c7421 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/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 diff --git a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py b/litellm/proxy/_experimental/mcp_server/ui_session_utils.py index 901259c18ad..b9e25259868 100644 --- a/litellm/proxy/_experimental/mcp_server/ui_session_utils.py +++ b/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)) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 0b687340ea5..0b470cf7bda 100644 --- a/litellm/proxy/_lazy_features.py +++ b/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} diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index afaa41e97a5..e7c7102c98f 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/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, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 0974a2a952d..8ed8ec8752e 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/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( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 86b37f31090..68f635f7799 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/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, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 27c03d2d5d7..9eb2a4444d6 100644 --- a/litellm/proxy/proxy_cli.py +++ b/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 diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 71f61079154..c9e0861935d 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/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. diff --git a/litellm/sandbox/__init__.py b/litellm/sandbox/__init__.py index e69de29bb2d..1974859fb1e 100644 --- a/litellm/sandbox/__init__.py +++ b/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}") diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 858123b5232..47f80c52845 100644 --- a/litellm/types/llms/custom_http.py +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index 0eb2754ed0c..05d5986885c 100644 --- a/litellm/utils.py +++ b/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, diff --git a/migrations/run.py b/migrations/run.py index 7ea80d48719..94e7cc7c0f0 100644 --- a/migrations/run.py +++ b/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, ) diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py index d7da48ce933..5c694702c6e 100644 --- a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py +++ b/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: ` 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: 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): diff --git a/tests/code_coverage_tests/check_provider_folders_documented.py b/tests/code_coverage_tests/check_provider_folders_documented.py index 60afc55331f..08fbde3d979 100644 --- a/tests/code_coverage_tests/check_provider_folders_documented.py +++ b/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", } diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 659dc438f2d..863934d76f7 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/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). diff --git a/tests/harness_e2e/__init__.py b/tests/harness_e2e/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/harness_e2e/conftest.py b/tests/harness_e2e/conftest.py new file mode 100644 index 00000000000..ba679d6d5bf --- /dev/null +++ b/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_. +""" + +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/`: 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 diff --git a/tests/harness_e2e/test_harness_e2e.py b/tests/harness_e2e/test_harness_e2e.py new file mode 100644 index 00000000000..8892085f4bc --- /dev/null +++ b/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"], + ) diff --git a/tests/integration/_support/mcp_grants.py b/tests/integration/_support/mcp_grants.py index 5fa9eeaa0b6..82e79b2c3bd 100644 --- a/tests/integration/_support/mcp_grants.py +++ b/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], }, ) diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index b0941672be1..8cfdf0db2c3 100644 --- a/tests/integration/_support/process.py +++ b/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) diff --git a/tests/integration/configuration/test_lazy_routes_flag.py b/tests/integration/configuration/test_lazy_routes_flag.py new file mode 100644 index 00000000000..6135deca1c8 --- /dev/null +++ b/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 diff --git a/tests/integration/database/test_request_log_indexes_at_boot.py b/tests/integration/database/test_request_log_indexes_at_boot.py new file mode 100644 index 00000000000..5f95c5c88df --- /dev/null +++ b/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) diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index 26f5ced4de7..01bf006a03e 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/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 diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 7a60c8ede30..60563e7aacd 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/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()) == () diff --git a/tests/integration/mcp/test_mcp_toolsets.py b/tests/integration/mcp/test_mcp_toolsets.py new file mode 100644 index 00000000000..3dd309db665 --- /dev/null +++ b/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()) == () diff --git a/tests/integration/providers/test_bedrock_converse_missing_content_wire.py b/tests/integration/providers/test_bedrock_converse_missing_content_wire.py new file mode 100644 index 00000000000..057596941a2 --- /dev/null +++ b/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: ]"}], +} +_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 diff --git a/tests/integration/spend/test_key_metadata_recovery_probe_bounds.py b/tests/integration/spend/test_key_metadata_recovery_probe_bounds.py index 81be9e56f30..fd67721cc6b 100644 --- a/tests/integration/spend/test_key_metadata_recovery_probe_bounds.py +++ b/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]: diff --git a/tests/proxy_migration_tests/test_db_schema_migration.py b/tests/proxy_migration_tests/test_db_schema_migration.py index b0d44cd3e1c..70498a3bcbd 100644 --- a/tests/proxy_migration_tests/test_db_schema_migration.py +++ b/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 `.\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 `.\n\n" + + drift + ) + else: + assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}" finally: shutil.rmtree(temp_base, ignore_errors=True) diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py index 741fa7386df..0c971b5d073 100644 --- a/tests/proxy_migration_tests/test_invalid_index_repair.py +++ b/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: diff --git a/tests/proxy_migration_tests/test_request_log_indexes.py b/tests/proxy_migration_tests/test_request_log_indexes.py new file mode 100644 index 00000000000..3b0c87b2a97 --- /dev/null +++ b/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 diff --git a/tests/unit/harness/__init__.py b/tests/unit/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/harness/core_fakes.py b/tests/unit/harness/core_fakes.py new file mode 100644 index 00000000000..c86cf39739b --- /dev/null +++ b/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() diff --git a/tests/unit/harness/handlers/__init__.py b/tests/unit/harness/handlers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/harness/handlers/test_deepagents_handler.py b/tests/unit/harness/handlers/test_deepagents_handler.py new file mode 100644 index 00000000000..6f0963cac46 --- /dev/null +++ b/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 diff --git a/tests/unit/harness/sandbox/__init__.py b/tests/unit/harness/sandbox/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/harness/sandbox/test_docker.py b/tests/unit/harness/sandbox/test_docker.py new file mode 100644 index 00000000000..10cb8a092f3 --- /dev/null +++ b/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() diff --git a/tests/unit/harness/sandbox/test_local.py b/tests/unit/harness/sandbox/test_local.py new file mode 100644 index 00000000000..5d08dac13a1 --- /dev/null +++ b/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',''));" + "print(os.environ.get('OPENAI_BASE_URL',''));" + "print(os.environ.get('GITHUB_TOKEN',''));" + "print(os.environ.get('HARNESS_TEST_PLAIN',''));" + "print(os.environ.get('ANTHROPIC_BASE_URL',''))" + ) + result = await sbx.run( + [PY, "-c", script], env={"ANTHROPIC_BASE_URL": "http://127.0.0.1:1"} + ) + assert result.stdout.split() == [ + "", + "", + "", + "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"] diff --git a/tests/unit/harness/sandbox/test_snapshot.py b/tests/unit/harness/sandbox/test_snapshot.py new file mode 100644 index 00000000000..e3bfc8806db --- /dev/null +++ b/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() diff --git a/tests/unit/harness/test_endpoint.py b/tests/unit/harness/test_endpoint.py new file mode 100644 index 00000000000..8b56debdc29 --- /dev/null +++ b/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 diff --git a/tests/unit/harness/test_init.py b/tests/unit/harness/test_init.py new file mode 100644 index 00000000000..effdc238249 --- /dev/null +++ b/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 diff --git a/tests/unit/harness/test_runtime.py b/tests/unit/harness/test_runtime.py new file mode 100644 index 00000000000..e5f2235b794 --- /dev/null +++ b/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" diff --git a/tests/unit/harness/test_sync.py b/tests/unit/harness/test_sync.py new file mode 100644 index 00000000000..a92f2957aee --- /dev/null +++ b/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)) diff --git a/tests/unit/harness/test_types.py b/tests/unit/harness/test_types.py new file mode 100644 index 00000000000..036d6dd625c --- /dev/null +++ b/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") diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 1e12a973cdb..8d2e6b9fd0c 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/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", + ) + == [] + ) diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index 755c7617701..5540cf54193 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/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 diff --git a/tests/unit/litellm_proxy_extras/test_request_log_indexes.py b/tests/unit/litellm_proxy_extras/test_request_log_indexes.py new file mode 100644 index 00000000000..5cf7593c5cd --- /dev/null +++ b/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: __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("") == "" diff --git a/tests/unit/llms/base_llm/harness/__init__.py b/tests/unit/llms/base_llm/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 499096621c5..f6f98e3b9bd 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/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", diff --git a/tests/unit/llms/claude_code/__init__.py b/tests/unit/llms/claude_code/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/claude_code/harness/__init__.py b/tests/unit/llms/claude_code/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/claude_code/harness/fixtures/__init__.py b/tests/unit/llms/claude_code/harness/fixtures/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/claude_code/harness/fixtures/api_error.jsonl b/tests/unit/llms/claude_code/harness/fixtures/api_error.jsonl new file mode 100644 index 00000000000..649b44345ce --- /dev/null +++ b/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": "", "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} diff --git a/tests/unit/llms/claude_code/harness/fixtures/max_turns.jsonl b/tests/unit/llms/claude_code/harness/fixtures/max_turns.jsonl new file mode 100644 index 00000000000..f40c7eae8a2 --- /dev/null +++ b/tests/unit/llms/claude_code/harness/fixtures/max_turns.jsonl @@ -0,0 +1,11 @@ +{"type": "system", "subtype": "init", "cwd": "/workspace", "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "tools": ["Task", "Bash", "CronCreate", "CronDelete", "CronList", "Edit", "EnterWorktree", "ExitWorktree", "ListAgents", "NotebookEdit", "Read", "ReportFindings", "ScheduleWakeup", "SendMessage", "Skill", "TaskCreate", "TaskGet", "TaskList", "TaskStop", "TaskUpdate", "WebFetch", "WebSearch", "Workflow", "Write"], "mcp_servers": [], "model": "claude-haiku-4-5-20251001", "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": "a1474c31-391d-443a-8cb7-ff594f9bf72c", "fast_mode_state": "off", "fast_mode_disabled_reason": "sdk_opt_in_required", "per_turn_effort_active": false, "view_mode": "default"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 7, "estimated_tokens_delta": 7, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "b263f250-510d-44b0-a0aa-79d269fe38cd"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 11, "estimated_tokens_delta": 4, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "a84f3325-922e-4ca1-a0f2-acb8d5b60b91"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZz9LnyzBujxEY6yNoWA", "type": "message", "role": "assistant", "content": [{"type": "thinking", "thinking": "The user wants me to create three files:\n1. a.txt with content \"1\"\n2. b.txt with content \"2\"\n3. c.txt with content \"3\"\n\nThey want one tool call for each, which means I should make three separate Write tool calls.\n\nSince these are independent operations (no dependencies between them), I should make all three calls in parallel in a single function_calls block.\n\nThe files should be created in the primary working directory: /workspace"}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 2730, "cache_read_input_tokens": 26136, "cache_creation": {"ephemeral_5m_input_tokens": 2730, "ephemeral_1h_input_tokens": 0}, "output_tokens": 6, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "1a577616-8652-49e8-8eba-195b5ea6d7d2", "timestamp": "2026-09-30T17:03:49.895Z"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZz9LnyzBujxEY6yNoWA", "type": "message", "role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01AAviEQjsijaCfLqz9HJexD", "name": "Write", "input": {"file_path": "/workspace/a.txt", "content": "1"}, "caller": {"type": "direct"}}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 2730, "cache_read_input_tokens": 26136, "cache_creation": {"ephemeral_5m_input_tokens": 2730, "ephemeral_1h_input_tokens": 0}, "output_tokens": 6, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "aec1ad61-fdb1-4e47-9238-68fad55c06b4", "timestamp": "2026-09-30T17:03:50.170Z"} +{"type": "user", "message": {"role": "user", "content": [{"tool_use_id": "toolu_01AAviEQjsijaCfLqz9HJexD", "type": "tool_result", "content": "File created successfully at: /workspace/a.txt (file state is current in your context — no need to Read it back)"}]}, "parent_tool_use_id": null, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "cfa2046c-604b-45ea-bcc4-a605207e77b8", "timestamp": "2026-09-30T17:03:50.177Z", "tool_use_result": {"type": "create", "filePath": "/workspace/a.txt", "content": "1", "structuredPatch": [], "originalFile": null, "userModified": false}} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZz9LnyzBujxEY6yNoWA", "type": "message", "role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01E8ZcsQTKVnqC76NZzV3HCU", "name": "Write", "input": {"file_path": "/workspace/b.txt", "content": "2"}, "caller": {"type": "direct"}}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 2730, "cache_read_input_tokens": 26136, "cache_creation": {"ephemeral_5m_input_tokens": 2730, "ephemeral_1h_input_tokens": 0}, "output_tokens": 6, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "40e1dcdd-3c69-4318-9081-81c68ae13ae3", "timestamp": "2026-09-30T17:03:50.450Z"} +{"type": "user", "message": {"role": "user", "content": [{"tool_use_id": "toolu_01E8ZcsQTKVnqC76NZzV3HCU", "type": "tool_result", "content": "File created successfully at: /workspace/b.txt (file state is current in your context — no need to Read it back)"}]}, "parent_tool_use_id": null, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "a2b1dfa4-cde9-4bfb-8f2a-460a0eef2a18", "timestamp": "2026-09-30T17:03:50.456Z", "tool_use_result": {"type": "create", "filePath": "/workspace/b.txt", "content": "2", "structuredPatch": [], "originalFile": null, "userModified": false}} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZz9LnyzBujxEY6yNoWA", "type": "message", "role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01N9f8PbyFhmyZ9wgiTqg3uG", "name": "Write", "input": {"file_path": "/workspace/c.txt", "content": "3"}, "caller": {"type": "direct"}}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 2730, "cache_read_input_tokens": 26136, "cache_creation": {"ephemeral_5m_input_tokens": 2730, "ephemeral_1h_input_tokens": 0}, "output_tokens": 6, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "8222d589-7c3e-4233-8bd2-4a73ffabeed2", "timestamp": "2026-09-30T17:03:50.725Z"} +{"type": "user", "message": {"role": "user", "content": [{"tool_use_id": "toolu_01N9f8PbyFhmyZ9wgiTqg3uG", "type": "tool_result", "content": "File created successfully at: /workspace/c.txt (file state is current in your context — no need to Read it back)"}]}, "parent_tool_use_id": null, "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "uuid": "98d09b2c-80b3-4b4f-8933-cfa7225ba8dc", "timestamp": "2026-09-30T17:03:50.737Z", "tool_use_result": {"type": "create", "filePath": "/workspace/c.txt", "content": "3", "structuredPatch": [], "originalFile": null, "userModified": false}} +{"duration_api_ms": 3347, "stop_reason": "tool_use", "session_id": "2270f364-6991-4263-8761-cd601a27cb8a", "total_cost_usd": 0.0077101, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 2730, "cache_read_input_tokens": 26136, "output_tokens": 335, "output_tokens_details": {"thinking_tokens": 118}, "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": 2730}, "inference_geo": "not_available", "iterations": [], "speed": "standard", "fallback_credit": null}, "modelUsage": {"claude-haiku-4-5-20251001": {"inputTokens": 9, "outputTokens": 335, "cacheReadInputTokens": 26136, "cacheCreationInputTokens": 2730, "webSearchRequests": 0, "costUSD": 0.0077101, "contextWindow": 200000, "maxOutputTokens": 32000, "thinkingTokens": 118, "canonicalModel": "claude-haiku-4-5", "provider": "firstParty", "costBasis": "list"}}, "permission_denials": [], "terminal_reason": "max_turns", "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": 2, "subtype": "error_max_turns", "errors": ["Reached maximum number of turns (1)"], "type": "result", "duration_ms": 3381, "uuid": "db4afec5-5854-4439-89b4-f8d3539de3fd", "queued_turn_count": 0, "result_index": 0} diff --git a/tests/unit/llms/claude_code/harness/fixtures/resume_turn.jsonl b/tests/unit/llms/claude_code/harness/fixtures/resume_turn.jsonl new file mode 100644 index 00000000000..3b6994ff6a5 --- /dev/null +++ b/tests/unit/llms/claude_code/harness/fixtures/resume_turn.jsonl @@ -0,0 +1,6 @@ +{"type": "system", "subtype": "init", "cwd": "/workspace", "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "tools": ["Task", "Bash", "CronCreate", "CronDelete", "CronList", "Edit", "EnterWorktree", "ExitWorktree", "ListAgents", "NotebookEdit", "Read", "ReportFindings", "ScheduleWakeup", "SendMessage", "Skill", "TaskCreate", "TaskGet", "TaskList", "TaskStop", "TaskUpdate", "WebFetch", "WebSearch", "Workflow", "Write"], "mcp_servers": [], "model": "claude-haiku-4-5-20251001", "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": "753b841b-38d0-44e1-97a0-1cf3fc4875b0", "fast_mode_state": "off", "fast_mode_disabled_reason": "sdk_opt_in_required", "per_turn_effort_active": false, "view_mode": "default"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 7, "estimated_tokens_delta": 7, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "8c86350f-f3c9-44d8-98b1-b9bc64a5d89d"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 10, "estimated_tokens_delta": 3, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "533fe314-e617-4692-bc30-3985fd3cadc1"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZyrEoMRLGNJffWtEf6r", "type": "message", "role": "assistant", "content": [{"type": "thinking", "thinking": "The user is asking what file I created, and they want a one-word answer. I created hello.txt, so the one-word answer would be \"hello.txt\" but that's technically two parts with a dot. If they want literally one word, it could be \"hello\" without the extension. But \"hello.txt\" is the filename. Let me think about what makes sense - they said \"one word\" so they probably want just \"hello\" or \"hello.txt\". The filename is \"hello.txt\" but if forced to one word, \"hello\" works.\n\nActually, looking at it more carefully, they asked \"What file\" - the file is \"hello.txt\". If they want one word, the answer should be \"hello.txt\" as a single unit (no spaces), or just \"hello\" if being very literal about \"one word\". I'll go with the filename as written."}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 104, "cache_read_input_tokens": 29266, "cache_creation": {"ephemeral_5m_input_tokens": 104, "ephemeral_1h_input_tokens": 0}, "output_tokens": 6, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "dfb24f01-69d2-4a2c-8628-d78af662deb2", "timestamp": "2026-09-30T16:59:58.607Z"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZyrEoMRLGNJffWtEf6r", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "hello.txt"}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 104, "cache_read_input_tokens": 29266, "cache_creation": {"ephemeral_5m_input_tokens": 104, "ephemeral_1h_input_tokens": 0}, "output_tokens": 6, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "09aae2d0-4408-4a2f-b6d1-238d15920f49", "timestamp": "2026-09-30T16:59:58.629Z"} +{"duration_api_ms": 27333, "stop_reason": "end_turn", "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "total_cost_usd": 0.048190000000000004, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 104, "cache_read_input_tokens": 29266, "output_tokens": 206, "output_tokens_details": {"thinking_tokens": 197}, "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": 104}, "inference_geo": "not_available", "iterations": [], "speed": "standard", "fallback_credit": null}, "modelUsage": {"claude-haiku-4-5-20251001": {"inputTokens": 34, "outputTokens": 544, "thinkingTokens": 354, "cacheReadInputTokens": 87235, "cacheCreationInputTokens": 29370, "webSearchRequests": 0, "costUSD": 0.048190000000000004, "contextWindow": 200000, "maxOutputTokens": 32000, "canonicalModel": "claude-haiku-4-5", "provider": "firstParty", "costBasis": "list"}}, "permission_denials": [], "terminal_reason": "completed", "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": false, "num_turns": 1, "subtype": "success", "api_error_status": null, "result": "hello.txt", "ttft_ms": 7545, "type": "result", "duration_ms": 7578, "uuid": "701ed1c8-48e5-4aad-b3fe-3982ad7aed00", "ttft_stream_ms": 5547, "time_to_request_ms": 23, "first_content_frame_ms": 5547, "queued_turn_count": 0, "result_index": 0} diff --git a/tests/unit/llms/claude_code/harness/fixtures/structured_output.jsonl b/tests/unit/llms/claude_code/harness/fixtures/structured_output.jsonl new file mode 100644 index 00000000000..29ba78c80e4 --- /dev/null +++ b/tests/unit/llms/claude_code/harness/fixtures/structured_output.jsonl @@ -0,0 +1,7 @@ +{"type": "system", "subtype": "init", "cwd": "/workspace", "session_id": "e0b4fb7e-b899-44ac-81fd-62841efa5380", "tools": ["Task", "Bash", "CronCreate", "CronDelete", "CronList", "Edit", "EnterWorktree", "ExitWorktree", "ListAgents", "NotebookEdit", "Read", "ReportFindings", "ScheduleWakeup", "SendMessage", "Skill", "StructuredOutput", "TaskCreate", "TaskGet", "TaskList", "TaskStop", "TaskUpdate", "WebFetch", "WebSearch", "Workflow", "Write"], "mcp_servers": [], "model": "claude-haiku-4-5-20251001", "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": "a9fd450b-ce6f-4dc8-8d2c-1e2c1b654aa5", "fast_mode_state": "off", "fast_mode_disabled_reason": "sdk_opt_in_required", "per_turn_effort_active": false, "view_mode": "default"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 5, "estimated_tokens_delta": 5, "session_id": "e0b4fb7e-b899-44ac-81fd-62841efa5380", "uuid": "c52043cf-0474-42ce-ae1b-1275fda7ee57"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 9, "estimated_tokens_delta": 4, "session_id": "e0b4fb7e-b899-44ac-81fd-62841efa5380", "uuid": "dc90c9e0-4aef-4fe0-8dc9-25738677206c"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZzH4cNfTXjBxTRhX5Vm", "type": "message", "role": "assistant", "content": [{"type": "thinking", "thinking": "The user is asking me to compute 2+3 and provide the answer and a short word. This is a straightforward math question.\n\n2 + 3 = 5\n\nThey want me to use the StructuredOutput tool to provide the answer in a structured format with an integer answer and a word.\n\nLooking at the tool definition:\n```\n{\n \"answer\": {\"type\": \"integer\"},\n \"word\": {\"type\": \"string\"}\n}\n```\n\nSo I need to return:\n- answer: 5\n- word: some short word\n\nI'll choose a simple word like \"sum\" or \"total\" or just a generic short word."}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 28946, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 28946, "ephemeral_1h_input_tokens": 0}, "output_tokens": 4, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "e0b4fb7e-b899-44ac-81fd-62841efa5380", "uuid": "27e9734a-846e-4ff2-b525-b249b1ff7b7f", "timestamp": "2026-09-30T17:05:35.124Z"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZzH4cNfTXjBxTRhX5Vm", "type": "message", "role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01NoK5Gm9S9NTEPEjjrjdfCf", "name": "StructuredOutput", "input": {"answer": 5, "word": "sum"}, "caller": {"type": "direct"}}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 28946, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 28946, "ephemeral_1h_input_tokens": 0}, "output_tokens": 4, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "e0b4fb7e-b899-44ac-81fd-62841efa5380", "uuid": "9e627359-3d66-4ce4-b893-330ebc0546f3", "timestamp": "2026-09-30T17:05:35.433Z"} +{"type": "user", "message": {"role": "user", "content": [{"tool_use_id": "toolu_01NoK5Gm9S9NTEPEjjrjdfCf", "type": "tool_result", "content": "Structured output provided successfully"}]}, "parent_tool_use_id": null, "session_id": "e0b4fb7e-b899-44ac-81fd-62841efa5380", "uuid": "939bc606-79ae-492f-a142-fbca6e400489", "timestamp": "2026-09-30T17:05:35.436Z", "tool_use_result": "Structured output provided successfully"} +{"duration_api_ms": 3171, "stop_reason": "tool_use", "session_id": "e0b4fb7e-b899-44ac-81fd-62841efa5380", "total_cost_usd": 0.0373165, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 28946, "cache_read_input_tokens": 0, "output_tokens": 225, "output_tokens_details": {"thinking_tokens": 151}, "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": 28946}, "inference_geo": "not_available", "iterations": [], "speed": "standard", "fallback_credit": null}, "modelUsage": {"claude-haiku-4-5-20251001": {"inputTokens": 9, "outputTokens": 225, "cacheReadInputTokens": 0, "cacheCreationInputTokens": 28946, "webSearchRequests": 0, "costUSD": 0.0373165, "contextWindow": 200000, "maxOutputTokens": 32000, "thinkingTokens": 151, "canonicalModel": "claude-haiku-4-5", "provider": "firstParty", "costBasis": "list"}}, "permission_denials": [], "terminal_reason": "completed", "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": false, "num_turns": 2, "subtype": "success", "api_error_status": null, "result": "{\"answer\":5,\"word\":\"sum\"}", "structured_output": {"answer": 5, "word": "sum"}, "ttft_ms": 2886, "type": "result", "duration_ms": 3202, "uuid": "dca4570c-3d4e-4442-9544-5d47d9ce4268", "ttft_stream_ms": 1288, "time_to_request_ms": 31, "first_content_frame_ms": 1288, "queued_turn_count": 0, "result_index": 0} diff --git a/tests/unit/llms/claude_code/harness/fixtures/success_tools.jsonl b/tests/unit/llms/claude_code/harness/fixtures/success_tools.jsonl new file mode 100644 index 00000000000..b84d31447be --- /dev/null +++ b/tests/unit/llms/claude_code/harness/fixtures/success_tools.jsonl @@ -0,0 +1,12 @@ +{"type": "system", "subtype": "init", "cwd": "/workspace", "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "tools": ["Task", "Bash", "CronCreate", "CronDelete", "CronList", "Edit", "EnterWorktree", "ExitWorktree", "ListAgents", "NotebookEdit", "Read", "ReportFindings", "ScheduleWakeup", "SendMessage", "Skill", "TaskCreate", "TaskGet", "TaskList", "TaskStop", "TaskUpdate", "WebFetch", "WebSearch", "Workflow", "Write"], "mcp_servers": [], "model": "claude-haiku-4-5-20251001", "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": "71b745d7-51cf-49c7-9e0d-9fa08a069e09", "fast_mode_state": "off", "fast_mode_disabled_reason": "sdk_opt_in_required", "per_turn_effort_active": false, "view_mode": "default"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 1, "estimated_tokens_delta": 1, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "a2ed41b6-6db0-4667-aad5-f20ffcf19ce4"} +{"type": "system", "subtype": "thinking_tokens", "estimated_tokens": 4, "estimated_tokens_delta": 3, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "f89d8c8f-2c16-4e1e-9d0f-16af3d94bb2c"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZynHUDvD1eucc6XZnzm", "type": "message", "role": "assistant", "content": [{"type": "thinking", "thinking": "The user wants me to:\n1. Create a file hello.txt with the content \"hi\"\n2. Read it back\n3. Be brief\n\nI should use the Write tool to create the file, then the Read tool to read it back. Since these operations depend on each other (I need to write first), I should do them sequentially.\n\nLet me create the file first with Write, then read it back with Read."}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 28852, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 28852, "ephemeral_1h_input_tokens": 0}, "output_tokens": 1, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "d2fab3ed-aee9-4122-99d3-ec5279016428", "timestamp": "2026-09-30T16:59:04.048Z"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZynHUDvD1eucc6XZnzm", "type": "message", "role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01DFhmKzT5x1NzxuestG2Hkj", "name": "Write", "input": {"file_path": "/workspace/hello.txt", "content": "hi"}, "caller": {"type": "direct"}}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 9, "cache_creation_input_tokens": 28852, "cache_read_input_tokens": 0, "cache_creation": {"ephemeral_5m_input_tokens": 28852, "ephemeral_1h_input_tokens": 0}, "output_tokens": 1, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "8eca3aef-aa68-48c3-82cc-9073fb1fb653", "timestamp": "2026-09-30T16:59:04.408Z"} +{"type": "user", "message": {"role": "user", "content": [{"tool_use_id": "toolu_01DFhmKzT5x1NzxuestG2Hkj", "type": "tool_result", "content": "File created successfully at: /workspace/hello.txt (file state is current in your context — no need to Read it back)"}]}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "5c07cb84-d183-4d41-abbf-aa81e6bccf24", "timestamp": "2026-09-30T16:59:04.421Z", "tool_use_result": {"type": "create", "filePath": "/workspace/hello.txt", "content": "hi", "structuredPatch": [], "originalFile": null, "userModified": false}} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZynchNSG2pNMhNYitM5", "type": "message", "role": "assistant", "content": [{"type": "thinking", "thinking": "The Write tool already created the file successfully. Now I need to read it back using the Read tool to confirm the contents."}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 8, "cache_creation_input_tokens": 265, "cache_read_input_tokens": 28852, "cache_creation": {"ephemeral_5m_input_tokens": 265, "ephemeral_1h_input_tokens": 0}, "output_tokens": 2, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "4db547e1-d93c-47b7-bcc2-782e17f7e445", "timestamp": "2026-09-30T16:59:07.726Z"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZynchNSG2pNMhNYitM5", "type": "message", "role": "assistant", "content": [{"type": "tool_use", "id": "toolu_017jG7ow2c8KCBUTq4Z9cDG8", "name": "Read", "input": {"file_path": "/workspace/hello.txt"}, "caller": {"type": "direct"}}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 8, "cache_creation_input_tokens": 265, "cache_read_input_tokens": 28852, "cache_creation": {"ephemeral_5m_input_tokens": 265, "ephemeral_1h_input_tokens": 0}, "output_tokens": 2, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "fd6c4e76-87b1-46a9-80dd-98c1e0ee2244", "timestamp": "2026-09-30T16:59:07.881Z"} +{"type": "user", "message": {"role": "user", "content": [{"tool_use_id": "toolu_017jG7ow2c8KCBUTq4Z9cDG8", "type": "tool_result", "content": "1\thi"}]}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "ac34772d-30f2-4fed-a0ec-5c9e50e2c864", "timestamp": "2026-09-30T16:59:07.895Z", "tool_use_result": {"type": "text", "file": {"filePath": "/workspace/hello.txt", "content": "hi", "numLines": 1, "startLine": 1, "totalLines": 1}}} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZyoDMCMiseBH9a2J4bZ", "type": "message", "role": "assistant", "content": [{"type": "thinking", "thinking": "Perfect! I've created the file and read it back. The file contains \"hi\" as expected. The task is complete."}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 8, "cache_creation_input_tokens": 149, "cache_read_input_tokens": 29117, "cache_creation": {"ephemeral_5m_input_tokens": 149, "ephemeral_1h_input_tokens": 0}, "output_tokens": 3, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "ec814481-9a98-419f-8a2f-96e57be0f3d6", "timestamp": "2026-09-30T16:59:16.170Z"} +{"type": "assistant", "message": {"model": "claude-haiku-4-5-20251001", "id": "msg_011CfZyoDMCMiseBH9a2J4bZ", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "Done. Created `hello.txt` with content \"hi\" and confirmed it reads back correctly."}], "container": null, "stop_reason": null, "stop_sequence": null, "stop_details": null, "usage": {"input_tokens": 8, "cache_creation_input_tokens": 149, "cache_read_input_tokens": 29117, "cache_creation": {"ephemeral_5m_input_tokens": 149, "ephemeral_1h_input_tokens": 0}, "output_tokens": 3, "service_tier": "standard", "inference_geo": "not_available"}, "diagnostics": null, "context_management": null}, "parent_tool_use_id": null, "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "uuid": "563dcf5b-5a7d-4d10-9dbb-9924c2f0b09f", "timestamp": "2026-09-30T16:59:16.434Z"} +{"duration_api_ms": 19780, "stop_reason": "end_turn", "session_id": "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34", "total_cost_usd": 0.044094400000000006, "usage": {"input_tokens": 25, "cache_creation_input_tokens": 29266, "cache_read_input_tokens": 57969, "output_tokens": 338, "output_tokens_details": {"thinking_tokens": 157}, "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": 29266}, "inference_geo": "not_available", "iterations": [], "speed": "standard", "fallback_credit": null}, "modelUsage": {"claude-haiku-4-5-20251001": {"inputTokens": 25, "outputTokens": 338, "cacheReadInputTokens": 57969, "cacheCreationInputTokens": 29266, "webSearchRequests": 0, "costUSD": 0.044094400000000006, "contextWindow": 200000, "maxOutputTokens": 32000, "thinkingTokens": 157, "canonicalModel": "claude-haiku-4-5", "provider": "firstParty", "costBasis": "list"}}, "permission_denials": [], "terminal_reason": "completed", "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": false, "num_turns": 3, "subtype": "success", "api_error_status": null, "result": "Done. Created `hello.txt` with content \"hi\" and confirmed it reads back correctly.", "ttft_ms": 7217, "type": "result", "duration_ms": 19835, "uuid": "702bf2fe-16b0-41e8-afb9-efe37f17abe3", "ttft_stream_ms": 6540, "time_to_request_ms": 27, "first_content_frame_ms": 6541, "queued_turn_count": 0, "result_index": 0} diff --git a/tests/unit/llms/claude_code/harness/test_transformation.py b/tests/unit/llms/claude_code/harness/test_transformation.py new file mode 100644 index 00000000000..6348b4f6a3e --- /dev/null +++ b/tests/unit/llms/claude_code/harness/test_transformation.py @@ -0,0 +1,708 @@ +"""Unit tests for the Claude Code harness config. No network, no real CLI. + +Fixtures under fixtures/ are sanitized stream-json recorded from Claude Code +2.1.285 through a LiteLLM gateway. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import pytest +from pydantic import BaseModel + +from litellm.harness.context import SessionContext +from litellm.harness.errors import ( + HarnessError, + HarnessInstallFailed, + OptionsMismatch, +) +from litellm.harness.handlers.cli_handler import CLIHarnessHandler +from litellm.harness.options import ClaudeCodeOptions, CodexOptions +from litellm.harness.sandbox.base import CompletedRun +from litellm.harness.types import ( + Compaction, + Harness, + Reasoning, + Text, + ToolCall, + ToolResult, +) +from litellm.llms.base_llm.harness.transformation import ( + HarnessTurnError, + HarnessTurnRequest, +) +from litellm.llms.base_llm.harness.utils import ( + decode_json_line, + last_json_object, + native_tool_names, +) +from litellm.llms.claude_code.harness.transformation import ( + MANAGED_CONFIG_KEYS, + MANAGED_ENV_KEYS, + NORMALIZED_TO_NATIVE, + PERMISSION_MODES, + ClaudeCodeHarnessConfig, + ClaudeCodeStreamState, + build_system_prompt, + stringify_tool_output, + turn_error_message, +) + +FIXTURES = Path(__file__).parent / "fixtures" +SESSION_ID = "5ef64ff1-d2af-4c38-a7ca-17b4a9d07d34" +TOKEN = "per-session-token-abc" +PORT = 53211 +PRIV = "/priv" + + +def fixture_lines(name: str) -> list[str]: + return (FIXTURES / name).read_text().splitlines() + + +def parse_line(line: str, state: ClaudeCodeStreamState) -> list[Any]: + decoded = decode_json_line(line) + if decoded is None: + return [] + return ClaudeCodeHarnessConfig().transform_stream_line(decoded, state) + + +def parse_fixture(name: str) -> tuple[list[Any], ClaudeCodeStreamState]: + state = ClaudeCodeHarnessConfig().create_stream_state() + events: list[Any] = [] + for line in fixture_lines(name): + events.extend(parse_line(line, state)) + return events, state + + +class FakeEndpoint: + port = PORT + token = TOKEN + + +class FakeProcess: + def __init__(self, stdout: bytes, stderr: bytes, exit_code: int) -> None: + self.stdin_data = bytearray() + self.stdin_closed = False + self.killed = False + self._exit_code = exit_code + 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.stdin = FakeStdin(self) + + async def wait(self) -> int: + return self._exit_code + + async def kill(self) -> None: + self.killed = True + + +class FakeStdin: + def __init__(self, proc: FakeProcess) -> None: + self._proc = proc + + def write(self, data: bytes) -> None: + self._proc.stdin_data.extend(data) + + async def drain(self) -> None: + return None + + def close(self) -> None: + self._proc.stdin_closed = True + + +class FakeSandbox: + def __init__( + self, + workdir: str, + outputs: list[tuple[str, bytes, int]], + binary: str | None = "/usr/bin/claude", + tempdir: str | None = None, + ) -> None: + self.workdir = workdir + self.binary = binary + self.outputs = list(outputs) + self.calls: list[dict[str, Any]] = [] + self.runs: list[list[str]] = [] + self.procs: list[FakeProcess] = [] + self.written: dict[str, bytes] = {} + self._tempdir = tempdir or os.path.join(workdir, "_cfg") + + async def exec( + self, + cmd: list[str], + *, + env: Mapping[str, str] | None = None, + cwd: str | None = None, + ) -> FakeProcess: + self.calls.append({"cmd": cmd, "env": dict(env or {}), "cwd": cwd}) + fixture, stderr, code = self.outputs.pop(0) + stdout = (FIXTURES / fixture).read_bytes() if fixture else b"" + proc = FakeProcess(stdout, stderr, code) + self.procs.append(proc) + return proc + + async def run(self, cmd: list[str], **kwargs: Any) -> CompletedRun: + self.runs.append(cmd) + return CompletedRun("", "", 0) + + async def read(self, path: str) -> bytes: + return self.written[path] + + async def write(self, path: str, data: bytes) -> None: + self.written[path] = data + + def host_url(self, port: int) -> str: + return f"http://host.docker.internal:{port}" + + async def which(self, binary: str) -> str | None: + return self.binary + + async def tempdir(self) -> str: + return self._tempdir + + async def snapshot(self) -> dict[str, str]: + return {} + + async def close(self) -> None: + return None + + +class Answer(BaseModel): + answer: int + word: str + + +def make_ctx(sandbox: FakeSandbox, **overrides: Any) -> SessionContext: + values: dict[str, Any] = { + "harness": Harness.CLAUDE_CODE, + "sandbox": sandbox, + "session_id": "hs_1", + "model": "claude-haiku-4-5-20251001", + "endpoint": FakeEndpoint(), + **overrides, + } + return SessionContext(**values) + + +def pure_ctx(tmp_path: Path, **overrides: Any) -> SessionContext: + return make_ctx(FakeSandbox(str(tmp_path), []), **overrides) + + +def make_handler() -> CLIHarnessHandler: + return CLIHarnessHandler(ClaudeCodeHarnessConfig()) + + +def request_for( + ctx: SessionContext, native_session_id: str | None = None, prompt: str = "hi" +) -> HarnessTurnRequest: + cfg = ClaudeCodeHarnessConfig() + setup = cfg.transform_session_setup(ctx, PRIV) + return cfg.transform_turn_request(ctx, setup, PRIV, prompt, native_session_id) + + +async def run_turn(handler: CLIHarnessHandler, ctx: SessionContext, prompt: str): + return [event async for event in handler.turn(ctx, prompt)] + + +# --------------------------------------------------------------------------- +# Parsing +# --------------------------------------------------------------------------- + + +def test_parse_success_fixture_events(): + events, state = parse_fixture("success_tools.jsonl") + kinds = [type(e).__name__ for e in events] + assert kinds == [ + "Reasoning", + "ToolCall", + "ToolResult", + "Reasoning", + "ToolCall", + "ToolResult", + "Reasoning", + "Text", + ] + write_call, read_call = events[1], events[4] + assert write_call == ToolCall( + id="toolu_01DFhmKzT5x1NzxuestG2Hkj", + name="write", + native_name="Write", + input={"file_path": "/workspace/hello.txt", "content": "hi"}, + builtin=True, + ) + assert read_call.name == "read" and read_call.native_name == "Read" + assert events[2].id == write_call.id and events[2].is_error is False + assert events[5].output == "1\thi" + assert state.session_id == SESSION_ID + assert ClaudeCodeHarnessConfig().get_native_session_id(state) == SESSION_ID + assert state.result_seen and not state.is_error + assert state.final_text.startswith("Done. Created `hello.txt`") + + +def test_parse_api_error_fixture_skips_synthetic_text(): + events, state = parse_fixture("api_error.jsonl") + assert events == [] + assert state.is_error + assert "no healthy deployments" in (state.result_text or "") + + +def test_parse_max_turns_fixture(): + events, state = parse_fixture("max_turns.jsonl") + assert [e.native_name for e in events if isinstance(e, ToolCall)] == [ + "Write", + "Write", + "Write", + ] + assert state.is_error and state.result_text is None + assert state.errors == ["Reached maximum number of turns (1)"] + + +def test_parse_structured_output_fixture(): + _, state = parse_fixture("structured_output.jsonl") + assert state.structured_output == {"answer": 5, "word": "sum"} + + +def test_parse_compaction_and_garbage(): + state = ClaudeCodeStreamState() + line = json.dumps( + { + "type": "system", + "subtype": "compact_boundary", + "compact_metadata": {"trigger": "auto", "pre_tokens": 1234}, + } + ) + assert parse_line(line, state) == [ + Compaction(tokens_before=1234, tokens_after=None) + ] + assert parse_line("not json", state) == [] + assert parse_line("", state) == [] + assert parse_line("[1,2]", state) == [] + cfg = ClaudeCodeHarnessConfig() + assert cfg.transform_stream_line({"type": "unknown"}, state) == [] + + +def test_parse_skips_subagent_messages_and_maps_errors(): + cfg = ClaudeCodeHarnessConfig() + state = ClaudeCodeStreamState() + sub = { + "type": "assistant", + "parent_tool_use_id": "toolu_parent", + "message": {"content": [{"type": "text", "text": "inner"}]}, + } + assert cfg.transform_stream_line(sub, state) == [] + err = { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "t1", + "is_error": True, + "content": [{"type": "text", "text": "boom"}], + } + ] + }, + } + assert cfg.transform_stream_line(err, state) == [ + ToolResult(id="t1", output="boom", is_error=True) + ] + + +def test_parse_thinking_and_mcp_tools(): + state = ClaudeCodeStreamState() + msg = { + "type": "assistant", + "message": { + "content": [ + {"type": "thinking", "thinking": "hmm"}, + {"type": "tool_use", "id": "t", "name": "mcp__x__y", "input": {}}, + {"type": "tool_use", "id": "u", "name": "MultiEdit", "input": {}}, + ] + }, + } + events = ClaudeCodeHarnessConfig().transform_stream_line(msg, state) + assert events[0] == Reasoning(delta="hmm") + assert events[1].name == "mcp__x__y" and events[1].builtin is False + assert events[2].name == "edit" + + +def test_stringify_tool_output_variants(): + assert stringify_tool_output(None) == "" + assert stringify_tool_output("x") == "x" + assert stringify_tool_output([{"type": "text", "text": "a"}, "b"]) == "a\nb" + assert stringify_tool_output({"k": 1}) == '{"k": 1}' + + +def test_extract_last_json_object(): + text = 'first {"a": 1} then {not json} and finally {"b": {"c": 2}}' + assert json.loads(last_json_object(text) or "") == {"b": {"c": 2}} + assert last_json_object("no json here") is None + + +# --------------------------------------------------------------------------- +# Session setup / turn request (argv + env) +# --------------------------------------------------------------------------- + + +def test_native_disallowed_tools_mapping(): + natives = native_tool_names(["edit", "bash", "Task", "edit"], NORMALIZED_TO_NATIVE) + assert natives == ["Edit", "MultiEdit", "Bash", "Task"] + + +@pytest.mark.parametrize( + "permissions,native", + [ + ("read-only", "plan"), + ("edit", "acceptEdits"), + ("full", "bypassPermissions"), + ], +) +def test_turn_request_permission_modes(tmp_path, permissions, native): + assert PERMISSION_MODES[permissions] == native + argv = list(request_for(pure_ctx(tmp_path, permissions=permissions)).argv) + assert argv[argv.index("--permission-mode") + 1] == native + assert "--resume" not in argv + assert argv[argv.index("--setting-sources") + 1] == "user" + + +def test_session_setup_and_turn_request_env_and_command(tmp_path): + ctx = pure_ctx( + tmp_path, + instructions="Be terse.", + disable_tools=["bash", "web_search"], + max_turns=7, + options=ClaudeCodeOptions(config={"cleanupPeriodDays": 1}, env={"X": "1"}), + ) + cfg = ClaudeCodeHarnessConfig() + setup = cfg.transform_session_setup(ctx, PRIV) + assert setup.persisted_dirs == [("projects", "claude_code/projects")] + assert setup.skills_dir == "skills" + request = cfg.transform_turn_request(ctx, setup, PRIV, "do the thing", None) + env, cmd = request.env, list(request.argv) + assert request.stdin == "do the thing" + assert env["ANTHROPIC_AUTH_TOKEN"] == TOKEN + assert env["ANTHROPIC_API_KEY"] == "" + assert env["ANTHROPIC_BASE_URL"] == f"http://host.docker.internal:{PORT}" + assert env["ANTHROPIC_MODEL"] == "claude-haiku-4-5-20251001" + assert env["ANTHROPIC_SMALL_FAST_MODEL"] == "claude-haiku-4-5-20251001" + assert env["CLAUDE_CONFIG_DIR"] == PRIV + assert env["DISABLE_TELEMETRY"] == "1" + assert env["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"] == "1" + assert env["X"] == "1" + assert not any(TOKEN in a for a in cmd) + assert cmd[:7] == [ + "claude", + "-p", + "--output-format", + "stream-json", + "--verbose", + "--input-format", + "text", + ] + assert cmd[cmd.index("--model") + 1] == "claude-haiku-4-5-20251001" + assert cmd[cmd.index("--permission-mode") + 1] == "bypassPermissions" + assert cmd[cmd.index("--setting-sources") + 1] == "user" + assert cmd[cmd.index("--append-system-prompt") + 1] == "Be terse." + assert cmd[cmd.index("--max-turns") + 1] == "7" + assert json.loads(cmd[cmd.index("--settings") + 1]) == {"cleanupPeriodDays": 1} + assert cmd[cmd.index("--disallowedTools") + 1] == "Bash,WebSearch" + assert "--resume" not in cmd + + +def test_background_model_is_the_session_model(tmp_path): + env = ( + ClaudeCodeHarnessConfig().transform_session_setup(pure_ctx(tmp_path), PRIV).env + ) + assert env["ANTHROPIC_SMALL_FAST_MODEL"] == "claude-haiku-4-5-20251001" + + +def test_resume_argv(tmp_path): + argv = list(request_for(pure_ctx(tmp_path), "prior-session").argv) + assert argv[argv.index("--resume") + 1] == "prior-session" + + +def test_missing_endpoint_raises(tmp_path): + with pytest.raises(HarnessError, match="endpoint"): + ClaudeCodeHarnessConfig().transform_session_setup( + pure_ctx(tmp_path, endpoint=None), PRIV + ) + + +@pytest.mark.parametrize("key", sorted(MANAGED_ENV_KEYS)) +def test_options_env_cannot_override_managed_keys(tmp_path, key): + ctx = pure_ctx(tmp_path, options=ClaudeCodeOptions(env={key: "sk-real"})) + with pytest.raises(OptionsMismatch, match=key): + ClaudeCodeHarnessConfig().validate_environment(ctx) + + +def test_wrong_options_type_rejected(tmp_path): + with pytest.raises(OptionsMismatch): + ClaudeCodeHarnessConfig().validate_environment( + pure_ctx(tmp_path, options=CodexOptions()) + ) + + +def test_structured_output_system_prompt(tmp_path): + argv = list( + request_for(pure_ctx(tmp_path, output=Answer, instructions="Base.")).argv + ) + prompt = argv[argv.index("--append-system-prompt") + 1] + assert prompt.startswith("Base.\n\n") + assert json.dumps(Answer.model_json_schema()) in prompt + assert build_system_prompt(None, None) is None + + +# --------------------------------------------------------------------------- +# Turn response +# --------------------------------------------------------------------------- + + +def test_turn_response_api_error_includes_stderr(tmp_path): + _, state = parse_fixture("api_error.jsonl") + with pytest.raises(HarnessTurnError) as info: + ClaudeCodeHarnessConfig().transform_turn_response( + pure_ctx(tmp_path), state, 1, ["[claude-code:unrecognized_model] bad"] + ) + assert "no healthy deployments" in str(info.value) + assert "unrecognized_model" in str(info.value) + + +def test_turn_response_max_turns(tmp_path): + _, state = parse_fixture("max_turns.jsonl") + with pytest.raises(HarnessTurnError, match="maximum number of turns"): + ClaudeCodeHarnessConfig().transform_turn_response( + pure_ctx(tmp_path), state, 1, [] + ) + + +def test_turn_error_message_no_result(): + message = turn_error_message(ClaudeCodeStreamState(), 139, ["segfault", ""]) + assert message is not None + assert "code 139: no result event" in message and "segfault" in message + _, ok = parse_fixture("success_tools.jsonl") + assert turn_error_message(ok, 0, []) is None + + +def test_turn_response_structured_output_and_fallback(tmp_path): + cfg = ClaudeCodeHarnessConfig() + ctx = pure_ctx(tmp_path, output=Answer) + _, state = parse_fixture("structured_output.jsonl") + response = cfg.transform_turn_response(ctx, state, 0, []) + assert json.loads(response.output_json or "") == {"answer": 5, "word": "sum"} + + _, plain = parse_fixture("resume_turn.jsonl") + response = cfg.transform_turn_response(ctx, plain, 0, []) + assert response.final_text == "hello.txt" + assert response.output_json is None # "hello.txt" holds no JSON object + + text_json = ClaudeCodeStreamState( + result_seen=True, result_text='answer: {"answer": 1, "word": "x"}' + ) + response = cfg.transform_turn_response(ctx, text_json, 0, []) + assert json.loads(response.output_json or "") == {"answer": 1, "word": "x"} + + no_output = cfg.transform_turn_response(pure_ctx(tmp_path), state, 0, []) + assert no_output.output_json is None + + +# --------------------------------------------------------------------------- +# Through CLIHarnessHandler (start + turn) +# --------------------------------------------------------------------------- + + +async def test_start_and_turn_env_and_command(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [("success_tools.jsonl", b"", 0)]) + ctx = make_ctx(sandbox, options=ClaudeCodeOptions(env={"X": "1"})) + handler = make_handler() + await handler.start(ctx) + assert len(sandbox.runs) == 1 + assert sandbox.runs[0][:2] == ["sh", "-c"] + assert sandbox.runs[0][-2:] == [ + f"{tmp_path / '_cfg'}/projects", + "claude_code/projects", + ] + events = await run_turn(handler, ctx, "do the thing") + + call = sandbox.calls[0] + env, cmd = call["env"], call["cmd"] + assert env["ANTHROPIC_AUTH_TOKEN"] == TOKEN + assert env["CLAUDE_CONFIG_DIR"] == str(tmp_path / "_cfg") + assert env["X"] == "1" + assert not any(TOKEN in a for a in cmd) + assert cmd[cmd.index("--setting-sources") + 1] == "user" + + proc = sandbox.procs[0] + assert bytes(proc.stdin_data) == b"do the thing" and proc.stdin_closed + assert any(isinstance(e, Text) for e in events) + assert ctx.final_text.startswith("Done.") + assert handler.native_session_id() == SESSION_ID + + +async def test_second_turn_resumes_session(tmp_path): + sandbox = FakeSandbox( + str(tmp_path), + [("success_tools.jsonl", b"", 0), ("resume_turn.jsonl", b"", 0)], + ) + ctx = make_ctx(sandbox) + handler = make_handler() + await handler.start(ctx) + await run_turn(handler, ctx, "one") + await run_turn(handler, ctx, "two") + cmd = sandbox.calls[1]["cmd"] + assert cmd[cmd.index("--resume") + 1] == SESSION_ID + assert ctx.final_text == "hello.txt" + + +async def test_resume_sets_native_session_id(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [("resume_turn.jsonl", b"", 0)]) + ctx = make_ctx(sandbox) + handler = make_handler() + await handler.start(ctx) + await handler.resume(ctx, "prior-session") + assert handler.native_session_id() == "prior-session" + await run_turn(handler, ctx, "again") + cmd = sandbox.calls[0]["cmd"] + assert cmd[cmd.index("--resume") + 1] == "prior-session" + + +async def test_missing_binary_raises_install_failed(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [], binary=None) + with pytest.raises(HarnessInstallFailed, match="claude"): + await make_handler().start(make_ctx(sandbox)) + + +async def test_start_missing_endpoint_raises(tmp_path): + sandbox = FakeSandbox(str(tmp_path), []) + with pytest.raises(HarnessError, match="endpoint"): + await make_handler().start(make_ctx(sandbox, endpoint=None)) + + +async def test_start_rejects_managed_env(tmp_path): + sandbox = FakeSandbox(str(tmp_path), []) + options = ClaudeCodeOptions(env={"ANTHROPIC_API_KEY": "sk-real"}) + with pytest.raises(OptionsMismatch, match="ANTHROPIC_API_KEY"): + await make_handler().start(make_ctx(sandbox, options=options)) + assert sandbox.runs == [] and sandbox.written == {} + + +async def test_api_error_raises_turn_error_with_stderr(tmp_path): + stderr = b"[claude-code:unrecognized_model] bad model\n" + sandbox = FakeSandbox(str(tmp_path), [("api_error.jsonl", stderr, 1)]) + ctx = make_ctx(sandbox) + handler = make_handler() + await handler.start(ctx) + with pytest.raises(HarnessTurnError) as info: + await run_turn(handler, ctx, "hi") + assert "no healthy deployments" in str(info.value) + assert "unrecognized_model" in str(info.value) + + +async def test_max_turns_raises_turn_error(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [("max_turns.jsonl", b"", 1)]) + ctx = make_ctx(sandbox) + handler = make_handler() + await handler.start(ctx) + with pytest.raises(HarnessTurnError, match="maximum number of turns"): + await run_turn(handler, ctx, "hi") + + +async def test_nonzero_exit_without_result_raises(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [("", b"segfault\n", 139)]) + ctx = make_ctx(sandbox) + handler = make_handler() + await handler.start(ctx) + with pytest.raises(HarnessTurnError, match=r"code 139.*no result event") as info: + await run_turn(handler, ctx, "hi") + assert "segfault" in str(info.value) + + +async def test_structured_output_prompt_and_json(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [("structured_output.jsonl", b"", 0)]) + ctx = make_ctx(sandbox, output=Answer, instructions="Base.") + handler = make_handler() + await handler.start(ctx) + await run_turn(handler, ctx, "2+3?") + cmd = sandbox.calls[0]["cmd"] + prompt = cmd[cmd.index("--append-system-prompt") + 1] + assert prompt.startswith("Base.\n\n") + assert json.loads(ctx.output_json or "") == {"answer": 5, "word": "sum"} + + +async def test_structured_output_falls_back_to_final_text(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [("resume_turn.jsonl", b"", 0)]) + ctx = make_ctx(sandbox, output=Answer) + handler = make_handler() + await handler.start(ctx) + await run_turn(handler, ctx, "hi") + assert ctx.output_json is None # "hello.txt" holds no JSON object + + +async def test_skills_copied_into_private_config(tmp_path): + skill = tmp_path / "skills_src" / "demo" + (skill / "scripts").mkdir(parents=True) + (skill / "SKILL.md").write_text("---\nname: demo\n---\nSay DEMO.\n") + (skill / "scripts" / "run.sh").write_text("echo hi\n") + workdir = tmp_path / "work" + workdir.mkdir() + sandbox = FakeSandbox(str(workdir), [], tempdir="/cfg") + await make_handler().start(make_ctx(sandbox, skills=[str(skill)])) + assert sandbox.written == { + "/cfg/skills/demo/SKILL.md": b"---\nname: demo\n---\nSay DEMO.\n", + "/cfg/skills/demo/scripts/run.sh": b"echo hi\n", + } + + +async def test_skill_without_manifest_rejected(tmp_path): + skill = tmp_path / "bad" + skill.mkdir() + sandbox = FakeSandbox(str(tmp_path), []) + with pytest.raises(ValueError, match=r"SKILL\.md"): + await make_handler().start(make_ctx(sandbox, skills=[str(skill)])) + + +async def test_stop_kills_live_process(tmp_path): + sandbox = FakeSandbox(str(tmp_path), [("success_tools.jsonl", b"", 0)]) + ctx = make_ctx(sandbox) + handler = make_handler() + await handler.start(ctx) + stream = handler.turn(ctx, "hi") + await stream.__anext__() + proc = sandbox.procs[0] + await handler.stop(ctx) + assert proc.killed + await stream.aclose() + await handler.stop(ctx) # safe twice + + +def test_capabilities_match_spec(): + cfg = ClaudeCodeHarnessConfig() + caps = cfg.capabilities + assert cfg.harness is Harness.CLAUDE_CODE + assert cfg.options_type is ClaudeCodeOptions + assert cfg.get_binary() == "claude" + assert "@anthropic-ai/claude-code" in cfg.get_install_hint() + assert caps.structured_output and caps.tool_filtering and caps.skills + assert caps.resume + assert not (caps.tool_approval or caps.custom_tools or caps.history) + assert caps.permission_modes == frozenset({"read-only", "edit", "full"}) + + +@pytest.mark.parametrize("key", sorted(MANAGED_CONFIG_KEYS)) +def test_options_config_cannot_set_managed_keys(tmp_path, key): + ctx = pure_ctx(tmp_path, options=ClaudeCodeOptions(config={key: "x"})) + with pytest.raises(OptionsMismatch, match=key): + ClaudeCodeHarnessConfig().validate_environment(ctx) + + +def test_no_settings_flag_without_config(tmp_path): + assert "--settings" not in list(request_for(pure_ctx(tmp_path), None).argv) diff --git a/tests/unit/llms/codex/__init__.py b/tests/unit/llms/codex/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codex/harness/__init__.py b/tests/unit/llms/codex/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codex/harness/fixtures/__init__.py b/tests/unit/llms/codex/harness/fixtures/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codex/harness/fixtures/reasoning.jsonl b/tests/unit/llms/codex/harness/fixtures/reasoning.jsonl new file mode 100644 index 00000000000..e38a212baed --- /dev/null +++ b/tests/unit/llms/codex/harness/fixtures/reasoning.jsonl @@ -0,0 +1,5 @@ +{"type":"thread.started","thread_id":"01a0f347-4945-7f40-ae19-d1a2724ddee1"} +{"type":"turn.started"} +{"type":"item.completed","item":{"id":"item_0","type":"reasoning","text":"**Calculating multiplication**\n\nAlright, I need to respond with just the number. I multiply 17 and 23 to get 391. Let me check that: 20 times 23 equals 460, and if I subtract 3 times 23, which is 69, from 460, I get 391. So, yes, 391 is correct! I’ll provide the final answer as \"391\" only, without any extra text."}} +{"type":"item.completed","item":{"id":"item_1","type":"agent_message","text":"391"}} +{"type":"turn.completed","usage":{"input_tokens":9098,"cached_input_tokens":0,"output_tokens":68,"reasoning_output_tokens":0}} diff --git a/tests/unit/llms/codex/harness/fixtures/structured_output.jsonl b/tests/unit/llms/codex/harness/fixtures/structured_output.jsonl new file mode 100644 index 00000000000..0ecb79d9f1e --- /dev/null +++ b/tests/unit/llms/codex/harness/fixtures/structured_output.jsonl @@ -0,0 +1,6 @@ +{"type":"thread.started","thread_id":"01a0f344-34bf-7b82-9025-bc6db70f867e"} +{"type":"turn.started"} +{"type":"item.started","item":{"id":"item_0","type":"command_execution","command":"/bin/zsh -lc \"rg --files -g 'AGENTS.md' -g 'hello.txt' . && printf '\\\\n---\\\\n' && cat hello.txt\"","aggregated_output":"","exit_code":null,"status":"in_progress"}} +{"type":"item.completed","item":{"id":"item_0","type":"command_execution","command":"/bin/zsh -lc \"rg --files -g 'AGENTS.md' -g 'hello.txt' . && printf '\\\\n---\\\\n' && cat hello.txt\"","aggregated_output":"./hello.txt\n\n---\nhello world\n","exit_code":0,"status":"completed"}} +{"type":"item.completed","item":{"id":"item_1","type":"agent_message","text":"{\"file\":\"hello.txt\",\"content\":\"Name: `hello.txt`\\nContent: `hello world`\"}"}} +{"type":"turn.completed","usage":{"input_tokens":24271,"cached_input_tokens":3758,"output_tokens":96,"reasoning_output_tokens":0}} diff --git a/tests/unit/llms/codex/harness/fixtures/turn1_bash.jsonl b/tests/unit/llms/codex/harness/fixtures/turn1_bash.jsonl new file mode 100644 index 00000000000..110ae70999d --- /dev/null +++ b/tests/unit/llms/codex/harness/fixtures/turn1_bash.jsonl @@ -0,0 +1,7 @@ +{"type":"thread.started","thread_id":"01a0f341-fe37-7072-93b3-055358e8147f"} +{"type":"turn.started"} +{"type":"item.completed","item":{"id":"item_0","type":"agent_message","text":"I’ll create the file, then print it back to confirm."}} +{"type":"item.started","item":{"id":"item_1","type":"command_execution","command":"/bin/zsh -lc \"printf 'hi\n' > hello.txt && cat hello.txt\"","aggregated_output":"","exit_code":null,"status":"in_progress"}} +{"type":"item.completed","item":{"id":"item_1","type":"command_execution","command":"/bin/zsh -lc \"printf 'hi\n' > hello.txt && cat hello.txt\"","aggregated_output":"hi\n","exit_code":0,"status":"completed"}} +{"type":"item.completed","item":{"id":"item_2","type":"agent_message","text":"Done — `hello.txt` now contains `hi`, and `cat hello.txt` prints:\n\n```text\nhi\n```"}} +{"type":"turn.completed","usage":{"input_tokens":24372,"cached_input_tokens":13998,"output_tokens":103,"reasoning_output_tokens":0}} diff --git a/tests/unit/llms/codex/harness/fixtures/turn2_resume_apply_patch.jsonl b/tests/unit/llms/codex/harness/fixtures/turn2_resume_apply_patch.jsonl new file mode 100644 index 00000000000..4cb5d924352 --- /dev/null +++ b/tests/unit/llms/codex/harness/fixtures/turn2_resume_apply_patch.jsonl @@ -0,0 +1,7 @@ +{"type":"thread.started","thread_id":"01a0f341-fe37-7072-93b3-055358e8147f"} +{"type":"turn.started"} +{"type":"item.completed","item":{"id":"item_0","type":"agent_message","text":"I’ll patch `hello.txt` directly, then I’ll reply exactly as requested."}} +{"type":"item.started","item":{"id":"item_1","type":"command_execution","command":"/bin/zsh -lc \"apply_patch '*** Begin Patch\n*** Delete File: hello.txt\n*** Add File: hello.txt\n+hello world\n*** End Patch'\"","aggregated_output":"","exit_code":null,"status":"in_progress"}} +{"type":"item.completed","item":{"id":"item_1","type":"command_execution","command":"/bin/zsh -lc \"apply_patch '*** Begin Patch\n*** Delete File: hello.txt\n*** Add File: hello.txt\n+hello world\n*** End Patch'\"","aggregated_output":"Success. Updated the following files:\nA hello.txt\nD hello.txt\n","exit_code":0,"status":"completed"}} +{"type":"item.completed","item":{"id":"item_2","type":"agent_message","text":"done"}} +{"type":"turn.completed","usage":{"input_tokens":49569,"cached_input_tokens":38236,"output_tokens":204,"reasoning_output_tokens":0}} diff --git a/tests/unit/llms/codex/harness/fixtures/turn_failed.jsonl b/tests/unit/llms/codex/harness/fixtures/turn_failed.jsonl new file mode 100644 index 00000000000..7e4404f6ff5 --- /dev/null +++ b/tests/unit/llms/codex/harness/fixtures/turn_failed.jsonl @@ -0,0 +1,4 @@ +{"type": "thread.started", "thread_id": "01a0f344-5b83-7d40-a964-7d56c9e4ec9a"} +{"type": "turn.started"} +{"type": "error", "message": "{\"error\":{\"message\":\"litellm.BadRequestError: You passed in model=no-such-model-xyz. There are no healthy deployments for this model\",\"type\":\"invalid_request_error\",\"param\":null,\"code\":\"400\"}}"} +{"type": "turn.failed", "error": {"message": "{\"error\":{\"message\":\"litellm.BadRequestError: You passed in model=no-such-model-xyz. There are no healthy deployments for this model\",\"type\":\"invalid_request_error\",\"param\":null,\"code\":\"400\"}}"}} diff --git a/tests/unit/llms/codex/harness/test_transformation.py b/tests/unit/llms/codex/harness/test_transformation.py new file mode 100644 index 00000000000..6371f74fba6 --- /dev/null +++ b/tests/unit/llms/codex/harness/test_transformation.py @@ -0,0 +1,670 @@ +"""Unit tests for the Codex harness config. No network, no real CLI. + +Fixtures under fixtures/ are sanitized `codex exec --json` output recorded from +codex-cli through a LiteLLM gateway. +""" + +import asyncio +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +import pytest +from pydantic import BaseModel + +from litellm.harness.context import SessionContext +from litellm.harness.errors import HarnessError, HarnessInstallFailed, OptionsMismatch +from litellm.harness.handlers.cli_handler import CLIHarnessHandler +from litellm.harness.options import ClaudeCodeOptions, CodexOptions +from litellm.harness.sandbox.base import CompletedRun +from litellm.harness.sandbox.docker import DockerSandbox +from litellm.harness.types import Harness, Reasoning, Text, ToolCall, ToolResult +from litellm.llms.base_llm.harness.transformation import ( + HarnessTurnError, + HarnessTurnRequest, +) +from litellm.llms.base_llm.harness.utils import strict_json_schema +from litellm.llms.codex.harness.transformation import ( + CODEX_SCHEMA_FILENAME, + CODEX_TOKEN_ENV, + MANAGED_CONFIG_KEYS, + CodexHarnessConfig, + CodexStreamState, + config_overrides, + toml_key, + toml_value, +) + +FIXTURES = Path(__file__).parent / "fixtures" +TOKEN = "tok-secret-123" +HOME = "/tmp/codex-home" +THREAD_ID = "01a0f341-fe37-7072-93b3-055358e8147f" + + +def load_fixture(name: str) -> list[dict]: + return [ + json.loads(line) for line in (FIXTURES / name).read_text().splitlines() if line + ] + + +def parse_event(obj: dict, state: CodexStreamState) -> list: + return CodexHarnessConfig().transform_stream_line(obj, state) + + +def parse_all(name: str, state: Optional[CodexStreamState] = None): + state = state or CodexHarnessConfig().create_stream_state() + events = [] + for obj in load_fixture(name): + events.extend(parse_event(obj, state)) + return events, state + + +# --------------------------------------------------------------------------- fakes + + +class FakeProcess: + def __init__(self, stdout: bytes, stderr: bytes = b"", exit_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._exit_code = exit_code + self.killed = False + + async def wait(self) -> int: + return self._exit_code + + async def kill(self) -> None: + self.killed = True + + +class FakeStdin: + def __init__(self): + 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 + + +@dataclass +class FakeSandbox: + workdir: str = "/work" + has_codex: bool = True + outputs: list = field(default_factory=list) + files: dict = field(default_factory=dict) + execs: list = field(default_factory=list) + runs: list = field(default_factory=list) + processes: list = field(default_factory=list) + + async def exec(self, cmd, *, env=None, cwd=None): + self.execs.append({"cmd": cmd, "env": dict(env or {}), "cwd": cwd}) + proc = self.outputs.pop(0) + self.processes.append(proc) + return proc + + async def run(self, cmd, *, env=None, cwd=None, timeout=None): + self.runs.append(cmd) + return CompletedRun("", "", 0) + + async def read(self, path): + return self.files[path] + + async def write(self, path, data): + self.files[path] = data + + def host_url(self, port): + return f"http://127.0.0.1:{port}" + + async def which(self, binary): + return f"/usr/bin/{binary}" if self.has_codex else None + + async def tempdir(self): + return HOME + + async def snapshot(self): + return {} + + async def close(self): + return None + + +@dataclass +class FakeEndpoint: + port: int = 4555 + token: str = TOKEN + + +class Answer(BaseModel): + file: str + content: str + + +class Nested(BaseModel): + answer: Answer + tags: list[str] = [] + note: Optional[str] = None + + +def make_ctx(sandbox, **kwargs) -> SessionContext: + return SessionContext( + harness=Harness.CODEX, + sandbox=sandbox, + session_id="s1", + model=kwargs.pop("model", "gpt-5.4"), + endpoint=kwargs.pop("endpoint", FakeEndpoint()), + **kwargs, + ) + + +def request_for( + ctx: SessionContext, native_session_id: Optional[str] = None, prompt: str = "hi" +) -> HarnessTurnRequest: + cfg = CodexHarnessConfig() + setup = cfg.transform_session_setup(ctx, HOME) + return cfg.transform_turn_request(ctx, setup, HOME, prompt, native_session_id) + + +def argv_for(ctx: SessionContext, native_session_id: Optional[str] = None) -> list: + return list(request_for(ctx, native_session_id).argv) + + +def fixture_proc(name: str, **kwargs) -> FakeProcess: + return FakeProcess((FIXTURES / name).read_bytes(), **kwargs) + + +def config_values(argv: list[str]) -> list[str]: + return [argv[i + 1] for i, a in enumerate(argv) if a == "-c"] + + +def make_handler() -> CLIHarnessHandler: + return CLIHarnessHandler(CodexHarnessConfig()) + + +async def collect(handler, ctx, prompt): + return [e async for e in handler.turn(ctx, prompt)] + + +# --------------------------------------------------------------------------- parsing + + +def test_parse_bash_turn(): + events, state = parse_all("turn1_bash.jsonl") + assert state.thread_id == THREAD_ID + assert CodexHarnessConfig().get_native_session_id(state) == THREAD_ID + assert [type(e) for e in events] == [Text, ToolCall, ToolResult, Text] + call, result = events[1], events[2] + assert call.name == "bash" and call.native_name == "command_execution" + assert call.builtin is True + assert "hello.txt" in call.input["command"] + assert result.id == call.id == "item_1" + assert result.output == "hi\n" and result.is_error is False + assert state.final_text.startswith("Done") + assert not state.failed + + +def test_parse_reasoning(): + events, state = parse_all("reasoning.jsonl") + assert isinstance(events[0], Reasoning) and "391" in events[0].delta + assert events[1] == Text(delta="391") + assert state.final_text == "391" + + +def test_parse_turn_failed(): + events, state = parse_all("turn_failed.jsonl") + assert events == [] + assert state.failed + assert "no healthy deployments" in state.error + + +def test_parse_file_change_and_mcp_and_web_search(): + state = CodexStreamState() + change = { + "id": "i1", + "type": "file_change", + "changes": [{"path": "a.txt", "kind": "add"}], + "status": "completed", + } + events = parse_event({"type": "item.completed", "item": change}, state) + assert events[0] == ToolCall( + id="i1", + name="edit", + native_name="apply_patch", + input={"changes": [{"path": "a.txt", "kind": "add"}]}, + ) + assert events[1] == ToolResult(id="i1", output="add a.txt", is_error=False) + + mcp = { + "id": "i2", + "type": "mcp_tool_call", + "server": "docs", + "tool": "search", + "arguments": {"q": "x"}, + "status": "in_progress", + } + started = parse_event({"type": "item.started", "item": mcp}, state) + assert started == [ + ToolCall( + id="i2", + name="docs.search", + native_name="search", + input={"q": "x"}, + builtin=False, + ) + ] + done = {**mcp, "status": "failed", "error": {"message": "boom"}} + assert parse_event({"type": "item.completed", "item": done}, state) == [ + ToolResult(id="i2", output="boom", is_error=True) + ] + + web = {"id": "i3", "type": "web_search", "query": "litellm"} + events = parse_event({"type": "item.completed", "item": web}, state) + assert events[0].name == "web_search" and events[0].input == {"query": "litellm"} + + +def test_parse_failed_command_is_error_and_unknown_events_ignored(): + state = CodexStreamState() + item = { + "id": "c", + "type": "command_execution", + "command": "false", + "aggregated_output": "", + "exit_code": 1, + "status": "failed", + } + events = parse_event({"type": "item.completed", "item": item}, state) + assert events[1].is_error is True + usage = {"type": "turn.completed", "usage": {"input_tokens": 5}} + assert parse_event(usage, state) == [] + todo = {"type": "item.completed", "item": {"type": "todo_list"}} + assert parse_event(todo, state) == [] + assert parse_event({"type": "item.completed", "item": "nope"}, state) == [] + + +def test_parse_error_event_then_turn_failed(): + state = CodexStreamState() + assert parse_event({"type": "error", "message": "reconnecting"}, state) == [] + assert state.error == "reconnecting" and not state.failed + assert parse_event({"type": "turn.failed", "error": {"message": "x"}}, state) == [] + assert state.failed and state.error == "x" + + +# --------------------------------------------------------------------------- helpers + + +def test_strict_json_schema_recursive(): + schema = strict_json_schema(Nested.model_json_schema()) + assert schema["additionalProperties"] is False + assert schema["required"] == ["answer", "tags", "note"] + assert schema["properties"]["answer"] == {"$ref": "#/$defs/Answer"} + assert "default" not in schema["properties"]["tags"] + answer = schema["$defs"]["Answer"] + assert answer["additionalProperties"] is False + assert answer["required"] == ["file", "content"] + + +def test_config_overrides_rejects_managed_keys(): + for key in ( + "model_provider", + "model_providers.x.base_url", + "approval_policy", + "sandbox_mode", + "mcp_servers.a", + ): + with pytest.raises(OptionsMismatch): + config_overrides({key: "x"}) + for key in sorted(MANAGED_CONFIG_KEYS): + with pytest.raises(OptionsMismatch, match="managed by LiteLLM"): + config_overrides({key: "x"}) + for bad in ("", "a=b"): + with pytest.raises(OptionsMismatch, match="Invalid"): + config_overrides({bad: "x"}) + assert config_overrides( + { + "sandbox_workspace_write.network_access": True, + "notice": {"a b": 1}, + "x": ["y"], + } + ) == [ + "sandbox_workspace_write.network_access=true", + 'notice={"a b" = 1}', + 'x=["y"]', + ] + + +def test_toml_value_and_key(): + assert toml_value('say "hi"') == '"say \\"hi\\""' + assert toml_value(False) == "false" + assert toml_value(1.5) == "1.5" + assert toml_value(("a", 2)) == '["a", 2]' + assert toml_key("plain_key-1") == "plain_key-1" + assert toml_key("a b") == '"a b"' + with pytest.raises(OptionsMismatch): + toml_value(object()) + + +# --------------------------------------------------------------------------- session setup / turn request + + +def test_session_setup_env_and_schema(): + ctx = make_ctx(FakeSandbox(), output=Answer, options=CodexOptions(env={"X": "1"})) + setup = CodexHarnessConfig().transform_session_setup(ctx, HOME) + assert setup.env == {"X": "1", CODEX_TOKEN_ENV: TOKEN, "CODEX_HOME": HOME} + assert setup.persisted_dirs == [("sessions", "codex/sessions")] + assert setup.skills_dir == "skills" + schema = json.loads(setup.files[CODEX_SCHEMA_FILENAME]) + assert schema["additionalProperties"] is False + assert schema["required"] == ["file", "content"] + no_schema = CodexHarnessConfig().transform_session_setup( + make_ctx(FakeSandbox()), HOME + ) + assert no_schema.files == {} + + +def test_missing_endpoint_raises(): + ctx = make_ctx(FakeSandbox(), endpoint=None) + with pytest.raises(HarnessError, match="endpoint"): + CodexHarnessConfig().transform_session_setup(ctx, HOME) + + +def test_validate_environment_rejects_managed_config_and_wrong_options(): + cfg = CodexHarnessConfig() + with pytest.raises(OptionsMismatch): + cfg.validate_environment( + make_ctx( + FakeSandbox(), options=CodexOptions(config={"model_provider": "openai"}) + ) + ) + with pytest.raises(OptionsMismatch): + cfg.validate_environment(make_ctx(FakeSandbox(), options=ClaudeCodeOptions())) + + +def test_first_turn_argv_env(): + ctx = make_ctx( + FakeSandbox(), + instructions="Be terse.", + options=CodexOptions( + reasoning_effort="low", + config={"sandbox_workspace_write.network_access": True}, + ), + ) + request = request_for(ctx, prompt="create hello.txt") + argv, env = list(request.argv), request.env + assert request.stdin == "create hello.txt" + assert request.cwd == "/work" + assert argv[:4] == ["codex", "exec", "--json", "--skip-git-repo-check"] + assert argv[-1] == "-" and argv[argv.index("-C") + 1] == "/work" + assert argv[argv.index("-m") + 1] == "gpt-5.4" + assert argv[argv.index("--sandbox") + 1] == "workspace-write" + cfg = config_values(argv) + assert "model_provider=litellm" in cfg + assert 'model_providers.litellm.base_url="http://127.0.0.1:4555/v1"' in cfg + assert "model_providers.litellm.env_key=LITELLM_HARNESS_TOKEN" in cfg + assert "model_providers.litellm.wire_api=responses" in cfg + assert "approval_policy=never" in cfg + assert "model_reasoning_effort=low" in cfg + assert "model_reasoning_summary=auto" in cfg + assert "web_search=disabled" in cfg + assert 'developer_instructions="Be terse."' in cfg + assert "sandbox_workspace_write.network_access=true" in cfg + assert "--output-schema" not in argv + assert not any(TOKEN in a for a in argv) + assert env["LITELLM_HARNESS_TOKEN"] == TOKEN + assert env["CODEX_HOME"] == HOME + + +def test_resume_argv(): + argv = argv_for(make_ctx(FakeSandbox()), THREAD_ID) + assert argv[:4] == ["codex", "exec", "resume", THREAD_ID] + assert "--sandbox" not in argv and "-C" not in argv + assert 'sandbox_mode="workspace-write"' in config_values(argv) + assert argv[-1] == "-" + + +def test_permission_modes(): + ro = make_ctx(FakeSandbox(), permissions="read-only") + argv = argv_for(ro) + assert argv[argv.index("--sandbox") + 1] == "read-only" + assert 'sandbox_mode="read-only"' in config_values(argv_for(ro, "t")) + + # The container is the boundary: DockerSandbox opts out of codex's own sandbox. + assert DockerSandbox.is_container is True + container = FakeSandbox(workdir="/workspace") + container.is_container = True + argv = argv_for(make_ctx(container, permissions="full")) + assert "--dangerously-bypass-approvals-and-sandbox" in argv + assert "--sandbox" not in argv + assert argv[argv.index("-C") + 1] == "/workspace" + resumed = argv_for(make_ctx(container, permissions="full"), "t") + assert "--dangerously-bypass-approvals-and-sandbox" in resumed + assert not any(v.startswith("sandbox_mode=") for v in config_values(resumed)) + + # read-only wins even inside a container + argv = argv_for(make_ctx(container, permissions="read-only")) + assert argv[argv.index("--sandbox") + 1] == "read-only" + assert "--dangerously-bypass-approvals-and-sandbox" not in argv + + web = make_ctx(FakeSandbox(), options=CodexOptions(web_search=True)) + assert "web_search=live" in config_values(argv_for(web)) + + +def test_structured_output_argv(): + argv = argv_for(make_ctx(FakeSandbox(), output=Answer)) + assert argv[argv.index("--output-schema") + 1] == f"{HOME}/{CODEX_SCHEMA_FILENAME}" + + +def test_no_model_omits_flag(): + assert "-m" not in argv_for(make_ctx(FakeSandbox(), model=None)) + + +# --------------------------------------------------------------------------- turn response + + +def test_turn_response_failed_raises(): + _, state = parse_all("turn_failed.jsonl") + with pytest.raises(HarnessTurnError, match="no healthy deployments"): + CodexHarnessConfig().transform_turn_response( + make_ctx(FakeSandbox()), state, 1, [] + ) + + +def test_turn_response_nonzero_exit_uses_stderr_tail(): + with pytest.raises(HarnessTurnError, match=r"code 1: Error loading config\.toml"): + CodexHarnessConfig().transform_turn_response( + make_ctx(FakeSandbox()), + CodexStreamState(), + 1, + ["Error loading config.toml: bad", ""], + ) + with pytest.raises(HarnessTurnError, match="code 2: no output"): + CodexHarnessConfig().transform_turn_response( + make_ctx(FakeSandbox()), CodexStreamState(), 2, [] + ) + + +def test_turn_response_output_json_only_with_output(): + state = CodexStreamState(final_text='{"file": "a", "content": "b"}') + cfg = CodexHarnessConfig() + with_out = cfg.transform_turn_response( + make_ctx(FakeSandbox(), output=Answer), state, 0, [] + ) + assert with_out.output_json == state.final_text + without = cfg.transform_turn_response(make_ctx(FakeSandbox()), state, 0, []) + assert without.output_json is None and without.final_text == state.final_text + + +# --------------------------------------------------------------------------- handler + + +async def test_start_missing_binary(): + ctx = make_ctx(FakeSandbox(has_codex=False)) + with pytest.raises(HarnessInstallFailed, match="codex"): + await make_handler().start(ctx) + + +async def test_start_rejects_managed_config_and_wrong_options(): + with pytest.raises(OptionsMismatch): + await make_handler().start( + make_ctx( + FakeSandbox(), options=CodexOptions(config={"model_provider": "openai"}) + ) + ) + with pytest.raises(OptionsMismatch): + await make_handler().start(make_ctx(FakeSandbox(), options=ClaudeCodeOptions())) + + +async def test_start_writes_skills_and_schema(tmp_path): + skill = tmp_path / "my-skill" + (skill / "scripts").mkdir(parents=True) + (skill / "SKILL.md").write_text("---\nname: my-skill\n---\nbody") + (skill / "scripts" / "run.sh").write_text("echo hi") + sbx = FakeSandbox() + await make_handler().start(make_ctx(sbx, skills=[str(skill)], output=Answer)) + assert sbx.files[f"{HOME}/skills/my-skill/SKILL.md"].startswith(b"---") + assert sbx.files[f"{HOME}/skills/my-skill/scripts/run.sh"] == b"echo hi" + schema = json.loads(sbx.files[f"{HOME}/output_schema.json"]) + assert schema["additionalProperties"] is False + assert schema["required"] == ["file", "content"] + assert sbx.runs[0][:2] == ["sh", "-c"] + assert sbx.runs[0][-2:] == [f"{HOME}/sessions", "codex/sessions"] + + +async def test_first_turn_then_resume_argv_env(): + sbx = FakeSandbox( + outputs=[ + fixture_proc("turn1_bash.jsonl"), + fixture_proc("turn2_resume_apply_patch.jsonl"), + ] + ) + handler = make_handler() + ctx = make_ctx( + sbx, + instructions="Be terse.", + options=CodexOptions( + reasoning_effort="low", + config={"sandbox_workspace_write.network_access": True}, + ), + ) + await handler.start(ctx) + events = await collect(handler, ctx, "create hello.txt") + + first = sbx.execs[0] + argv, env = first["cmd"], first["env"] + assert argv[:4] == ["codex", "exec", "--json", "--skip-git-repo-check"] + assert argv[-1] == "-" and argv[argv.index("-C") + 1] == "/work" + assert first["cwd"] == "/work" + assert "model_provider=litellm" in config_values(argv) + assert not any(TOKEN in a for a in argv) + assert env["LITELLM_HARNESS_TOKEN"] == TOKEN + assert env["CODEX_HOME"] == HOME + assert sbx.processes[0].stdin.data == b"create hello.txt" + assert sbx.processes[0].stdin.closed + + assert isinstance(events[-1], Text) + assert ctx.final_text.startswith("Done") + assert handler.native_session_id() == THREAD_ID + + await collect(handler, ctx, "edit it") + argv2 = sbx.execs[1]["cmd"] + assert argv2[:4] == ["codex", "exec", "resume", THREAD_ID] + assert "--sandbox" not in argv2 and "-C" not in argv2 + assert 'sandbox_mode="workspace-write"' in config_values(argv2) + assert ctx.final_text == "done" + + +async def test_resume_sets_thread_id(): + sbx = FakeSandbox(outputs=[fixture_proc("reasoning.jsonl")]) + handler = make_handler() + ctx = make_ctx(sbx) + await handler.start(ctx) + await handler.resume(ctx, "thread-9") + assert handler.native_session_id() == "thread-9" + await collect(handler, ctx, "again") + assert sbx.execs[0]["cmd"][:4] == ["codex", "exec", "resume", "thread-9"] + + +async def test_structured_output_sets_output_json(): + sbx = FakeSandbox(outputs=[fixture_proc("structured_output.jsonl")]) + handler = make_handler() + ctx = make_ctx(sbx, output=Answer, permissions="read-only") + await handler.start(ctx) + await collect(handler, ctx, "read hello.txt") + argv = sbx.execs[0]["cmd"] + assert argv[argv.index("--output-schema") + 1] == f"{HOME}/output_schema.json" + assert argv[argv.index("--sandbox") + 1] == "read-only" + assert Answer.model_validate_json(ctx.output_json).file == "hello.txt" + + +async def test_turn_failed_raises(): + sbx = FakeSandbox(outputs=[fixture_proc("turn_failed.jsonl", exit_code=1)]) + handler = make_handler() + ctx = make_ctx(sbx) + await handler.start(ctx) + with pytest.raises(HarnessTurnError, match="no healthy deployments"): + await collect(handler, ctx, "hi") + + +async def test_nonzero_exit_raises_with_stderr_tail(): + sbx = FakeSandbox( + outputs=[ + FakeProcess(b"", stderr=b"Error loading config.toml: bad\n", exit_code=1) + ] + ) + handler = make_handler() + ctx = make_ctx(sbx) + await handler.start(ctx) + with pytest.raises(HarnessTurnError, match=r"code 1: Error loading config\.toml"): + await collect(handler, ctx, "hi") + + +async def test_early_close_kills_process_and_stop_is_idempotent(): + sbx = FakeSandbox(outputs=[fixture_proc("turn1_bash.jsonl")]) + handler = make_handler() + ctx = make_ctx(sbx) + await handler.start(ctx) + gen = handler.turn(ctx, "hi") + await gen.__anext__() + await gen.aclose() + assert sbx.processes[0].killed + await handler.stop(ctx) + await handler.stop(ctx) + + +async def test_long_jsonl_line_is_parsed(): + text = "x" * 200_000 + line = json.dumps( + { + "type": "item.completed", + "item": {"id": "a", "type": "agent_message", "text": text}, + } + ) + sbx = FakeSandbox(outputs=[FakeProcess(line.encode() + b"\n")]) + handler = make_handler() + ctx = make_ctx(sbx) + await handler.start(ctx) + events = await collect(handler, ctx, "hi") + assert events == [Text(delta=text)] + + +def test_capabilities(): + cfg = CodexHarnessConfig() + caps = cfg.capabilities + assert cfg.harness is Harness.CODEX + assert cfg.options_type is CodexOptions + assert cfg.get_binary() == "codex" + assert "@openai/codex" in cfg.get_install_hint() + assert caps.structured_output and caps.skills and caps.resume + assert not ( + caps.tool_approval or caps.tool_filtering or caps.custom_tools or caps.history + ) + assert caps.permission_modes == frozenset({"read-only", "full"}) diff --git a/tests/unit/llms/deepagents/__init__.py b/tests/unit/llms/deepagents/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepagents/harness/__init__.py b/tests/unit/llms/deepagents/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/deepagents/harness/test_sandbox_backend_symlinks.py b/tests/unit/llms/deepagents/harness/test_sandbox_backend_symlinks.py new file mode 100644 index 00000000000..022771400e5 --- /dev/null +++ b/tests/unit/llms/deepagents/harness/test_sandbox_backend_symlinks.py @@ -0,0 +1,78 @@ +"""A repository must not be able to reach host files through symlinks, in any file tool.""" + +import asyncio +import os +from pathlib import Path + +import pytest + +from litellm.harness.sandbox.local import LocalSandbox + +backend = pytest.importorskip("litellm.llms.deepagents.harness.sandbox_backend") + +SECRET = "AWS_SECRET_ACCESS_KEY=leaked-from-host" + + +@pytest.fixture +def repo_with_escape_links(tmp_path: Path) -> Path: + host = tmp_path / "host_home" + host.mkdir() + (host / "credentials").write_text(SECRET + "\n") + repo = tmp_path / "repo" + repo.mkdir() + (repo / "README.md").write_text("hello\n") + os.symlink(host / "credentials", repo / "creds_link") + os.symlink(host, repo / "home_link") + return repo + + +async def _backend(repo: Path) -> object: + return backend.SandboxBackend(LocalSandbox(str(repo)), loop=asyncio.get_running_loop(), writable=False) + + +async def test_grep_whole_repo_skips_symlinks_out_of_workspace(repo_with_escape_links: Path) -> None: + b = await _backend(repo_with_escape_links) + result = await b.agrep("AWS_SECRET") + assert not result.matches, f"grep followed a symlink out of the repo: {result}" + + +async def test_grep_rooted_at_symlink_dir_is_refused(repo_with_escape_links: Path) -> None: + b = await _backend(repo_with_escape_links) + result = await b.agrep("AWS_SECRET", path="/home_link") + assert result.error and "outside the workspace" in result.error + assert not result.matches + + +async def test_read_through_symlink_is_refused(repo_with_escape_links: Path) -> None: + b = await _backend(repo_with_escape_links) + result = await b.aread("/creds_link") + assert result.error and "outside the workspace" in result.error + assert SECRET not in str(result.file_data) + + +async def test_glob_does_not_list_files_behind_symlinks(repo_with_escape_links: Path) -> None: + b = await _backend(repo_with_escape_links) + result = await b.aglob("**/*") + paths = [m["path"] for m in result.matches or []] + assert paths == ["/README.md"] + + +async def test_grep_still_finds_real_repo_files(repo_with_escape_links: Path) -> None: + b = await _backend(repo_with_escape_links) + result = await b.agrep("hello") + assert [(m["path"], m["line"]) for m in result.matches] == [("/README.md", 1)] + + +async def test_write_into_new_nested_directory_is_allowed(tmp_path: Path) -> None: + repo = tmp_path / "repo" + repo.mkdir() + b = backend.SandboxBackend(LocalSandbox(str(repo)), loop=asyncio.get_running_loop(), writable=True) + result = await b.awrite("/new_dir/sub/file.py", "print('hi')\n") + assert result.error is None, result.error + assert (repo / "new_dir" / "sub" / "file.py").read_text() == "print('hi')\n" + + +async def test_write_under_symlinked_dir_is_refused(repo_with_escape_links: Path) -> None: + b = backend.SandboxBackend(LocalSandbox(str(repo_with_escape_links)), loop=asyncio.get_running_loop(), writable=True) + result = await b.awrite("/home_link/new_dir/evil.txt", "x") + assert result.error and "outside the workspace" in result.error diff --git a/tests/unit/llms/deepagents/harness/test_transformation.py b/tests/unit/llms/deepagents/harness/test_transformation.py new file mode 100644 index 00000000000..4409cead299 --- /dev/null +++ b/tests/unit/llms/deepagents/harness/test_transformation.py @@ -0,0 +1,185 @@ +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +from pydantic import BaseModel + +from litellm.harness.context import GatewayTarget, SessionContext +from litellm.harness.errors import OptionsMismatch +from litellm.harness.options import CodexOptions, DeepAgentsOptions +from litellm.harness.sandbox.local import LocalSandbox +from litellm.harness.types import Harness, Reasoning, Text, ToolCall, ToolResult +from litellm.llms.deepagents.harness import transformation as da + + +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 msg(kind: str, **fields: Any) -> SimpleNamespace: + return SimpleNamespace(type=kind, **fields) + + +def test_blocked_tools_modes() -> None: + assert da.blocked_tools("full", []) == frozenset() + assert da.blocked_tools("edit", []) == frozenset({"execute"}) + assert {"write_file", "edit_file", "delete", "execute"} <= da.blocked_tools( + "read-only", [] + ) + assert da.blocked_tools("full", ["read", "ls"]) == frozenset({"read_file", "ls"}) + assert da.blocked_tools("full", ["bash", "grep"]) == frozenset({"execute", "grep"}) + + +def test_interrupt_config_only_for_ask() -> None: + assert da.interrupt_config("full", frozenset()) is None + config = da.interrupt_config("ask", frozenset({"execute"})) + assert set(config) == {"write_file", "edit_file", "delete"} + assert all( + v == {"allowed_decisions": ["approve", "reject"]} for v in config.values() + ) + + +def test_normalized_tool_name() -> None: + assert da.normalized_tool_name("write_file") == "write" + assert da.normalized_tool_name("read_file") == "read" + assert da.normalized_tool_name("edit_file") == "edit" + assert da.normalized_tool_name("execute") == "bash" + assert da.normalized_tool_name("add") == "add" + + +def test_chat_model_kwargs_gateway_and_sdk(tmp_path: Path) -> None: + gw = GatewayTarget(api_base="https://gw.example.com", api_key="sk-virtual") + ctx = make_ctx(tmp_path, gateway=gw, metadata={"team": "a"}) + kwargs = da.chat_model_kwargs(ctx) + assert kwargs["model"] == "litellm_proxy/gpt-4o-mini" + assert kwargs["api_base"] == "https://gw.example.com" + assert kwargs["api_key"] == "sk-virtual" + assert kwargs["extra_headers"]["x-litellm-tags"] == "harness,deepagents" + assert '"team": "a"' in kwargs["extra_headers"]["x-litellm-spend-logs-metadata"] + no_meta = da.chat_model_kwargs(make_ctx(tmp_path, gateway=gw)) + assert "x-litellm-spend-logs-metadata" not in no_meta["extra_headers"] + + sdk = da.chat_model_kwargs(make_ctx(tmp_path, api_key="k", api_base="http://b")) + assert sdk == {"model": "gpt-4o-mini", "api_key": "k", "api_base": "http://b"} + with pytest.raises(ValueError, match="needs model="): + da.chat_model_kwargs(make_ctx(tmp_path, model=None)) + + +def test_recursion_limit(tmp_path: Path) -> None: + assert ( + da.recursion_limit(make_ctx(tmp_path)) == da.DEEPAGENTS_DEFAULT_RECURSION_LIMIT + ) + assert da.recursion_limit(make_ctx(tmp_path, max_turns=2)) == ( + da.DEEPAGENTS_BASE_RECURSION_LIMIT + 2 * da.DEEPAGENTS_STEPS_PER_TURN + ) + opts = DeepAgentsOptions(recursion_limit=7) + assert da.recursion_limit(make_ctx(tmp_path, max_turns=2, options=opts)) == 7 + + +def test_stream_events_text_and_reasoning() -> None: + assert da.stream_events(msg("human", content="hi")) == [] + events = da.stream_events( + msg( + "AIMessageChunk", + content=[ + {"type": "thinking", "thinking": "hmm"}, + {"type": "text", "text": "a"}, + "b", + ], + additional_kwargs={}, + ) + ) + assert events == [Reasoning(delta="hmm"), Text(delta="ab")] + extra = da.stream_events( + msg("ai", content="x", additional_kwargs={"reasoning_content": "r"}) + ) + assert extra == [Reasoning(delta="r"), Text(delta="x")] + + +def test_update_events_tool_calls_results_and_skip() -> None: + ai = msg( + "ai", + tool_calls=[ + {"name": "write_file", "args": {"file_path": "/a"}, "id": "c1"}, + {"name": "Answer", "args": {"city": "Paris"}, "id": "c2"}, + {"name": "add", "args": None, "id": "c3"}, + ], + ) + tool = msg("tool", name="write_file", tool_call_id="c1", content="ok", status=None) + err = msg("tool", name="execute", tool_call_id="c4", content="x", status="error") + skipped = msg("tool", name="Answer", tool_call_id="c2", content="", status=None) + update = { + "model": {"messages": [ai]}, + "tools": {"messages": [tool, err, skipped]}, + "SomeMiddleware.after_model": {"messages": [ai]}, + } + events = da.update_events(update, frozenset({"Answer"})) + assert events == [ + ToolCall( + id="c1", + name="write", + native_name="write_file", + input={"file_path": "/a"}, + builtin=True, + ), + ToolCall( + id="c3", name="add", native_name="add", input={"args": None}, builtin=False + ), + ToolResult(id="c1", output="ok", is_error=False), + ToolResult(id="c4", output="x", is_error=True), + ] + assert da.update_events(None, frozenset()) == [] + assert da.update_events({"model": None}, frozenset()) == [] + + +def test_interrupts_and_approval_requests() -> None: + assert da.interrupts_in({"__interrupt__": ("i",)}) == ["i"] + assert da.interrupts_in({}) == [] and da.interrupts_in(None) == [] + value = {"action_requests": [{"name": "write_file", "args": {}}, "junk"]} + assert da.approval_requests(value) == [{"name": "write_file", "args": {}}] + assert da.approval_requests(None) == [] + assert da.approval_requests({"action_requests": "x"}) == [] + + +def test_decision() -> None: + assert da.decision(True, "") == {"type": "approve"} + assert da.decision(False, "no") == {"type": "reject", "message": "no"} + assert da.decision(False, "")["message"] + + +class Answer(BaseModel): + city: str + + +def test_final_ai_text_and_structured_json() -> None: + messages = [ + msg("ai", content="first"), + msg("tool", content="t"), + msg("ai", content=""), + ] + assert da.final_ai_text(messages) == "first" + assert da.final_ai_text([]) == "" + assert da.structured_json(None) is None + assert Answer.model_validate_json(da.structured_json(Answer(city="Paris"))) + assert da.structured_json({"city": "Paris"}) == '{"city": "Paris"}' + + +def test_config_capabilities_and_validation(tmp_path: Path) -> None: + config = da.DeepAgentsHarnessConfig() + assert config.uses_model_endpoint is False + assert config.capabilities.tool_approval and config.capabilities.history + assert "ask" in config.capabilities.permission_modes + config.validate_environment(make_ctx(tmp_path)) + with pytest.raises(ValueError, match="needs model="): + config.validate_environment(make_ctx(tmp_path, model=None)) + with pytest.raises(OptionsMismatch): + config.validate_environment(make_ctx(tmp_path, options=CodexOptions())) + assert "pip install deepagents langchain-litellm" in da.INSTALL_HINT diff --git a/tests/unit/llms/opencode/__init__.py b/tests/unit/llms/opencode/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opencode/harness/__init__.py b/tests/unit/llms/opencode/harness/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opencode/harness/fixtures/__init__.py b/tests/unit/llms/opencode/harness/fixtures/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opencode/harness/fixtures/api_error.jsonl b/tests/unit/llms/opencode/harness/fixtures/api_error.jsonl new file mode 100644 index 00000000000..b2b3148285e --- /dev/null +++ b/tests/unit/llms/opencode/harness/fixtures/api_error.jsonl @@ -0,0 +1 @@ +{"type":"error","timestamp":1790788205744,"sessionID":"ses_f0cb48565ffeMWhVl1J584kSti","error":{"name":"APIError","data":{"message":"litellm.BadRequestError: You passed in model=no-such-model-xyz. There are no healthy deployments for this model","statusCode":400,"isRetryable":false}}} diff --git a/tests/unit/llms/opencode/harness/fixtures/endpoint_requests.jsonl b/tests/unit/llms/opencode/harness/fixtures/endpoint_requests.jsonl new file mode 100644 index 00000000000..0b4c3d22bf4 --- /dev/null +++ b/tests/unit/llms/opencode/harness/fixtures/endpoint_requests.jsonl @@ -0,0 +1,4 @@ +{"method":"POST","path":"/v1/chat/completions","headers":{"Content-Type":"application/json","User-Agent":"opencode/1.14.41 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.13","x-session-affinity":"ses_f0cb977cdffeoCMeplOiw1KY25","Connection":"keep-alive","Accept":"*/*"},"auth_prefix":"Bearer ","body_keys":["max_tokens","messages","model","stream","stream_options"],"model":"claude-haiku-4-5-20251001","stream":true,"stream_options":{"include_usage":true},"tools":[],"n_messages":3,"roles":["system","user","user"]} +{"method":"POST","path":"/v1/chat/completions","headers":{"Content-Type":"application/json","User-Agent":"opencode/1.14.41 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.13","x-session-affinity":"ses_f0cb977cdffeoCMeplOiw1KY25","Connection":"keep-alive","Accept":"*/*"},"auth_prefix":"Bearer ","body_keys":["max_tokens","messages","model","stream","stream_options","tool_choice","tools"],"model":"claude-haiku-4-5-20251001","stream":true,"stream_options":{"include_usage":true},"tools":["bash","read","glob","grep","edit","write","task","webfetch","todowrite","skill"],"n_messages":2,"roles":["system","user"]} +{"method":"POST","path":"/v1/chat/completions","headers":{"Content-Type":"application/json","User-Agent":"opencode/1.14.41 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.13","x-session-affinity":"ses_f0cb977cdffeoCMeplOiw1KY25","Connection":"keep-alive","Accept":"*/*"},"auth_prefix":"Bearer ","body_keys":["max_tokens","messages","model","stream","stream_options","tool_choice","tools"],"model":"claude-haiku-4-5-20251001","stream":true,"stream_options":{"include_usage":true},"tools":["bash","read","glob","grep","edit","write","task","webfetch","todowrite","skill"],"n_messages":4,"roles":["system","user","assistant","tool"]} +{"method":"POST","path":"/v1/chat/completions","headers":{"Content-Type":"application/json","User-Agent":"opencode/1.14.41 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.13","x-session-affinity":"ses_f0cb977cdffeoCMeplOiw1KY25","Connection":"keep-alive","Accept":"*/*"},"auth_prefix":"Bearer ","body_keys":["max_tokens","messages","model","stream","stream_options","tool_choice","tools"],"model":"claude-haiku-4-5-20251001","stream":true,"stream_options":{"include_usage":true},"tools":["bash","read","glob","grep","edit","write","task","webfetch","todowrite","skill"],"n_messages":6,"roles":["system","user","assistant","tool","assistant","tool"]} diff --git a/tests/unit/llms/opencode/harness/fixtures/readonly_denied_bash.jsonl b/tests/unit/llms/opencode/harness/fixtures/readonly_denied_bash.jsonl new file mode 100644 index 00000000000..e31cb95503c --- /dev/null +++ b/tests/unit/llms/opencode/harness/fixtures/readonly_denied_bash.jsonl @@ -0,0 +1,7 @@ +{"type":"step_start","timestamp":1790787993225,"sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","part":{"id":"prt_0f3483e81001WZXwo1sKkhgOIo","messageID":"msg_0f3483ac90012uFburOGJS91gV","sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","type":"step-start"}} +{"type":"tool_use","timestamp":1790787993571,"sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","part":{"type":"tool","tool":"invalid","callID":"toolu_0126GuE9NKXyF4HoXLY3rREx","state":{"status":"completed","input":{"tool":"bash","error":"Model tried to call unavailable tool 'bash'. Available tools: invalid, read, glob, grep, task, todowrite, skill."},"output":"The arguments provided to the tool are invalid: Model tried to call unavailable tool 'bash'. Available tools: invalid, read, glob, grep, task, todowrite, skill.","metadata":{"truncated":false},"title":"Invalid Tool","time":{"start":1790787993567,"end":1790787993570}},"id":"prt_0f3483f5e001tfzTsZE0IrrCxH","sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","messageID":"msg_0f3483ac90012uFburOGJS91gV"}} +{"type":"text","timestamp":1790787993573,"sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","part":{"id":"prt_0f3483e85001y7vW8x70lZJBwu","messageID":"msg_0f3483ac90012uFburOGJS91gV","sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","type":"text","text":"I'll run that shell command for you using bash.","time":{"start":1790787993221,"end":1790787993572}}} +{"type":"step_finish","timestamp":1790787993575,"sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","part":{"id":"prt_0f3483fe5001VEbV0Sy9d0ybb0","reason":"tool-calls","messageID":"msg_0f3483ac90012uFburOGJS91gV","sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","type":"step-finish","tokens":{"total":7484,"input":7416,"output":68,"reasoning":0,"cache":{"write":0,"read":0}},"cost":0}} +{"type":"step_start","timestamp":1790787994554,"sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","part":{"id":"prt_0f34843b6001QhtCAWotsESvRK","messageID":"msg_0f3483fea0021zWRKRkh32701g","sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","type":"step-start"}} +{"type":"text","timestamp":1790787994748,"sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","part":{"id":"prt_0f34843b900143p3b7263tu3pc","messageID":"msg_0f3483fea0021zWRKRkh32701g","sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","type":"text","text":"FAILED\n\nThe bash tool is not available in this environment, so I cannot execute that shell command.","time":{"start":1790787994553,"end":1790787994747}}} +{"type":"step_finish","timestamp":1790787994749,"sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","part":{"id":"prt_0f348447c001hnUGx5uag7BAYm","reason":"stop","messageID":"msg_0f3483fea0021zWRKRkh32701g","sessionID":"ses_f0cb7c593ffesLm4De78bjOZjn","type":"step-finish","tokens":{"total":7596,"input":158,"output":24,"reasoning":0,"cache":{"write":0,"read":7414}},"cost":0}} diff --git a/tests/unit/llms/opencode/harness/fixtures/turn1_write_read.jsonl b/tests/unit/llms/opencode/harness/fixtures/turn1_write_read.jsonl new file mode 100644 index 00000000000..83eab9e49a8 --- /dev/null +++ b/tests/unit/llms/opencode/harness/fixtures/turn1_write_read.jsonl @@ -0,0 +1,11 @@ +{"type":"step_start","timestamp":1790787882230,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f3468cf4001Lp09mwOdY2Nk0v","messageID":"msg_0f3468893001gnm1GTIsdViRXS","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-start"}} +{"type":"text","timestamp":1790787882762,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f3468cf5001EnUgA050AAqsKb","messageID":"msg_0f3468893001gnm1GTIsdViRXS","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"text","text":"I'll create a hello.txt file containing \"hi\" and then read it.","time":{"start":1790787882229,"end":1790787882762}}} +{"type":"tool_use","timestamp":1790787882770,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"type":"tool","tool":"write","callID":"toolu_015FFUwEf2dazoWfCrMbMCnm","state":{"status":"completed","input":{"filePath":"/workspace/hello.txt","content":"hi"},"output":"Wrote file successfully.","metadata":{"diagnostics":{},"filepath":"/workspace/hello.txt","exists":false,"truncated":false},"title":"private/workspace/hello.txt","time":{"start":1790787882760,"end":1790787882768}},"id":"prt_0f3468dcf001rxCR4QNHtzY272","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","messageID":"msg_0f3468893001gnm1GTIsdViRXS"}} +{"type":"step_finish","timestamp":1790787882770,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f3468f11001woc6vGYjk4ErHI","reason":"tool-calls","messageID":"msg_0f3468893001gnm1GTIsdViRXS","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-finish","tokens":{"total":11018,"input":10918,"output":100,"reasoning":0,"cache":{"write":0,"read":0}},"cost":0}} +{"type":"step_start","timestamp":1790787905598,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f346e8390015Ikck6AJPLpFr0","messageID":"msg_0f3468f14001j9Au0eVNVNWev1","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-start"}} +{"type":"text","timestamp":1790787905940,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f346e83d001bMWqsgMxNP5FGG","messageID":"msg_0f3468f14001j9Au0eVNVNWev1","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"text","text":"Now let me read the file:","time":{"start":1790787905597,"end":1790787905940}}} +{"type":"tool_use","timestamp":1790787905949,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"type":"tool","tool":"read","callID":"toolu_017766NprB4499fkNFoqLx6k","state":{"status":"completed","input":{"filePath":"/workspace/hello.txt"},"output":"/workspace/hello.txt\nfile\n\n1: hi\n\n(End of file - total 1 lines)\n","metadata":{"preview":"hi","truncated":false,"loaded":[]},"title":"private/workspace/hello.txt","time":{"start":1790787905937,"end":1790787905947}},"id":"prt_0f346e8be001TkCzxNO3iCWyHh","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","messageID":"msg_0f3468f14001j9Au0eVNVNWev1"}} +{"type":"step_finish","timestamp":1790787905949,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f346e99c0018zQ7l02y6UXvHZ","reason":"tool-calls","messageID":"msg_0f3468f14001j9Au0eVNVNWev1","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-finish","tokens":{"total":11107,"input":118,"output":73,"reasoning":0,"cache":{"write":0,"read":10916}},"cost":0}} +{"type":"step_start","timestamp":1790787907026,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f346edd1001UEsH2Hitcy599I","messageID":"msg_0f346e99f0010o3y8k52od8t3I","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-start"}} +{"type":"text","timestamp":1790787907596,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f346edd10036kWCbbi5l5qxja","messageID":"msg_0f346e99f0010o3y8k52od8t3I","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"text","text":"Done! I've created hello.txt with the content \"hi\" and read it back. The file contains exactly what was requested.","time":{"start":1790787907025,"end":1790787907595}}} +{"type":"step_finish","timestamp":1790787907597,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f346f00c002zjKza4p4eIg1Ox","reason":"stop","messageID":"msg_0f346e99f0010o3y8k52od8t3I","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-finish","tokens":{"total":11198,"input":139,"output":30,"reasoning":0,"cache":{"write":0,"read":11029}},"cost":0}} diff --git a/tests/unit/llms/opencode/harness/fixtures/turn2_session_skill.jsonl b/tests/unit/llms/opencode/harness/fixtures/turn2_session_skill.jsonl new file mode 100644 index 00000000000..76d40e08c99 --- /dev/null +++ b/tests/unit/llms/opencode/harness/fixtures/turn2_session_skill.jsonl @@ -0,0 +1,7 @@ +{"type":"step_start","timestamp":1790787953184,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f347a21e001EmD408wvQkLCan","messageID":"msg_0f3479cf2001JlcVj8lSmhqAum","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-start"}} +{"type":"text","timestamp":1790787953736,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f347a21f001Bu2RKXM9urS1v9","messageID":"msg_0f3479cf2001JlcVj8lSmhqAum","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"text","text":"I created `/workspace/hello.txt`.\n\nFor the secret greeting word, let me load the greeter skill:","time":{"start":1790787953183,"end":1790787953736}}} +{"type":"tool_use","timestamp":1790787953752,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"type":"tool","tool":"skill","callID":"toolu_01Q1ip1d6snmazhbYCaEqcsT","state":{"status":"completed","input":{"name":"greeter"},"output":"\n# Skill: greeter\n\nThe secret greeting word is PINEAPPLE.\n\nBase directory for this skill: file:///tmp/litellm-harness-x/skills/greeter\nRelative paths in this skill (e.g., scripts/, reference/) are relative to this base directory.\nNote: file list is sampled.\n\n\n\n\n","metadata":{"name":"greeter","dir":"/tmp/litellm-harness-x/skills/greeter","truncated":false},"title":"Loaded skill: greeter","time":{"start":1790787953733,"end":1790787953749}},"id":"prt_0f347a3ab001knlkXMZcS3jV9G","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","messageID":"msg_0f3479cf2001JlcVj8lSmhqAum"}} +{"type":"step_finish","timestamp":1790787953752,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f347a457001JH8PaFwNLIRfA9","reason":"tool-calls","messageID":"msg_0f3479cf2001JlcVj8lSmhqAum","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-finish","tokens":{"total":11531,"input":11446,"output":85,"reasoning":0,"cache":{"write":0,"read":0}},"cost":0}} +{"type":"step_start","timestamp":1790787970574,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f347e60c001CUPeXPcH733B3G","messageID":"msg_0f347a45b001XQWc2pM6IPZrsO","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-start"}} +{"type":"text","timestamp":1790787970675,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f347e60d001pQ2aAgMwHkYRx3","messageID":"msg_0f347a45b001XQWc2pM6IPZrsO","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"text","text":"The secret greeting word is **PINEAPPLE**.","time":{"start":1790787970573,"end":1790787970674}}} +{"type":"step_finish","timestamp":1790787970676,"sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","part":{"id":"prt_0f347e673002k2mm651nyRtr07","reason":"stop","messageID":"msg_0f347a45b001XQWc2pM6IPZrsO","sessionID":"ses_f0cb977cdffeoCMeplOiw1KY25","type":"step-finish","tokens":{"total":11663,"input":204,"output":15,"reasoning":0,"cache":{"write":0,"read":11444}},"cost":0}} diff --git a/tests/unit/llms/opencode/harness/test_transformation.py b/tests/unit/llms/opencode/harness/test_transformation.py new file mode 100644 index 00000000000..2caf1ef80b8 --- /dev/null +++ b/tests/unit/llms/opencode/harness/test_transformation.py @@ -0,0 +1,726 @@ +import asyncio +import json +from dataclasses import dataclass, field +from pathlib import Path + +import pytest +from pydantic import BaseModel + +from litellm.harness.context import SessionContext +from litellm.harness.errors import ( + CapabilityUnsupported, + HarnessError, + HarnessInstallFailed, + OptionsMismatch, +) +from litellm.harness.handlers.cli_handler import PERSIST_DIR_SCRIPT, CLIHarnessHandler +from litellm.harness.options import CodexOptions, OpenCodeOptions +from litellm.harness.sandbox.base import CompletedRun +from litellm.harness.types import Harness, Reasoning, Text, ToolCall, ToolResult +from litellm.llms.base_llm.harness.transformation import ( + HarnessSessionSetup, + HarnessTurnError, +) +from litellm.llms.opencode.harness.transformation import ( + INSTRUCTIONS_FILENAME, + OPENCODE_ISOLATION_ENV, + OPENCODE_SESSION_TITLE, + TOKEN_FILENAME, + XDG_DIRNAME, + OpenCodeHarnessConfig, + OpenCodeStreamState, + build_instructions, + build_opencode_config, + permission_rules, + turn_prompt, + validate_user_config, +) + +FIXTURES = Path(__file__).parent / "fixtures" +TOKEN = "tok-secret-123" +SESSION = "ses_f0cb977cdffeoCMeplOiw1KY25" +PRIVATE = "/tmp/oc-1" +CONFIG = OpenCodeHarnessConfig() + + +def load_fixture(name: str) -> list[dict]: + return [ + json.loads(line) for line in (FIXTURES / name).read_text().splitlines() if line + ] + + +def parse(obj: dict, state: OpenCodeStreamState) -> list: + return CONFIG.transform_stream_line(obj, state) + + +def parse_all(name: str, state: OpenCodeStreamState | None = None): + state = state or CONFIG.create_stream_state() + events = [] + for obj in load_fixture(name): + events.extend(parse(obj, state)) + return events, state + + +# --------------------------------------------------------------------------- fakes + + +class FakeStdin: + def __init__(self): + 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 FakeProcess: + def __init__(self, stdout: bytes, stderr: bytes = b"", exit_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._exit_code = exit_code + self.killed = False + + async def wait(self) -> int: + return self._exit_code + + async def kill(self) -> None: + self.killed = True + + +@dataclass +class FakeSandbox: + workdir: str = "/work" + has_binary: bool = True + persist_ok: bool = True + outputs: list = field(default_factory=list) + files: dict = field(default_factory=dict) + execs: list = field(default_factory=list) + runs: list = field(default_factory=list) + tempdirs: int = 0 + + async def exec(self, cmd, *, env=None, cwd=None): + self.execs.append({"cmd": cmd, "env": dict(env or {}), "cwd": cwd}) + return self.outputs.pop(0) + + async def run(self, cmd, *, env=None, cwd=None, timeout=None): + self.runs.append(cmd) + if self.persist_ok: + return CompletedRun("", "", 0) + return CompletedRun("", "read-only fs", 1) + + async def read(self, path): + return self.files[path] + + async def write(self, path, data): + self.files[path] = data + + def host_url(self, port): + return f"http://host.docker.internal:{port}" + + async def which(self, binary): + return f"/usr/bin/{binary}" if self.has_binary else None + + async def tempdir(self): + self.tempdirs += 1 + return f"/tmp/oc-{self.tempdirs}" + + async def snapshot(self): + return {} + + async def close(self): + return None + + +@dataclass +class FakeEndpoint: + port: int = 4555 + token: str = TOKEN + model: str | None = None + + +class Answer(BaseModel): + file: str + content: str + + +def make_ctx(sandbox=None, **kwargs) -> SessionContext: + return SessionContext( + harness=Harness.OPENCODE, + sandbox=sandbox or FakeSandbox(), + session_id="s1", + model=kwargs.pop("model", "claude-haiku-4-5-20251001"), + endpoint=kwargs.pop("endpoint", FakeEndpoint()), + **kwargs, + ) + + +def setup_for(ctx: SessionContext) -> HarnessSessionSetup: + return CONFIG.transform_session_setup(ctx, PRIVATE) + + +def setup_config(setup: HarnessSessionSetup) -> dict: + return json.loads(setup.env["OPENCODE_CONFIG_CONTENT"]) + + +def fixture_proc(name: str, **kwargs) -> FakeProcess: + return FakeProcess((FIXTURES / name).read_bytes(), **kwargs) + + +async def collect(handler, ctx, prompt): + return [e async for e in handler.turn(ctx, prompt)] + + +async def started(sandbox=None, **kwargs): + sandbox = sandbox or FakeSandbox() + handler = CLIHarnessHandler(OpenCodeHarnessConfig()) + ctx = make_ctx(sandbox, **kwargs) + await handler.start(ctx) + return handler, ctx, sandbox + + +def exec_config(sandbox, index=0) -> dict: + return json.loads(sandbox.execs[index]["env"]["OPENCODE_CONFIG_CONTENT"]) + + +# --------------------------------------------------------------------------- parsing + + +def test_parse_write_read_turn(): + events, state = parse_all("turn1_write_read.jsonl") + assert CONFIG.get_native_session_id(state) == SESSION + assert [type(e) for e in events] == [ + Text, + ToolCall, + ToolResult, + Text, + ToolCall, + ToolResult, + Text, + ] + write, write_result = events[1], events[2] + assert write.name == "write" and write.native_name == "write" + assert write.builtin is True + assert write.input == {"filePath": "/workspace/hello.txt", "content": "hi"} + assert write_result.id == write.id == "toolu_015FFUwEf2dazoWfCrMbMCnm" + assert write_result.is_error is False + read, read_result = events[4], events[5] + assert read.name == "read" and "1: hi" in read_result.output + assert state.final_text.startswith("Done!") + assert state.error is None + + +def test_parse_skill_tool_on_continued_session(): + events, state = parse_all("turn2_session_skill.jsonl") + assert state.session_id == SESSION + skill = next(e for e in events if isinstance(e, ToolCall)) + assert skill.name == "skill" and skill.input == {"name": "greeter"} + assert skill.builtin is True + assert "PINEAPPLE" in state.final_text + + +def test_parse_denied_tool_is_error_result(): + events, state = parse_all("readonly_denied_bash.jsonl") + call = next(e for e in events if isinstance(e, ToolCall)) + result = next(e for e in events if isinstance(e, ToolResult)) + assert call.native_name == "invalid" and call.input["tool"] == "bash" + assert result.is_error is True + assert state.final_text.startswith("FAILED") + + +def test_parse_api_error_records_error(): + events, state = parse_all("api_error.jsonl") + assert events == [] + assert "no healthy deployments" in state.error + + +def test_parse_reasoning_and_tool_error_and_name_mapping(): + state = OpenCodeStreamState() + reasoning = parse( + {"type": "reasoning", "sessionID": "s", "part": {"text": "thinking hard"}}, + state, + ) + assert reasoning == [Reasoning(delta="thinking hard")] + failed = parse( + { + "type": "tool_use", + "part": { + "tool": "bash", + "callID": "c1", + "state": { + "status": "error", + "input": {"command": "x"}, + "error": "boom", + }, + }, + }, + state, + ) + assert failed[1] == ToolResult(id="c1", output="boom", is_error=True) + for native, normalized in [ + ("list", "ls"), + ("webfetch", "web_search"), + ("glob", "glob"), + ("grep", "grep"), + ("apply_patch", "edit"), + ]: + call = parse( + { + "type": "tool_use", + "part": { + "tool": native, + "callID": "x", + "state": {"status": "completed", "input": {}, "output": ""}, + }, + }, + state, + )[0] + assert call.name == normalized + mcp = parse( + { + "type": "tool_use", + "part": { + "tool": "github_search", + "callID": "m", + "state": {"status": "completed", "input": {}, "output": {"a": 1}}, + }, + }, + state, + ) + assert mcp[0].builtin is False and mcp[1].output == '{"a": 1}' + assert state.session_id == "s" + + +def test_final_text_is_last_step_text(): + state = OpenCodeStreamState() + parse({"type": "step_start"}, state) + parse({"type": "text", "part": {"text": "working"}}, state) + parse({"type": "step_start"}, state) + parse({"type": "text", "part": {"text": "a"}}, state) + parse({"type": "text", "part": {"text": "b"}}, state) + assert state.final_text == "a\n\nb" + + +def test_error_event_message_shapes(): + state = OpenCodeStreamState() + parse({"type": "error", "error": {"data": {"message": "m1"}}}, state) + parse({"type": "error", "error": {"name": "APIError"}}, state) + parse({"type": "error", "error": "raw"}, state) + assert state.error == "m1\nAPIError\nraw" + + +# --------------------------------------------------------------------------- config + + +def test_permission_mapping(): + assert permission_rules("full", ()) == {"*": "allow"} + assert permission_rules("read-only", ()) == { + "edit": "deny", + "bash": "deny", + "webfetch": "deny", + } + edit = permission_rules("edit", ()) + assert edit["edit"] == "allow" and edit["bash"] == "deny" + with pytest.raises(CapabilityUnsupported): + permission_rules("ask", ()) + + +def test_disable_tools_map_to_native_denies_after_wildcard(): + rules = permission_rules("full", ["bash", "web_search", "ls", "write"]) + assert list(rules)[0] == "*" + assert rules["bash"] == "deny" + assert rules["webfetch"] == rules["websearch"] == "deny" + assert rules["list"] == "deny" + assert rules["edit"] == "deny" + + +def test_build_config_merges_user_config_under_managed_keys(): + config = build_opencode_config( + model="m1", + base_url="http://h:1/v1", + token_path="/tmp/p/token", + permissions="full", + user_config={"instructions": ["RULES.md"], "compaction": {"auto": False}}, + instructions_path="/tmp/p/instructions.md", + skills_path="/tmp/p/skills", + ) + provider = config["provider"]["litellm"] + assert provider["npm"] == "@ai-sdk/openai-compatible" + assert provider["options"] == { + "baseURL": "http://h:1/v1", + "apiKey": "{file:/tmp/p/token}", + } + assert provider["models"] == {"m1": {}} + assert config["model"] == config["small_model"] == "litellm/m1" + assert config["enabled_providers"] == ["litellm"] + assert config["instructions"] == ["RULES.md", "/tmp/p/instructions.md"] + assert config["skills"] == {"paths": ["/tmp/p/skills"]} + assert config["compaction"] == {"auto": False} + + +@pytest.mark.parametrize( + "config", + [ + {"provider": {}}, + {"model": "openai/gpt-5"}, + {"permission": {"*": "allow"}}, + {"tools": {"bash": True}}, + {"agent": {"build": {"permission": {"bash": "allow"}}}}, + {"mode": {"x": {"model": "a/b"}}}, + {"agent": "not-a-mapping"}, + ], +) +def test_managed_keys_rejected(config): + with pytest.raises(OptionsMismatch): + validate_user_config(config) + + +def test_config_metadata(): + assert CONFIG.get_binary() == "opencode" + assert "opencode" in CONFIG.get_install_hint() + assert CONFIG.uses_model_endpoint is True + assert CONFIG.capabilities.permission_modes == {"read-only", "edit", "full"} + + +def test_validate_environment_rejects_wrong_options_and_managed_config(): + CONFIG.validate_environment(make_ctx()) + with pytest.raises(OptionsMismatch): + CONFIG.validate_environment(make_ctx(options=CodexOptions())) + with pytest.raises(OptionsMismatch): + CONFIG.validate_environment( + make_ctx(options=OpenCodeOptions(config={"model": "openai/x"})) + ) + + +def test_session_setup_token_only_in_private_file(): + setup = setup_for(make_ctx()) + assert setup.files == {TOKEN_FILENAME: TOKEN.encode()} + assert TOKEN not in json.dumps(dict(setup.env)) + config = setup_config(setup) + assert config["provider"]["litellm"]["options"] == { + "baseURL": "http://host.docker.internal:4555/v1", + "apiKey": "{file:/tmp/oc-1/token}", + } + assert config["permission"] == {"*": "allow"} + + +def test_session_setup_env_and_persisted_xdg(): + setup = setup_for(make_ctx(options=OpenCodeOptions(env={"FOO": "1"}))) + assert list(setup.persisted_dirs) == [(XDG_DIRNAME, "opencode")] + assert setup.skills_dir == "skills" + env = setup.env + for sub in ("config", "data", "state", "cache"): + assert env[f"XDG_{sub.upper()}_HOME"] == f"{PRIVATE}/xdg/{sub}" + for key, value in OPENCODE_ISOLATION_ENV.items(): + assert env[key] == value + assert env["OPENCODE_CONFIG"] == "" and env["OPENCODE_PERMISSION"] == "" + assert env["FOO"] == "1" + + +def test_session_setup_errors(): + with pytest.raises(HarnessError): + setup_for(make_ctx(endpoint=None)) + with pytest.raises(ValueError, match="needs model="): + setup_for(make_ctx(model=None, endpoint=FakeEndpoint(model=None))) + + +def test_session_setup_read_only_and_disable_tools(): + setup = setup_for(make_ctx(permissions="read-only", disable_tools=["grep"])) + assert setup_config(setup)["permission"] == { + "edit": "deny", + "bash": "deny", + "webfetch": "deny", + "grep": "deny", + } + + +def test_session_setup_instructions_and_skills(): + ctx = make_ctx(instructions="Be terse.", output=Answer, skills=["/s/greeter"]) + setup = setup_for(ctx) + written = setup.files[INSTRUCTIONS_FILENAME].decode() + assert written == build_instructions(ctx) + assert written.startswith("Be terse.") + assert '"file"' in written and "single JSON object" in written + config = setup_config(setup) + assert config["instructions"] == ["/tmp/oc-1/instructions.md"] + assert config["skills"] == {"paths": ["/tmp/oc-1/skills"]} + assert build_instructions(make_ctx()) is None + assert "skills" not in setup_config(setup_for(make_ctx())) + + +def test_turn_request_argv_and_session_continuation(): + ctx = make_ctx(options=OpenCodeOptions(agent="build")) + setup = setup_for(ctx) + first = CONFIG.transform_turn_request(ctx, setup, PRIVATE, "hello", None) + assert list(first.argv) == [ + "opencode", + "run", + "--pure", + "--format", + "json", + "--thinking", + "-m", + "litellm/claude-haiku-4-5-20251001", + "--agent", + "build", + "--title", + OPENCODE_SESSION_TITLE, + ] + assert first.cwd == "/work" + assert first.stdin == "hello" + assert first.env == setup.env + second = CONFIG.transform_turn_request(ctx, setup, PRIVATE, "again", SESSION) + argv = list(second.argv) + assert argv[argv.index("--session") + 1] == SESSION + assert "--title" not in argv + assert "again" not in " ".join(argv) + + +def test_turn_prompt_repeats_schema_when_output_set(): + assert turn_prompt(make_ctx(), "hi") == "hi" + prompt = turn_prompt(make_ctx(output=Answer), "hi") + assert prompt.startswith("hi\n\n") and '"file"' in prompt + ctx = make_ctx(output=Answer) + request = CONFIG.transform_turn_request(ctx, setup_for(ctx), PRIVATE, "hi", None) + assert request.stdin == prompt + + +def test_turn_response_paths(): + state = OpenCodeStreamState(final_text='Here: {"file": "a", "content": "hi"}') + ok = CONFIG.transform_turn_response(make_ctx(output=Answer), state, 0, []) + assert json.loads(ok.output_json) == {"file": "a", "content": "hi"} + plain = CONFIG.transform_turn_response(make_ctx(), state, 0, []) + assert plain.output_json is None and plain.final_text == state.final_text + with pytest.raises(HarnessTurnError, match="boom"): + CONFIG.transform_turn_response( + make_ctx(), OpenCodeStreamState(error="boom"), 0, [] + ) + with pytest.raises(HarnessTurnError, match="code 3: no output"): + CONFIG.transform_turn_response(make_ctx(), OpenCodeStreamState(), 3, []) + + +# --------------------------------------------------------------------------- handler + + +async def test_start_writes_token_only_in_private_file(): + handler, ctx, sandbox = await started() + assert sandbox.files["/tmp/oc-1/token"] == TOKEN.encode() + sandbox.outputs.append(fixture_proc("turn1_write_read.jsonl")) + await collect(handler, ctx, "create hello.txt containing hi then read it") + call = sandbox.execs[0] + assert TOKEN not in json.dumps(call["cmd"]) + assert TOKEN not in json.dumps(call["env"]) + config = exec_config(sandbox) + assert config["provider"]["litellm"]["options"] == { + "baseURL": "http://host.docker.internal:4555/v1", + "apiKey": "{file:/tmp/oc-1/token}", + } + assert config["permission"] == {"*": "allow"} + + +async def test_start_persists_xdg_dir(): + _, _, sandbox = await started() + assert sandbox.runs == [ + ["sh", "-c", PERSIST_DIR_SCRIPT, "sh", "/tmp/oc-1/xdg", "opencode"] + ] + + +async def test_persist_failure_still_uses_private_xdg(): + handler, ctx, sandbox = await started(FakeSandbox(persist_ok=False)) + sandbox.outputs.append(fixture_proc("turn1_write_read.jsonl")) + await collect(handler, ctx, "x") + assert sandbox.execs[0]["env"]["XDG_DATA_HOME"] == "/tmp/oc-1/xdg/data" + + +async def test_turn_argv_env_and_session_continuation(): + handler, ctx, sandbox = await started( + options=OpenCodeOptions(agent="build", env={"FOO": "1"}) + ) + sandbox.outputs.append(fixture_proc("turn1_write_read.jsonl")) + events = await collect(handler, ctx, "create hello.txt containing hi then read it") + first = sandbox.execs[0] + assert first["cmd"] == [ + "opencode", + "run", + "--pure", + "--format", + "json", + "--thinking", + "-m", + "litellm/claude-haiku-4-5-20251001", + "--agent", + "build", + "--title", + OPENCODE_SESSION_TITLE, + ] + assert first["cwd"] == "/work" + env = first["env"] + assert env["XDG_CONFIG_HOME"] == "/tmp/oc-1/xdg/config" + assert env["XDG_DATA_HOME"] == "/tmp/oc-1/xdg/data" + assert env["XDG_STATE_HOME"] == "/tmp/oc-1/xdg/state" + assert env["XDG_CACHE_HOME"] == "/tmp/oc-1/xdg/cache" + assert env["OPENCODE_DISABLE_AUTOUPDATE"] == "1" + assert env["OPENCODE_DISABLE_MODELS_FETCH"] == "1" + assert env["OPENCODE_CONFIG"] == "" and env["OPENCODE_PERMISSION"] == "" + assert env["FOO"] == "1" + assert any(isinstance(e, ToolCall) for e in events) + assert ctx.final_text.startswith("Done!") + assert handler.native_session_id() == SESSION + + sandbox.outputs.append(fixture_proc("turn2_session_skill.jsonl")) + await collect(handler, ctx, "what file did you create?") + second = sandbox.execs[1]["cmd"] + assert second[second.index("--session") + 1] == SESSION + assert "--title" not in second + assert "what file" not in " ".join(second) + + +async def test_prompt_is_sent_on_stdin_not_argv(): + handler, ctx, sandbox = await started() + proc = fixture_proc("turn1_write_read.jsonl") + sandbox.outputs.append(proc) + await collect(handler, ctx, "secret prompt text") + assert proc.stdin.data == b"secret prompt text" and proc.stdin.closed + assert "secret prompt text" not in sandbox.execs[0]["cmd"] + + +async def test_resume_sets_session(): + handler, ctx, sandbox = await started() + await handler.resume(ctx, "ses_prev") + sandbox.outputs.append(fixture_proc("turn2_session_skill.jsonl")) + await collect(handler, ctx, "hi") + cmd = sandbox.execs[0]["cmd"] + assert cmd[cmd.index("--session") + 1] == "ses_prev" + + +async def test_read_only_and_disable_tools_config(): + handler, ctx, sandbox = await started( + permissions="read-only", disable_tools=["grep"] + ) + sandbox.outputs.append(fixture_proc("readonly_denied_bash.jsonl")) + await collect(handler, ctx, "x") + assert exec_config(sandbox)["permission"] == { + "edit": "deny", + "bash": "deny", + "webfetch": "deny", + "grep": "deny", + } + + +async def test_instructions_and_structured_output(): + handler, ctx, sandbox = await started(instructions="Be terse.", output=Answer) + written = sandbox.files["/tmp/oc-1/instructions.md"].decode() + assert written.startswith("Be terse.") + assert '"file"' in written and "single JSON object" in written + lines = [ + {"type": "step_start", "sessionID": "s"}, + { + "type": "text", + "sessionID": "s", + "part": {"text": 'Here: {"file": "a", "content": "hi"}'}, + }, + ] + proc = FakeProcess("\n".join(json.dumps(line) for line in lines).encode()) + sandbox.outputs.append(proc) + await collect(handler, ctx, "x") + assert exec_config(sandbox)["instructions"] == ["/tmp/oc-1/instructions.md"] + assert '"file"' in proc.stdin.data.decode() + assert json.loads(ctx.output_json) == {"file": "a", "content": "hi"} + + +async def test_skills_copied_to_private_skills_path(tmp_path): + skill = tmp_path / "greeter" + (skill / "ref").mkdir(parents=True) + (skill / "SKILL.md").write_text("---\nname: greeter\ndescription: d\n---\nbody") + (skill / "ref" / "notes.txt").write_text("n") + handler, ctx, sandbox = await started(skills=[str(skill)]) + assert sandbox.files["/tmp/oc-1/skills/greeter/SKILL.md"].startswith(b"---") + assert sandbox.files["/tmp/oc-1/skills/greeter/ref/notes.txt"] == b"n" + sandbox.outputs.append(fixture_proc("turn2_session_skill.jsonl")) + await collect(handler, ctx, "x") + assert exec_config(sandbox)["skills"] == {"paths": ["/tmp/oc-1/skills"]} + + +async def test_missing_binary(): + with pytest.raises(HarnessInstallFailed, match="opencode"): + await started(FakeSandbox(has_binary=False)) + + +async def test_wrong_options_and_managed_config_rejected(): + with pytest.raises(OptionsMismatch): + await started(options=CodexOptions()) + with pytest.raises(OptionsMismatch): + await started(options=OpenCodeOptions(config={"model": "openai/x"})) + + +async def test_turn_before_start_raises(): + handler = CLIHarnessHandler(OpenCodeHarnessConfig()) + with pytest.raises(RuntimeError, match="before start"): + await collect(handler, make_ctx(), "x") + + +async def test_api_error_event_raises_even_on_exit_zero(): + handler, ctx, sandbox = await started() + sandbox.outputs.append(fixture_proc("api_error.jsonl")) + with pytest.raises(HarnessTurnError, match="no healthy deployments"): + await collect(handler, ctx, "x") + + +async def test_nonzero_exit_raises_with_stderr_tail(): + handler, ctx, sandbox = await started() + sandbox.outputs.append( + FakeProcess(b"", stderr=b"line1\nfatal: bad config\n", exit_code=2) + ) + with pytest.raises(HarnessTurnError, match="code 2: line1\nfatal: bad config"): + await collect(handler, ctx, "x") + + +async def test_early_close_kills_process_and_stop_is_idempotent(): + handler, ctx, sandbox = await started() + proc = fixture_proc("turn1_write_read.jsonl") + sandbox.outputs.append(proc) + gen = handler.turn(ctx, "x") + await gen.__anext__() + await gen.aclose() + assert proc.killed + await handler.stop(ctx) + await handler.stop(ctx) + + +async def test_model_falls_back_to_endpoint_model(): + handler, ctx, sandbox = await started( + model=None, endpoint=FakeEndpoint(model="gw-model") + ) + sandbox.outputs.append(fixture_proc("turn1_write_read.jsonl")) + await collect(handler, ctx, "x") + assert "litellm/gw-model" in sandbox.execs[0]["cmd"] + assert exec_config(sandbox)["provider"]["litellm"]["models"] == {"gw-model": {}} + + +def test_turn_request_never_loads_plugins(): + """A repo's .opencode/plugin/*.js would run as the host user at startup; --pure blocks it.""" + ctx = make_ctx() + argv = list(CONFIG.transform_turn_request(ctx, setup_for(ctx), PRIVATE, "hi", None).argv) + assert argv[:3] == ["opencode", "run", "--pure"] + + +def test_options_config_cannot_add_plugins(): + with pytest.raises(OptionsMismatch, match="plugin"): + validate_user_config({"plugin": ["./evil.js"]}) + + +def test_endpoint_request_fixture_documents_contract(): + requests = load_fixture("endpoint_requests.jsonl") + assert {r["path"] for r in requests} == {"/v1/chat/completions"} + assert all(r["stream"] is True for r in requests) + assert all(r["stream_options"] == {"include_usage": True} for r in requests) diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 0d0c65e3650..02ac1540071 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -8201,7 +8201,7 @@ class TestGatewaySessionAdmission: assert not any(k.lower() == "authorization" for k in (raw_headers or {})) -def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("sso-user",)): +def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=("sso-user",), toolsets=None): from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable, Member return LiteLLM_TeamTable( @@ -8210,7 +8210,10 @@ def _make_team(team_id, mcp_servers, *, org_id=None, tool_perms=None, members=(" members_with_roles=[Member(user_id=u, role="user") for u in members], access_group_ids=[], object_permission=LiteLLM_ObjectPermissionTable( - object_permission_id=f"op-{team_id}", mcp_servers=mcp_servers, mcp_tool_permissions=tool_perms + object_permission_id=f"op-{team_id}", + mcp_servers=mcp_servers, + mcp_tool_permissions=tool_perms, + mcp_toolsets=toolsets, ), ) @@ -8271,6 +8274,64 @@ class TestUserSubjectTeamUnion: result = await MCPRequestHandler.get_allowed_mcp_servers(auth) assert set(result) == {"srv1", "srv2", "srv3"} + async def test_toolsets_of_a_team_that_dropped_the_user_from_its_roster_are_not_granted(self): + """The user's cached team list still names team-revoked, but its live roster no longer lists + the user, so its toolset is withheld exactly as its servers are on the aggregate /mcp.""" + from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids + + teams = { + "team-kept": _make_team("team-kept", [], toolsets=["ts-kept"]), + "team-revoked": _make_team("team-revoked", [], toolsets=["ts-revoked"], members=("someone-else",)), + } + auth = _make_admitted_subject("sso-user") + with self._patch(teams_by_id=teams, user_teams=["team-kept", "team-revoked"]): + granted = await granted_toolset_ids(auth) + assert granted == {"ts-kept"} + + async def test_a_pinned_toolset_narrows_every_source_to_the_toolset_servers_and_tools(self): + """On /toolset/{name}/mcp the admitted subject carries mcp_toolset_id; team-a's grant on srv1 and + srv2 with every tool collapses to the toolset's srv1 and its one tool, and team-b's srv3 drops.""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"]), "team-b": _make_team("team-b", ["srv3"])} + auth = _make_admitted_subject("sso-user") + pinned = auth.model_copy(update={"mcp_toolset_id": "ts-1"}) + resolve = AsyncMock(return_value={"srv1": ["add"]}) + with ( + self._patch(teams_by_id=teams, user_teams=["team-a", "team-b"]), + patch.object(global_mcp_server_manager, "resolve_toolset_tool_permissions", resolve), + ): + servers = await MCPRequestHandler.resolve_admitted_subject_servers(pinned) + tools = await MCPRequestHandler.resolve_admitted_subject_tools("srv1", pinned) + unpinned_servers = await MCPRequestHandler.resolve_admitted_subject_servers(auth) + unpinned_tools = await MCPRequestHandler.resolve_admitted_subject_tools("srv1", auth) + assert servers == ["srv1"] + assert tools == ["add"] + assert set(unpinned_servers) == {"srv1", "srv2", "srv3"} + assert unpinned_tools is None + assert {call.kwargs["toolset_ids"][0] for call in resolve.await_args_list} == {"ts-1"} + + async def test_a_fresh_policy_pinned_toolset_bypasses_the_toolset_permission_cache(self): + """A session admitted under requires_fresh_policy reads the pinned toolset from the writer, so a + tool revoked from the toolset is gone on the very next request (Devin Review 4150024092).""" + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + teams = {"team-a": _make_team("team-a", ["srv1", "srv2"])} + auth = _make_admitted_subject("sso-user") + auth.requires_fresh_policy = True + pinned = auth.model_copy(update={"mcp_toolset_id": "ts-1"}) + resolve = AsyncMock(return_value={"srv1": ["add"]}) + with ( + self._patch(teams_by_id=teams, user_teams=["team-a"]), + patch.object(global_mcp_server_manager, "resolve_toolset_tool_permissions", resolve), + ): + servers = await MCPRequestHandler.resolve_admitted_subject_servers(pinned) + tools = await MCPRequestHandler.resolve_admitted_subject_tools("srv1", pinned) + assert servers == ["srv1"] + assert tools == ["add"] + assert resolve.await_args_list + assert all(call.kwargs == {"toolset_ids": ["ts-1"], "requires_fresh_policy": True} for call in resolve.await_args_list) + async def test_key_based_caller_uses_single_team_only(self): """A key-based caller (api_key set) with a team_id sees ONLY that team, even though the same user belongs to other teams: key auth must be byte-identical to before.""" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py index 1398884783e..95be2b8b12b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_toolset_scope.py @@ -1,10 +1,12 @@ """Tests for MCP toolset scope enforcement.""" import asyncio +from collections.abc import Awaitable, Callable from typing import Dict, List, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi import HTTPException from litellm.proxy._types import ( LiteLLM_ObjectPermissionTable, @@ -30,6 +32,19 @@ def _make_auth( ) +def _granted_through_team(*team_toolset_ids: str) -> Callable[[UserAPIKeyAuth], Awaitable[frozenset[str]]]: + """The real grant resolver over a team that holds ``team_toolset_ids``, with no key access rule.""" + from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids + + async def team_permission(context: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable: + return LiteLLM_ObjectPermissionTable(object_permission_id="team-op", mcp_toolsets=list(team_toolset_ids)) + + async def granted(context: UserAPIKeyAuth) -> frozenset[str]: + return await granted_toolset_ids(context, team_object_permission=team_permission, require_key_access=False) + + return granted + + class TestApplyToolsetScope: """Tests for _apply_toolset_scope helper.""" @@ -97,6 +112,122 @@ class TestApplyToolsetScope: assert op.mcp_servers == ["server-a"] assert op.mcp_tool_permissions == toolset_perms + @pytest.mark.asyncio + async def test_team_granted_toolset_is_served_to_a_key_without_its_own_grant(self): + """A team key whose own row carries no toolset grant is admitted to the toolset its team + holds (LIT-6029), scoped to that toolset's servers and tools.""" + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + toolset_perms = {"server-a": ["tool1"]} + auth = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=None) + with patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.resolve_toolset_tool_permissions", + new=AsyncMock(return_value=toolset_perms), + ): + result = await _apply_toolset_scope(auth, "toolset-123", granted=_granted_through_team("toolset-123")) + + assert result.mcp_toolset_id == "toolset-123" + assert result.object_permission is not None + assert result.object_permission.mcp_servers == ["server-a"] + assert result.object_permission.mcp_tool_permissions == toolset_perms + + @pytest.mark.asyncio + async def test_a_non_admin_dashboard_session_is_pinned_as_its_admitted_user_instead_of_rewritten(self): + """The dashboard session acts as its admitted user, whose team grants resolve per source, so a + team-granted toolset is not capped by the user's own row: the row stays intact and the toolset + rides along as mcp_toolset_id (LIT-6029).""" + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1") + own_row = LiteLLM_ObjectPermissionTable(object_permission_id="user-op", mcp_servers=["server-own"]) + admitted = UserAPIKeyAuth(user_id="user-1", object_permission=own_row) + admitted.mcp_admitted_user_subject = True + granted = AsyncMock(return_value=frozenset({"toolset-123"})) + resolve = AsyncMock(return_value={"server-team": ["tool1"]}) + with patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.resolve_toolset_tool_permissions", + new=resolve, + ): + result = await _apply_toolset_scope( + session, "toolset-123", acting_user=AsyncMock(return_value=admitted), granted=granted + ) + + assert granted.await_args is not None and granted.await_args.args[0].mcp_admitted_user_subject is True + assert result.mcp_admitted_user_subject is True + assert result.mcp_toolset_id == "toolset-123" + assert result.object_permission == own_row + resolve.assert_not_awaited() + + @pytest.mark.asyncio + async def test_a_gateway_admitted_user_without_the_toolset_in_any_source_is_denied(self): + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None) + admitted.mcp_admitted_user_subject = True + granted = AsyncMock(return_value=frozenset({"toolset-other"})) + with pytest.raises(HTTPException) as exc_info: + await _apply_toolset_scope(admitted, "toolset-123", granted=granted) + + assert exc_info.value.status_code == 403 + granted.assert_awaited_once_with(admitted) + + @pytest.mark.asyncio + async def test_a_resource_scoped_admitted_user_is_denied_a_team_toolset_on_another_server(self): + """A gateway bearer scoped to server-own (RFC 8707 resource) cannot open a team toolset whose + servers lie outside that resource, even though the team grants it (Devin Review 4150024267).""" + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None) + admitted.mcp_admitted_user_subject = True + admitted.mcp_session_resource_server_id = "server-own" + admitted.requires_fresh_policy = True + granted = AsyncMock(return_value=frozenset({"toolset-123"})) + resolve = AsyncMock(return_value={"server-team": ["tool1"]}) + with patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.resolve_toolset_tool_permissions", + new=resolve, + ): + with pytest.raises(HTTPException) as exc_info: + await _apply_toolset_scope(admitted, "toolset-123", granted=granted) + + assert exc_info.value.status_code == 403 + resolve.assert_awaited_once_with(toolset_ids=["toolset-123"], requires_fresh_policy=True) + + @pytest.mark.asyncio + async def test_a_resource_scoped_admitted_user_opens_a_toolset_inside_its_resource(self): + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + admitted = UserAPIKeyAuth(user_id="user-1", object_permission=None) + admitted.mcp_admitted_user_subject = True + admitted.mcp_session_resource_server_id = "server-team" + granted = AsyncMock(return_value=frozenset({"toolset-123"})) + resolve = AsyncMock(return_value={"server-team": ["tool1"], "server-other": ["tool2"]}) + with patch( + "litellm.proxy._experimental.mcp_server.server." + "global_mcp_server_manager.resolve_toolset_tool_permissions", + new=resolve, + ): + result = await _apply_toolset_scope(admitted, "toolset-123", granted=granted) + + assert result.mcp_toolset_id == "toolset-123" + assert result.mcp_session_resource_server_id == "server-team" + resolve.assert_awaited_once_with(toolset_ids=["toolset-123"], requires_fresh_policy=False) + + @pytest.mark.asyncio + async def test_team_grant_for_another_toolset_does_not_admit_a_key_to_this_one(self): + from litellm.proxy._experimental.mcp_server.server import _apply_toolset_scope + + auth = _make_auth(mcp_toolsets=[]) + auth.team_id = "team-a" + with pytest.raises(HTTPException) as exc_info: + await _apply_toolset_scope(auth, "toolset-123", granted=_granted_through_team("toolset-other")) + + assert exc_info.value.status_code == 403 + @pytest.mark.asyncio async def test_non_admin_no_object_permission_raises_403(self): """Non-admin key with object_permission=None is denied (no grants configured).""" @@ -250,6 +381,131 @@ class TestFetchMCPToolsetsAccess: assert len(result) == 2 mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-1", "ts-2"]) + @pytest.mark.asyncio + async def test_team_granted_toolsets_are_listed_for_a_key_without_its_own_grant(self): + """GET /v1/mcp/toolset for a team key lists the team's toolsets (LIT-6029).""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=None) + team_permission = LiteLLM_ObjectPermissionTable(object_permission_id="team-op", mcp_toolsets=["ts-team"]) + fake_toolsets = [MagicMock(toolset_id="ts-team")] + mock_client = MagicMock() + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=fake_toolsets), + ) as mock_list, + patch.object( + MCPRequestHandler, + "_get_team_object_permission", + new=AsyncMock(return_value=team_permission), + ), + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == fake_toolsets + mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-team"]) + + @pytest.mark.asyncio + async def test_admin_with_own_grants_is_not_narrowed_by_a_team_lookup(self): + """An admin's own grant list is the only filter; no team lookup runs for admins.""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolsets, + ) + + auth = _make_auth(mcp_toolsets=["ts-1"]) + auth.user_role = LitellmUserRoles.PROXY_ADMIN + mock_client = MagicMock() + own_toolsets = [{"toolset_id": "ts-1", "toolset_name": "own"}] + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=mock_client, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.list_mcp_toolsets", + new=AsyncMock(return_value=own_toolsets), + ) as mock_list, + patch.object( + MCPRequestHandler, "_get_team_object_permission", new=AsyncMock(return_value=None) + ) as team_lookup, + ): + result = await fetch_mcp_toolsets(user_api_key_dict=auth) + + assert result == own_toolsets + mock_list.assert_called_once_with(mock_client, toolset_ids=["ts-1"]) + team_lookup.assert_not_awaited() + + +class TestFetchMCPToolsetAccess: + """Tests for GET /v1/mcp/toolset/{toolset_id} access control.""" + + @staticmethod + async def _fetch(auth: UserAPIKeyAuth, toolset_id: str, team_toolsets: list[str] | None): + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + fetch_mcp_toolset, + ) + + team_permission = ( + LiteLLM_ObjectPermissionTable(object_permission_id="team-op", mcp_toolsets=team_toolsets) + if team_toolsets is not None + else None + ) + toolset = MagicMock(toolset_id=toolset_id) + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_toolset", + new=AsyncMock(return_value=toolset), + ), + patch.object( + MCPRequestHandler, + "_get_team_object_permission", + new=AsyncMock(return_value=team_permission), + ), + ): + return await fetch_mcp_toolset(toolset_id=toolset_id, user_api_key_dict=auth) + + @pytest.mark.asyncio + async def test_team_granted_toolset_detail_is_served_to_a_key_without_its_own_grant(self): + auth = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=None) + + toolset = await self._fetch(auth, "ts-team", team_toolsets=["ts-team"]) + + assert toolset.toolset_id == "ts-team" + + @pytest.mark.asyncio + async def test_toolset_detail_stays_forbidden_when_neither_key_nor_team_holds_it(self): + from fastapi import HTTPException + + auth = _make_auth(mcp_toolsets=["ts-own"]) + auth.team_id = "team-a" + + with pytest.raises(HTTPException) as exc_info: + await self._fetch(auth, "ts-withheld", team_toolsets=["ts-team"]) + + assert exc_info.value.status_code == 403 + class TestToolsetPrefixResolution: """Regression for LIT-3419. diff --git a/tests/unit/proxy/_experimental/mcp_server/test_ui_session_utils.py b/tests/unit/proxy/_experimental/mcp_server/test_ui_session_utils.py index 816ccc5e7e6..293d9443ced 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_ui_session_utils.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_ui_session_utils.py @@ -6,13 +6,14 @@ import pytest from fastapi import HTTPException from litellm.constants import UI_SESSION_TOKEN_TEAM_ID -from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy._experimental.mcp_server.ui_session_utils import ( build_effective_auth_contexts, clone_user_api_key_auth_with_team, + granted_toolset_ids, + toolset_grant_contexts, resolve_ui_session_team_ids, ) +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth def test_clone_user_api_key_auth_with_team_creates_independent_copy(): @@ -258,3 +259,234 @@ async def test_admitted_user_context_carries_the_request_span(monkeypatch): assert (await acting_user_auth(user_auth)).parent_otel_span is parent_span assert (await build_effective_auth_contexts(user_auth))[-1].parent_otel_span is parent_span + + +def _toolset_permission(*toolset_ids: str) -> LiteLLM_ObjectPermissionTable: + return LiteLLM_ObjectPermissionTable( + object_permission_id=f"op-{'-'.join(toolset_ids)}", mcp_toolsets=list(toolset_ids) + ) + + +@pytest.mark.asyncio +async def test_granted_toolset_ids_unions_own_and_team_grants_over_every_effective_context(): + """A dashboard session of a user in two teams holds the toolsets of both teams plus the ones on + the user row itself, exactly the grant sources the aggregate /mcp listing expands.""" + session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1") + team_a = UserAPIKeyAuth(team_id="team-a", user_id="user-1") + team_b = UserAPIKeyAuth(team_id="team-b", user_id="user-1", object_permission=_toolset_permission()) + admitted = UserAPIKeyAuth(user_id="user-1", object_permission=_toolset_permission("ts-user")) + team_grants = {"team-a": _toolset_permission("ts-a", "ts-shared"), "team-b": _toolset_permission("ts-b")} + + async def effective_contexts(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + assert auth is session + return [team_a, team_b, admitted] + + async def team_permission(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + return team_grants.get(auth.team_id or "") + + granted = await granted_toolset_ids(session, effective_contexts, team_permission) + + assert granted == frozenset({"ts-a", "ts-shared", "ts-b", "ts-user"}) + + +@pytest.mark.asyncio +async def test_granted_toolset_ids_is_empty_when_neither_key_nor_team_grants_a_toolset(): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=_toolset_permission()) + + async def effective_contexts(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + return [auth] + + async def no_team_permission(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + return None + + assert await granted_toolset_ids(key, effective_contexts, no_team_permission) == frozenset() + + +async def _same_context(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + return [auth] + + +async def _team_grants_ts_team(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + return _toolset_permission("ts-team") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "own", + [ + LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_toolsets=["ts-own"]), + LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_servers=["srv-own"]), + LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_tool_permissions={"srv-own": ["add"]}), + LiteLLM_ObjectPermissionTable(object_permission_id="op", mcp_access_groups=["group-own"]), + ], +) +async def test_a_key_declaring_its_own_mcp_grant_does_not_inherit_the_team_toolsets(own): + """The key/team rule of the aggregate listing: a key's own MCP grant is a ceiling the team cannot widen.""" + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=own) + + granted = await granted_toolset_ids(key, _same_context, _team_grants_ts_team, require_key_access=False) + + assert granted == frozenset(own.mcp_toolsets or ()) + + +@pytest.mark.asyncio +async def test_require_key_mcp_access_defined_stops_a_key_inheriting_team_toolsets_but_not_a_session(): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a") + session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1") + + assert await granted_toolset_ids(key, _same_context, _team_grants_ts_team, require_key_access=False) == {"ts-team"} + assert await granted_toolset_ids(key, _same_context, _team_grants_ts_team, require_key_access=True) == frozenset() + assert await granted_toolset_ids(session, _same_context, _team_grants_ts_team, require_key_access=True) == { + "ts-team" + } + + +@pytest.mark.asyncio +async def test_toolset_grant_contexts_of_a_virtual_key_is_the_key_alone(): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a") + + async def never(auth: UserAPIKeyAuth) -> None: + raise AssertionError("a virtual key has no admitted sources") + + assert await toolset_grant_contexts(key, admitted_context=never, admitted_sources=never) == (key,) + + +def _admitted(user_id: str, own: LiteLLM_ObjectPermissionTable | None = None) -> UserAPIKeyAuth: + subject = UserAPIKeyAuth(user_id=user_id, object_permission=own) + subject.mcp_admitted_user_subject = True + return subject + + +@pytest.mark.asyncio +async def test_toolset_grant_contexts_of_a_dashboard_session_are_its_admitted_users_grant_sources(): + """The dashboard session fans out through the same roster-checked source builder as the aggregate + /mcp resolution, applied to the admitted user it acts as, so a cached membership a team has since + revoked never reaches the toolset check.""" + session = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="user-1") + admitted = _admitted("user-1") + own_source = UserAPIKeyAuth(user_id="user-1") + team_source = UserAPIKeyAuth(user_id="user-1", team_id="team-a") + + async def admitted_context(auth: UserAPIKeyAuth) -> UserAPIKeyAuth: + assert auth is session + return admitted + + async def admitted_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + assert auth is admitted + return [own_source, team_source] + + assert await toolset_grant_contexts(session, admitted_context, admitted_sources) == (own_source, team_source) + + +@pytest.mark.asyncio +async def test_toolset_grant_contexts_of_a_gateway_admitted_user_are_its_own_grant_sources(): + admitted = _admitted("user-1") + team_source = UserAPIKeyAuth(user_id="user-1", team_id="team-a") + + async def no_dashboard_context(auth: UserAPIKeyAuth) -> None: + return None + + async def admitted_sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + assert auth is admitted + return [team_source] + + assert await toolset_grant_contexts(admitted, no_dashboard_context, admitted_sources) == (team_source,) + + +@pytest.mark.asyncio +async def test_a_source_declaring_its_own_mcp_grant_never_reads_its_team(): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission=_toolset_permission("ts-own")) + team_reads: list[str | None] = [] # mutable-ok: records the lookups the code under test performs + + async def team_permission(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + team_reads.append(auth.team_id) + return _toolset_permission("ts-team") + + granted = await granted_toolset_ids(key, _same_context, team_permission, require_key_access=False) + + assert granted == {"ts-own"} + assert team_reads == [] + + +async def _team_a_unreadable(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + if auth.team_id == "team-a": + raise RuntimeError("team row unreadable") + return _toolset_permission("ts-b") + + +@pytest.mark.asyncio +async def test_an_unreadable_team_grants_nothing_while_the_direct_and_other_team_grants_still_count(): + """A dashboard user whose own row grants ts-user and who sits on team-a and team-b keeps ts-user and + ts-b when team-a cannot be read; team-a itself contributes nothing rather than failing the lookup.""" + admitted = _admitted("user-1", _toolset_permission("ts-user")) + team_a = UserAPIKeyAuth(user_id="user-1", team_id="team-a") + team_b = UserAPIKeyAuth(user_id="user-1", team_id="team-b") + + async def sources(auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + return [admitted, team_a, team_b] + + assert await granted_toolset_ids(admitted, sources, _team_a_unreadable) == {"ts-user", "ts-b"} + + +@pytest.mark.asyncio +async def test_a_key_whose_only_grant_source_is_an_unreadable_team_is_granted_nothing(): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a") + + assert await granted_toolset_ids(key, _same_context, _team_a_unreadable, require_key_access=False) == frozenset() + + +async def _hydrates_op_key_to_srv_own(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + if auth.object_permission is not None: + return auth.object_permission + if auth.object_permission_id == "op-key": + return LiteLLM_ObjectPermissionTable(object_permission_id="op-key", mcp_servers=["srv-own"]) + return None + + +@pytest.mark.asyncio +async def test_a_key_cached_with_its_own_grant_unhydrated_is_scoped_to_that_grant_not_its_team(): + """The main auth flow can cache a key with object_permission_id set and object_permission None. The + row it names is the key's ceiling, so it is loaded and read as the key's own grant instead of letting the + key inherit its team's toolsets.""" + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission_id="op-key") + + granted = await granted_toolset_ids( + key, + _same_context, + _team_grants_ts_team, + require_key_access=False, + own_object_permission=_hydrates_op_key_to_srv_own, + ) + + assert granted == frozenset() + + +async def _own_row_unreadable(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + raise RuntimeError("object permission row unreadable") + + +async def _own_row_gone(auth: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable | None: + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("load_own", [_own_row_unreadable, _own_row_gone]) +async def test_a_key_naming_an_own_grant_that_cannot_be_read_is_granted_nothing_rather_than_its_team(load_own): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a", object_permission_id="op-key") + + granted = await granted_toolset_ids( + key, _same_context, _team_grants_ts_team, require_key_access=False, own_object_permission=load_own + ) + + assert granted == frozenset() + + +@pytest.mark.asyncio +async def test_a_key_naming_no_own_grant_is_not_hydrated_before_inheriting_its_team(): + key = UserAPIKeyAuth(api_key="sk-test", team_id="team-a") + + granted = await granted_toolset_ids( + key, _same_context, _team_grants_ts_team, require_key_access=False, own_object_permission=_own_row_unreadable + ) + + assert granted == {"ts-team"} diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 99e494fccd5..a34d3c0c27d 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/tests/unit/proxy/db/test_prisma_client.py @@ -454,3 +454,19 @@ def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashi assert PrismaManager.setup_database(use_migrate=False) is False assert fake_prisma_cli.calls == [] + + +@pytest.mark.parametrize( + ("run", "outcome"), + ( + (PrismaManager.build_request_log_indexes, False), + (PrismaManager.start_request_log_index_build, None), + ), + ids=("wait-for-the-build", "start-the-build"), +) +def test_without_proxy_extras_the_index_build_reports_failure_instead_of_raising(monkeypatch, run, outcome): + """The migration job exits non-zero and a serving proxy keeps booting when the extras + package that owns the index build is not installed.""" + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.utils", None) + + assert run() is outcome diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 32b3c919b07..6d56d325dc3 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -13596,6 +13596,63 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp assert mock_audit.call_args.kwargs["team_alias"] == "list-audit" +@pytest.mark.asyncio +async def test_team_member_add_evicts_the_cached_team_roster(monkeypatch): + """Roster checks read the team through get_team_object, so a cached pre-add roster must be dropped.""" + from litellm.proxy._types import TeamMemberAddRequest + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.team_endpoints import team_member_add + + team_id = "team-roster-evict" + team_row = LiteLLM_TeamTable(team_id=team_id, team_alias="roster-evict", members_with_roles=[]) + cache = UserApiKeyCache() + cache.set_cache(key=f"team_id:{team_id}", value=team_row) + cache.set_cache(key="team_alias:roster-evict", value=team_row) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id") + + joined_user = LiteLLM_UserTable(user_id="joiner", max_budget=None, spend=0.0, models=[]) + updated_team = MagicMock() + updated_team.model_dump.return_value = {"team_id": team_id, "members_with_roles": []} + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=team_row, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._validate_team_member_add_permissions", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._validate_and_populate_member_user_info", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._resolve_existing_member_user_ids", + new_callable=AsyncMock, + return_value=frozenset(), + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + new_callable=AsyncMock, + return_value=(updated_team, [joined_user], []), + ), + patch("litellm.proxy.management_endpoints.team_endpoints._schedule_team_member_add_audit_logs"), + ): + await team_member_add( + data=TeamMemberAddRequest(team_id=team_id, member=Member(user_id="joiner", role="user")), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + assert cache.get_cache(key=f"team_id:{team_id}") is None + assert cache.get_cache(key="team_alias:roster-evict") is None + + class _RecordingAuditLogger(CustomLogger): def __init__(self) -> None: super().__init__() diff --git a/tests/unit/proxy/test__lazy_features.py b/tests/unit/proxy/test__lazy_features.py new file mode 100644 index 00000000000..d2bb8244c2f --- /dev/null +++ b/tests/unit/proxy/test__lazy_features.py @@ -0,0 +1,211 @@ +import sys +from collections.abc import AsyncGenerator, Mapping +from contextlib import asynccontextmanager +from types import ModuleType +from typing import Final + +import pytest +from fastapi import APIRouter, FastAPI +from fastapi.testclient import TestClient +from pydantic import BaseModel + +from litellm.proxy._lazy_features import ( + LazyFeature, + LazyFeatureMiddleware, + attach_lazy_features, + lazy_tag_to_prefix, + loaded_lazy_modules, +) + +FLAG: Final = "LITELLM_DISABLE_LAZY_ROUTES" +WARMUP_PATH: Final = "/lazy/warm/{name}" + + +class _Operation(BaseModel): + tags: tuple[str, ...] + + +class _WarmupBody(BaseModel): + stub_path: str + paths: Mapping[str, Mapping[str, _Operation]] + + +def _feature_module(monkeypatch: pytest.MonkeyPatch, name: str, path: str) -> LazyFeature: + async def served() -> dict[str, str]: + return {"feature": name} + + router: Final = APIRouter() + router.add_api_route(path, served, methods=["GET"]) + module: Final = ModuleType(f"tests.unit.proxy.lazy_fixture_{name}") + module.router = router # pyright: ignore[reportAttributeAccessIssue] # fixture module built at test time + monkeypatch.setitem(sys.modules, module.__name__, module) + return LazyFeature(name=name, module_path=module.__name__, path_prefixes=(path,)) + + +def _paths(app: FastAPI) -> tuple[str, ...]: + return tuple(str(getattr(route, "path", "")) for route in app.routes) + + +def _has_lazy_middleware(app: FastAPI) -> bool: + return any(middleware.cls is LazyFeatureMiddleware for middleware in app.user_middleware) + + +@pytest.mark.parametrize("value", ("1", "true", "TRUE", "yes", "on")) +def test_flag_registers_every_feature_at_startup(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv(FLAG, value) + features: Final = ( + _feature_module(monkeypatch, "alpha", "/alpha/list"), + _feature_module(monkeypatch, "beta", "/beta/list"), + ) + app: Final = FastAPI() + + attach_lazy_features(app, features) + + assert WARMUP_PATH not in _paths(app) + assert not _has_lazy_middleware(app) + assert loaded_lazy_modules(app) == set() + with TestClient(app) as client: + at_startup: Final = _paths(app) + assert {"/alpha/list", "/beta/list"} <= set(at_startup) + assert loaded_lazy_modules(app) == {features[0].module_path, features[1].module_path} + assert client.get("/beta/list").json() == {"feature": "beta"} + assert client.post("/lazy/warm/alpha").status_code == 404 + assert _paths(app) == at_startup, "first feature request changed the table" + + +def test_flag_registers_before_the_inner_lifespan_and_after_late_routes(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(FLAG, "true") + features: Final = (_feature_module(monkeypatch, "epsilon", "/epsilon/{name}"),) + seen_by_inner_lifespan: Final[list[tuple[str, ...]]] = [] # mutable-ok: captured from inside the lifespan + + @asynccontextmanager + async def inner_lifespan(app_: FastAPI) -> AsyncGenerator[None]: + seen_by_inner_lifespan.append(_paths(app_)) + yield + + async def late() -> dict[str, str]: + return {"feature": "late"} + + app: Final = FastAPI(lifespan=inner_lifespan) + attach_lazy_features(app, features) + app.add_api_route("/epsilon/list", late, methods=["GET"]) + + with TestClient(app) as client: + assert client.get("/epsilon/list").json() == {"feature": "late"}, "late eager route must win, as in lazy mode" + assert client.get("/epsilon/x").json() == {"feature": "epsilon"} + assert seen_by_inner_lifespan == [_paths(app)], "startup hooks inside the proxy lifespan must see the full table" + + +def test_flag_lets_a_route_added_during_startup_beat_an_overlapping_feature_route( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv(FLAG, "true") + features: Final = (_feature_module(monkeypatch, "zeta", "/zeta/{endpoint:path}"),) + + async def configured() -> dict[str, str]: + return {"feature": "configured"} + + @asynccontextmanager + async def adds_a_pass_through(app_: FastAPI) -> AsyncGenerator[None]: + app_.add_api_route("/zeta/{subpath:path}", configured, methods=["GET"]) + yield + + app: Final = FastAPI(lifespan=adds_a_pass_through) + attach_lazy_features(app, features) + + with TestClient(app) as client: + assert client.get("/zeta/health").json() == {"feature": "configured"}, ( + "lazy mode routes this to startup's route" + ) + + +def test_flag_does_not_bring_back_a_feature_route_removed_during_startup(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(FLAG, "true") + features: Final = ( + _feature_module(monkeypatch, "eta", "/eta/list"), + _feature_module(monkeypatch, "theta", "/theta/list"), + ) + + @asynccontextmanager + async def drops_eta(app_: FastAPI) -> AsyncGenerator[None]: + app_.router.routes[:] = [route for route in app_.router.routes if getattr(route, "path", "") != "/eta/list"] + yield + + app: Final = FastAPI(lifespan=drops_eta) + attach_lazy_features(app, features) + + with TestClient(app) as client: + assert client.get("/eta/list").status_code == 404 + assert client.get("/theta/list").json() == {"feature": "theta"} + assert "/eta/list" not in _paths(app) + + +@pytest.mark.parametrize("value", (None, "", "0", "false", "off")) +def test_without_the_flag_features_still_mount_on_first_request( + monkeypatch: pytest.MonkeyPatch, value: str | None +) -> None: + if value is None: + monkeypatch.delenv(FLAG, raising=False) + else: + monkeypatch.setenv(FLAG, value) + features: Final = (_feature_module(monkeypatch, "gamma", "/gamma/list"),) + app: Final = FastAPI() + + attach_lazy_features(app, features) + + assert "/gamma/list" not in _paths(app) + assert WARMUP_PATH in _paths(app) + assert _has_lazy_middleware(app) + with TestClient(app) as client: + assert client.get("/gamma/list").json() == {"feature": "gamma"} + assert "/gamma/list" in _paths(app) + + +def test_flag_keeps_registering_after_one_feature_fails_to_import(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(FLAG, "true") + broken: Final = LazyFeature( + name="broken", module_path="tests.unit.proxy.lazy_fixture_does_not_exist", path_prefixes=("/broken",) + ) + healthy: Final = _feature_module(monkeypatch, "delta", "/delta/list") + app: Final = FastAPI() + + attach_lazy_features(app, (broken, healthy)) + + with TestClient(app) as client: + assert "/delta/list" in _paths(app) + assert loaded_lazy_modules(app) == {broken.module_path, healthy.module_path} + assert client.get("/delta/list").json() == {"feature": "delta"} + assert client.get("/broken").status_code == 404 + + +def test_flag_hides_the_swagger_warmup_plugin(monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy._lazy_openapi_snapshot as snapshot + + monkeypatch.setattr(snapshot, "SNAPSHOT_FILE", snapshot.SNAPSHOT_FILE.with_name("missing-snapshot.json")) + monkeypatch.setenv(FLAG, "false") + assert lazy_tag_to_prefix() != {}, "control: without the flag and without a snapshot the plugin has tags" + monkeypatch.setenv(FLAG, "true") + assert lazy_tag_to_prefix() == {} + + +def test_without_the_flag_the_warmup_route_registers_a_feature_and_returns_its_paths( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv(FLAG, raising=False) + features: Final = ( + _feature_module(monkeypatch, "alpha", "/alpha/list"), + _feature_module(monkeypatch, "beta", "/beta/list"), + ) + app: Final = FastAPI() + attach_lazy_features(app, features) + + with TestClient(app) as client: + assert client.post("/lazy/warm/zeta").status_code == 404 + warmed: Final = client.post("/lazy/warm/alpha") + assert warmed.status_code == 200, warmed.text + body: Final = _WarmupBody.model_validate_json(warmed.text) + assert body.stub_path == "/alpha/list" + assert set(body.paths) == {"/alpha/list"} + assert body.paths["/alpha/list"]["get"].tags == ("alpha",) + assert loaded_lazy_modules(app) == {features[0].module_path} + assert "/alpha/list" in _paths(app) and "/beta/list" not in _paths(app) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 7007189aab9..9520a94d0ea 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -1,5 +1,6 @@ import inspect import os +from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -2128,12 +2129,14 @@ class TestRunServerDbSetup: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") def test_use_prisma_db_push_flag_behavior( self, mock_should_update_schema, mock_check_schema_diff, + mock_build_indexes, mock_setup_database, mock_atexit_register, mock_subprocess_run, @@ -2187,9 +2190,7 @@ class TestRunServerDbSetup: # Test 1: Without --use_prisma_db_push flag (default behavior) # use_prisma_db_push should be False (default), so use_migrate should be True run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) - mock_setup_database.assert_called_with( - use_migrate=True, use_v2_resolver=True - ) + mock_setup_database.assert_called_with(use_migrate=True, use_v2_resolver=True) # Reset mocks mock_setup_database.reset_mock() @@ -2202,18 +2203,18 @@ class TestRunServerDbSetup: ["--local", "--skip_server_startup", "--use_prisma_db_push"], standalone_mode=False, ) - mock_setup_database.assert_called_with( - use_migrate=False, use_v2_resolver=True - ) + mock_setup_database.assert_called_with(use_migrate=False, use_v2_resolver=True) @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above def test_migrations_run_when_the_prisma_cli_is_not_on_path( self, mock_should_update_schema, mock_check_schema_diff, + mock_build_indexes, mock_setup_database, mock_atexit_register, tmp_path, @@ -2262,19 +2263,19 @@ class TestRunServerDbSetup: run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) assert "prisma CLI is neither on PATH" not in capsys.readouterr().out - mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=True - ) + mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") def test_startup_fails_when_db_setup_fails( self, mock_should_update_schema, mock_check_schema_diff, + mock_build_indexes, mock_setup_database, mock_atexit_register, mock_subprocess_run, @@ -2329,19 +2330,19 @@ class TestRunServerDbSetup: standalone_mode=False, ) assert exc_info.value.code == 1 - mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=True - ) + mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") def test_startup_exits_on_non_postgres_database_url( self, mock_should_update_schema, mock_check_schema_diff, + mock_build_indexes, mock_setup_database, mock_atexit_register, mock_subprocess_run, @@ -2387,12 +2388,14 @@ class TestRunServerDbSetup: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") def test_v2_migration_resolver_opts_in_via_env_var( self, mock_should_update_schema, mock_check_schema_diff, + mock_build_indexes, mock_setup_database, mock_atexit_register, mock_subprocess_run, @@ -2439,9 +2442,7 @@ class TestRunServerDbSetup: ["--local", "--skip_server_startup"], standalone_mode=False ) - mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=True - ) + mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out @pytest.mark.parametrize( @@ -2479,12 +2480,14 @@ class TestRunServerDbSetup: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") def test_legacy_resolver_flag_reaches_database_setup( self, mock_should_update_schema, mock_check_schema_diff, + mock_build_indexes, mock_setup_database, mock_atexit_register, mock_subprocess_run, @@ -2533,9 +2536,76 @@ class TestRunServerDbSetup: standalone_mode=False, ) - mock_setup_database.assert_called_once_with( - use_migrate=True, use_v2_resolver=False + mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=False) + + @pytest.mark.parametrize( + ("arguments", "migrated", "exits", "waits_for_the_build"), + ( + (("--local", "--skip_server_startup"), True, True, True), + (("--local",), True, False, False), + (("--local",), False, True, False), + ), + ids=("migration-job", "serving-proxy", "serving-proxy-whose-migrations-failed"), + ) + @patch("uvicorn.run") + @patch("subprocess.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True) + @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes", return_value=False) + @patch("litellm.proxy.db.prisma_client.PrismaManager.start_request_log_index_build") + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=True) + def test_the_migration_job_waits_for_the_index_build_and_a_serving_proxy_starts_it_in_the_background( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_start_build, + mock_build_indexes, + mock_setup_database, + mock_atexit_register, + mock_subprocess_run, + mock_uvicorn_run, + arguments, + migrated, + exits, + waits_for_the_build, + ): + """`--skip_server_startup` is the migration job: it waits for the index build after the + migrations and exits 1 when one could not be built. A serving proxy that ran the + migrations starts the build in the background and serves whatever the build does; one + whose migrations failed exits 1 under `--enforce_prisma_migration_check` and starts no build.""" + from litellm.proxy.proxy_cli import run_server + + mock_setup_database.return_value = migrated + mock_subprocess_run.return_value = MagicMock(returncode=0) + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), ) + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} + clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" + outcome = pytest.raises(SystemExit) if exits else nullcontext() + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, + ), + patch( + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + outcome as exc_info, + ): + mock_get_args.return_value = {"app": "litellm.proxy.proxy_server:app", "host": "localhost", "port": 8000} + run_server.main([*arguments, "--enforce_prisma_migration_check"], standalone_mode=False) + + assert (exc_info is not None and exc_info.value.code == 1) is exits + mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) + assert mock_build_indexes.call_count == int(migrated and waits_for_the_build) + assert mock_start_build.call_count == int(migrated and not waits_for_the_build) # --- Module-level helpers for worker startup hook tests --- diff --git a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py index 643944673a2..2699f9445c9 100644 --- a/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/unit/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -8,7 +8,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException -from mcp.types import CallToolResult, TextContent, Tool as MCPTool +from mcp.types import CallToolResult, TextContent +from mcp.types import Tool as MCPTool from openai.types.responses.tool_param import Mcp from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing @@ -110,9 +111,7 @@ def test_extract_tool_calls_from_chat_response_handles_tool_calls(): object="chat.completion", ) - tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( - response - ) + tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response(response) assert len(tool_calls) == 1 assert tool_calls[0]["function"]["name"] == "foo" @@ -182,9 +181,7 @@ def test_transform_mcp_tools_to_openai_uses_chat_format(monkeypatch): fake_transform_responses, ) - chat_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( - ["tool"], target_format="chat" - ) + chat_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(["tool"], target_format="chat") resp_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(["tool"]) assert chat_tools == [{"chat": True}] @@ -304,9 +301,7 @@ async def test_execute_tool_calls_strips_prefix_when_alias_differs_from_server_n ) from litellm.proxy._experimental.mcp_server import mcp_server_manager as _msm - _msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock( - return_value=fake_server - ) + _msm.global_mcp_server_manager._get_mcp_server_from_tool_name = MagicMock(return_value=fake_server) tool_name = "my_deepwiki-read_wiki_structure" tool_calls = [ @@ -380,7 +375,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey fake_manager = types.SimpleNamespace( get_registry=MagicMock(return_value={}), - call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")) + call_tool=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom")), ) monkeypatch.setattr( "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", @@ -388,9 +383,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey ) tool_name = "deepwiki-read_wiki_structure" - tool_calls = [ - {"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}} - ] + tool_calls = [{"id": "call-err", "function": {"name": tool_name, "arguments": "{}"}}] user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") @@ -408,10 +401,7 @@ async def test_execute_tool_calls_logs_failure_via_post_call_failure_hook(monkey post_call_failure_hook.assert_awaited_once() assert post_call_failure_hook.await_args is not None - assert ( - post_call_failure_hook.await_args.kwargs.get("route") - == "/responses/mcp/call_tool" - ) + assert post_call_failure_hook.await_args.kwargs.get("route") == "/responses/mcp/call_tool" @pytest.mark.asyncio @@ -434,9 +424,7 @@ async def test_execute_tool_calls_passes_litellm_call_id_and_trace_id_to_functio # NOTE: Don't patch via dotted string path here because `litellm.responses` # is a function attribute on the `litellm` package (shadowing the submodule), # which breaks monkeypatch's importpath resolution. - handler_module = importlib.import_module( - "litellm.responses.mcp.litellm_proxy_mcp_handler" - ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -516,7 +504,9 @@ async def test_execute_tool_calls_applies_post_call_hook_content(monkeypatch): logging_obj = MagicMock() logging_obj.model_call_details = {} - logging_obj.async_post_mcp_tool_call_hook = AsyncMock(return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True)) + logging_obj.async_post_mcp_tool_call_hook = AsyncMock( + return_value=CallToolResult(content=[TextContent(type="text", text="[REDACTED]")], is_error=True) + ) logging_obj.async_success_handler = AsyncMock() handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", lambda *_args, **_kwargs: (logging_obj, None)) @@ -679,9 +669,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch user_auth = types.SimpleNamespace(api_key="test_key", user_id="test_user") tools, _server_names = await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( user_api_key_auth=user_auth, - mcp_tools_with_litellm_proxy=[ - {"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"} - ], + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], ) forwarded: Final = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(tools) @@ -700,9 +688,7 @@ async def test_get_mcp_tools_from_manager_enables_list_tools_logging(monkeypatch def test_get_parent_request_tags_from_metadata(): - tags = LiteLLM_Proxy_MCP_Handler._get_parent_request_tags( - {"metadata": {"tags": ["team-a", "prod"]}} - ) + tags = LiteLLM_Proxy_MCP_Handler._get_parent_request_tags({"metadata": {"tags": ["team-a", "prod"]}}) assert tags == ["team-a", "prod"] @@ -739,9 +725,7 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch): await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( user_api_key_auth=types.SimpleNamespace(api_key="k", user_id="u"), - mcp_tools_with_litellm_proxy=[ - {"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"} - ], + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/deepwiki"}], request_tags=["team-a"], ) @@ -761,9 +745,7 @@ async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(mo captured.update(kwargs) return None, None - handler_module = importlib.import_module( - "litellm.responses.mcp.litellm_proxy_mcp_handler" - ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -789,9 +771,7 @@ async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monk captured.update(kwargs) return None, None - handler_module = importlib.import_module( - "litellm.responses.mcp.litellm_proxy_mcp_handler" - ) + handler_module = importlib.import_module("litellm.responses.mcp.litellm_proxy_mcp_handler") monkeypatch.setattr(handler_module, "function_setup", fake_function_setup) tool_name = "deepwiki-read_wiki_structure" @@ -1171,7 +1151,9 @@ def test_create_follow_up_input_keeps_each_reasoning_item_before_its_function_ca "function_call_output", "function_call_output", ] - assert [cast(dict[str, Any], item).get("id") or cast(dict[str, Any], item).get("call_id") for item in follow_up[1:5]] == [ + assert [ + cast(dict[str, Any], item).get("id") or cast(dict[str, Any], item).get("call_id") for item in follow_up[1:5] + ] == [ "rs_1", "call-1", "rs_2", @@ -1229,16 +1211,20 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false( async def fake_aresponses(**kwargs: Any) -> ResponsesAPIResponse: captured_calls.append(kwargs) - return first_response if len(captured_calls) == 1 else ResponsesAPIResponse( - id="resp_follow_up", - created_at=1234567891, - model="gpt-5", - object="response", - status="completed", - output=[], - parallel_tool_calls=False, - tool_choice="auto", - tools=[], + return ( + first_response + if len(captured_calls) == 1 + else ResponsesAPIResponse( + id="resp_follow_up", + created_at=1234567891, + model="gpt-5", + object="response", + status="completed", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ) ) async def fake_process(**kwargs: Any) -> tuple[list[Any], dict[str, str]]: @@ -1279,12 +1265,14 @@ async def test_mcp_follow_up_call_is_stateless_when_store_is_false( @pytest.mark.asyncio async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: pytest.MonkeyPatch): - from litellm.proxy._experimental.mcp_server import operations - from litellm.proxy._experimental.mcp_server import mcp_server_manager + from litellm.proxy._experimental.mcp_server import mcp_server_manager, operations headers: Final = { - "x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a", - "x-mcp-deepwiki-authorization": "upstream-sentinel", "authorization": "proxy-sentinel", + "x-app-id": "app-a", + "x-nuid": "user-a", + "x-user-id": "identity-a", + "x-mcp-deepwiki-authorization": "upstream-sentinel", + "authorization": "proxy-sentinel", } manager: Final = types.SimpleNamespace( get_registry=MagicMock(return_value={}), @@ -1298,12 +1286,21 @@ async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: py monkeypatch.setattr(operations, "_get_allowed_mcp_servers", AsyncMock(return_value=[])) monkeypatch.setattr(operations, "function_setup", setup) response: Final = ResponsesAPIResponse( - id="resp_test", created_at=1234567891, model="test-model", object="response", - status="completed", output=[], parallel_tool_calls=False, tool_choice="auto", tools=[], + id="resp_test", + created_at=1234567891, + model="test-model", + object="response", + status="completed", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], ) monkeypatch.setattr(responses_main, "aresponses", AsyncMock(return_value=response)) result: Final = await responses_main.aresponses_api_with_mcp( - input="hi", model="test-model", tools=[{"type": "mcp", "server_url": "litellm_proxy"}], + input="hi", + model="test-model", + tools=[{"type": "mcp", "server_url": "litellm_proxy"}], secret_fields={"raw_headers": headers}, ) assert result is response @@ -1311,3 +1308,84 @@ async def test_responses_discovery_logs_sanitized_caller_headers(monkeypatch: py logged: Final = setup.call_args.kwargs["metadata"]["headers"] assert logged == {"x-app-id": "app-a", "x-nuid": "user-a", "x-user-id": "identity-a"} assert headers["x-mcp-deepwiki-authorization"] == "upstream-sentinel" + + +def _toolset_gateway_manager(toolset_id: str, server_id: str) -> types.SimpleNamespace: + return types.SimpleNamespace( + get_registry=MagicMock(return_value={}), + get_allowed_mcp_servers=AsyncMock(return_value=[]), + get_mcp_servers_from_ids=MagicMock(return_value=[]), + get_mcp_server_by_name=MagicMock(return_value=None), + get_toolset_by_name_cached=AsyncMock(return_value=types.SimpleNamespace(toolset_id=toolset_id)), + resolve_toolset_tool_permissions=AsyncMock(return_value={server_id: ["add"]}), + ) + + +async def _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id: str) -> dict[str, object]: + from litellm.proxy._experimental.mcp_server.ui_session_utils import granted_toolset_ids + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth + + mock_get_tools = AsyncMock(return_value=AggregateToolListing(tools=[], outcomes={})) + monkeypatch.setattr("litellm.proxy._experimental.mcp_server.server._get_tools_from_mcp_servers", mock_get_tools) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + _toolset_gateway_manager("ts-granted", "srv-1"), + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + + async def team_permission(context: UserAPIKeyAuth) -> LiteLLM_ObjectPermissionTable: + return LiteLLM_ObjectPermissionTable(object_permission_id="op-team", mcp_toolsets=[team_toolset_id]) + + async def granted_through_team(context: UserAPIKeyAuth) -> frozenset[str]: + return await granted_toolset_ids(context, team_object_permission=team_permission, require_key_access=False) + + team_key: Final = UserAPIKeyAuth(api_key="sk-team", team_id="team-1", user_role=LitellmUserRoles.INTERNAL_USER) + await LiteLLM_Proxy_MCP_Handler._get_mcp_tools_from_manager( + user_api_key_auth=team_key, + mcp_tools_with_litellm_proxy=[{"type": "mcp", "server_url": "litellm_proxy/mcp/team-toolset"}], + granted_toolsets=granted_through_team, + ) + assert mock_get_tools.await_args is not None + return mock_get_tools.await_args.kwargs + + +@pytest.mark.asyncio +async def test_toolset_gateway_url_scopes_a_team_granted_toolset_for_a_key_without_its_own_grant(monkeypatch): + kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-granted") + scoped = kwargs["user_api_key_auth"].object_permission + assert scoped is not None + assert scoped.mcp_servers == ["srv-1"] + assert scoped.mcp_tool_permissions == {"srv-1": ["add"]} + assert kwargs["mcp_servers"] is None + + +@pytest.mark.asyncio +async def test_toolset_gateway_url_skips_a_toolset_the_team_does_not_grant(monkeypatch): + kwargs: Final = await _tools_listing_kwargs_for_toolset_url(monkeypatch, team_toolset_id="ts-other") + assert kwargs["user_api_key_auth"].object_permission is None + assert kwargs["mcp_servers"] is None + + +@pytest.mark.asyncio +async def test_apply_toolset_permissions_pins_the_auth_to_explicit_grants_only(monkeypatch: pytest.MonkeyPatch): + """A toolset gateway URL must not widen to operator-open (allow_all_keys) servers.""" + from litellm.proxy._types import UserAPIKeyAuth + + fake_manager = types.SimpleNamespace( + resolve_toolset_tool_permissions=AsyncMock(return_value={"srv-1": ["add"]}), + ) + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + fake_manager, + ) + + scoped = await LiteLLM_Proxy_MCP_Handler._apply_toolset_permissions( + resolved_toolset_ids=["ts-1"], + resolved_mcp_servers=[], + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", user_id="u1"), + ) + + assert scoped.mcp_explicit_grants_only is True + assert scoped.object_permission is not None + assert scoped.object_permission.mcp_servers == ["srv-1"] + assert scoped.object_permission.mcp_tool_permissions == {"srv-1": ["add"]} diff --git a/tests/unit/test_check_migrations_no_data_rewrites.py b/tests/unit/test_check_migrations_no_data_rewrites.py index c5d3cdd9073..fb1053ff887 100644 --- a/tests/unit/test_check_migrations_no_data_rewrites.py +++ b/tests/unit/test_check_migrations_no_data_rewrites.py @@ -10,6 +10,8 @@ import importlib.util import sys from pathlib import Path +import pytest + _CHECKER_PATH = Path(__file__).resolve().parents[1] / "code_coverage_tests" / "check_migrations_no_data_rewrites.py" _SPEC = importlib.util.spec_from_file_location("check_migrations_no_data_rewrites", _CHECKER_PATH) assert _SPEC is not None and _SPEC.loader is not None @@ -205,6 +207,96 @@ class TestDefaultedColumnsOnRequestLogTables: assert 'ADD COLUMN ... DEFAULT on "LiteLLM_SpendLogs" rewrites existing rows at boot' in rendered +class TestIndexesOnLogTables: + """Every CREATE INDEX on a request-log table is rejected: a plain one blocks writes for + the whole build and a concurrent one fails on a partitioned parent, so the migration job + (litellm_proxy_extras/request_log_indexes.py) builds those instead.""" + + def test_the_original_spend_log_index_statement_is_flagged(self, tmp_path): + sql = ( + "-- CreateIndex\n" + 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ' + 'ON "LiteLLM_SpendLogs"("api_key", "startTime");' + ) + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + + def test_the_original_concurrent_call_id_index_statement_is_flagged(self, tmp_path): + sql = ( + 'CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ' + 'ON "LiteLLM_SpendLogs"("litellm_call_id");' + ) + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + + def test_unique_index_with_if_not_exists_on_error_logs_is_flagged(self, tmp_path): + sql = 'CREATE UNIQUE INDEX IF NOT EXISTS "ix" ON "LiteLLM_ErrorLogs" ("request_id");' + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_ErrorLogs"',) + + def test_a_unique_concurrent_index_on_error_logs_is_flagged(self, tmp_path): + sql = 'CREATE UNIQUE INDEX CONCURRENTLY "ix" ON "LiteLLM_ErrorLogs" ("request_id");' + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_ErrorLogs"',) + + def test_lowercase_schema_qualified_and_only_forms_are_flagged(self, tmp_path): + sql = ( + 'create index on "public"."LiteLLM_SpendLogs" ("api_key");\n' + 'CREATE INDEX "ix" ON ONLY "LiteLLM_SpendLogs" ("api_key");\n' + 'CREATE INDEX CONCURRENTLY "iy" ON "public"."LiteLLM_SpendLogs" ("api_key");' + ) + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) * 3 + + def test_a_comment_between_on_and_the_table_is_flagged(self, tmp_path): + sql = 'CREATE INDEX CONCURRENTLY "ix" ON /* table */ "LiteLLM_SpendLogs" ("api_key");' + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + + def test_a_concurrent_index_with_comments_and_line_breaks_is_flagged(self, tmp_path): + sql = ( + "-- CreateIndex\n" + 'CREATE INDEX CONCURRENTLY IF NOT EXISTS "ix"\n' + ' ON "LiteLLM_SpendLogs" /* partitioned in some deployments */\n' + ' ("api_key", "startTime");\n' + ) + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + + def test_indexes_on_a_non_log_table_pass_concurrent_or_not(self, tmp_path): + sql = ( + 'CREATE INDEX "ix" ON "LiteLLM_VerificationToken" ("token");\n' + 'CREATE INDEX CONCURRENTLY "iy" ON "LiteLLM_VerificationToken" ("token");' + ) + assert _keywords(tmp_path, sql) == () + + def test_an_index_run_by_execute_is_flagged(self, tmp_path): + sql = 'DO $$ BEGIN EXECUTE \'CREATE INDEX "ix" ON "LiteLLM_SpendLogs" ("api_key")\'; END $$;' + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + + def test_a_marker_does_not_exempt_the_index(self, tmp_path): + sql = ( + '-- data-migration-ok: table is empty at this point\nCREATE INDEX "ix" ON "LiteLLM_SpendLogs" ("api_key");' + ) + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + + def test_a_marker_on_a_rewrite_still_leaves_the_index_below_it_flagged(self, tmp_path): + sql = ( + "-- data-migration-ok: one row\n" + 'UPDATE "LiteLLM_SpendLogs" SET "api_key" = \'k\' WHERE "request_id" = \'r\';\n' + 'CREATE INDEX CONCURRENTLY "ix" ON "LiteLLM_SpendLogs" ("api_key");' + ) + assert _keywords(tmp_path, sql) == ('CREATE INDEX on "LiteLLM_SpendLogs"',) + + def test_render_points_at_the_migration_job_index_list(self, tmp_path): + sql = 'CREATE INDEX CONCURRENTLY "ix" ON "public"."LiteLLM_SpendLogs" ("api_key");' + rendered = _scan(tmp_path, sql)[0].render() + assert "20260101000000_fixture/migration.sql:1" in rendered + assert "blocks writes until the build finishes, or fails on a partitioned table" in rendered + assert "REQUEST_LOG_INDEXES in litellm_proxy_extras/request_log_indexes.py" in rendered + + @pytest.mark.parametrize( + "name", + ("20260823000000_add_spend_logs_api_key_starttime_index", "20260831120001_spend_logs_litellm_call_id_index"), + ) + def test_the_inert_index_migrations_scan_clean_without_a_grandfather(self, name): + assert checker.scan_migration(checker.MIGRATIONS_DIR / name) == () + assert name not in checker.GRANDFATHERED + + class TestInsert: def test_insert_values_is_bounded_and_passes(self, tmp_path): assert _keywords(tmp_path, "INSERT INTO \"Foo\" (\"id\") VALUES ('a'), ('b');") == () diff --git a/ui/litellm-dashboard/public/assets/agent-traces-preview.png b/ui/litellm-dashboard/public/assets/agent-traces-preview.png deleted file mode 100644 index 34569e26331..00000000000 Binary files a/ui/litellm-dashboard/public/assets/agent-traces-preview.png and /dev/null differ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx index 282262b6322..26726691a48 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/lens/page.tsx @@ -37,7 +37,12 @@ export default function LensPage() { - + {isProxyAdminTierRole(userRole ?? "") ? ( diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 97d1782b08c..da902d4a9e2 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2122,6 +2122,9 @@ export const agentTraceListCall = async ({ return apiClient.get(`/v1/traces`, { accessToken, query }); }; +export const sendOtlpTraceCall = async (accessToken: string, exportRequest: object): Promise => + apiClient.post(`/v1/traces`, { accessToken, body: exportRequest }); + export const agentTraceCall = async (accessToken: string, traceId: string, traceRef?: string): Promise => apiClient.get(`/v1/traces/${encodeURIComponent(traceId)}`, { accessToken, diff --git a/ui/litellm-dashboard/src/components/view_logs/TraceView/ActiveDot.tsx b/ui/litellm-dashboard/src/components/view_logs/TraceView/ActiveDot.tsx new file mode 100644 index 00000000000..b3d4cf55482 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/TraceView/ActiveDot.tsx @@ -0,0 +1,8 @@ +export function ActiveDot() { + return ( +