scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

refactor(workflows): reuse gateway owners and require explicit opt-in

Browse source
This commit is contained in:
ishaan-j-berri 2026-11-06 18:43:48 +08:00
parent 7891e30a74
commit bb147ec202
30 changed files with 1418 additions and 893 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/litellm_core_utils/prompt_templates/json_schema_utils.py
View file

@ -239,7 +239,7 @@ def _schema_levels(schema: Mapping[str, object]) -> Iterator[tuple[Mapping[str,
frontier: tuple[Mapping[str, object], ...] = (schema,) # rebind-ok: level-order cursor, one level a round
while frontier:
yield frontier
frontier = tuple(child for node in frontier for child in _subschemas(node))
frontier = tuple(child for node in frontier for child in schema_children(node))
def schema_fragment_ref(ref: str) -> str:
@ -359,7 +359,8 @@ def _is_open_object(node: Mapping[str, object]) -> bool:
return is_object_schema(node) and "properties" not in node and additional_properties is not False
def _subschemas(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]:
def schema_children(node: Mapping[str, object]) -> Iterator[Mapping[str, object]]:
"""Yield schema-bearing children, excluding annotations and literal JSON values."""
for key, value in node.items():
if key in _SUBSCHEMA_MAP_KEYWORDS and isinstance(value, dict):
yield from (sub for sub in value.values() if isinstance(sub, dict))

2
litellm/models/workflow_definition.py
View file

@ -167,4 +167,4 @@ class WorkflowDefinitionStore(Protocol):
self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str
) -> WorkflowSnapshot | DefinitionError: ...
async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError: ...
async def prepare_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError: ...

3
litellm/models/workflow_execution.py
View file

@ -7,6 +7,8 @@ from uuid import NAMESPACE_URL, uuid5
from pydantic import BaseModel, ConfigDict, Field, JsonValue
from litellm.models.workflow_definition import WorkflowSnapshot
class RunState(str, Enum):
PENDING = "pending"
@ -142,6 +144,7 @@ class WorkflowExecutionStore(Protocol):
input: RunInput,
digest: str,
idempotency_key: str | None,
draft_snapshot: WorkflowSnapshot | None = None,
) -> RunAdmission | ExecutionError: ...
async def get(self, run_id: str) -> EngineRun | None: ...

4
litellm/proxy/_types.py
View file

@ -2601,6 +2601,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
None,
description="opt-in to copy the x-litellm-call-id response header's value into JSON error bodies, as error.litellm_call_id on the OpenAI-shaped and /v1/messages routes and as a top-level litellm_call_id on pass-through routes, so an error a client prints names the request to look up. Off by default",
)
enable_workflows: bool = Field(
default=False,
description="Enable LiteFlow workflow authoring and execution. Off by default; restart the proxy after changing this setting.",
)
enable_claude_code_gateway: bool | None = Field(
None,
description="serve the Claude Code gateway protocol (https://code.claude.com/docs/en/claude-apps-gateway) under /claude_code_gateway: OAuth device-flow sign-in reusing proxy SSO, plus managed settings and OTLP telemetry ingestion. Off by default",

28
litellm/proxy/auth/user_api_key_auth.py
View file

@ -4873,7 +4873,7 @@ async def _run_post_custom_auth_checks(
async def authorize_internal_virtual_key(
key_hash: str,
request: Request,
request_data: dict[str, object],
request_data: dict[str, object], # mutable-ok: existing auth enriches this attempt-owned request body
*,
route: str | None = None,
expected_team_id: str | None = None,
@ -4890,6 +4890,28 @@ async def authorize_internal_virtual_key(
raise HTTPException(status_code=403, detail="The execution key no longer belongs to its original team")
execution_route: Final = route if route is not None else get_request_route(request=request)
await pre_db_read_auth_checks(request_data=request_data, request=request, route=execution_route)
identity.end_user_id = await resolve_and_validate_end_user_id(
raw_end_user_id=identity.end_user_id
or get_end_user_id_from_request_body(request_data, safe_get_request_headers(request)),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,
proxy_logging_obj=proxy_logging_obj,
route=execution_route,
key_end_user_budget_id=get_key_end_user_budget_id(identity.metadata),
)
end_user_object: Final[LiteLLM_EndUserTable | None]
if identity.end_user_id is not None:
_, end_user_object = await _lookup_end_user_and_apply_budget(
valid_token=identity,
route=execution_route,
parent_otel_span=None,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
else:
end_user_object = None
auth: Final = await validate_resolved_virtual_key(
request=request,
request_data=request_data,
@ -4898,9 +4920,9 @@ async def authorize_internal_virtual_key(
route=execution_route,
start_time=datetime.now(timezone.utc),
parent_otel_span=None,
end_user_id=None,
end_user_id=identity.end_user_id,
end_user_params={},
_end_user_object=None,
_end_user_object=end_user_object,
)
auth.budget_reservation = None
recovered: Final = await _authorize_authenticated_request(

22
litellm/proxy/deps/repositories.py
View file

@ -1,4 +1,4 @@
from typing import Annotated, TypeAlias, cast # noqa: TID251 # Prisma writer delegates are generated dynamically
from typing import Annotated, TypeAlias
from fastapi import Depends
@ -15,8 +15,8 @@ from litellm.repositories.gateway_activity_repository import GatewayActivityRepo
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
from litellm.repositories.skill_submission_repository import SkillSubmissionRepository
from litellm.repositories.user_banner_repository import UserBannerRepository
from litellm.repositories.workflow_definition_repository import DefinitionDatabase, WorkflowDefinitionRepository
from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from litellm.repositories.workflow_execution_repository import WorkflowExecutionRepository
async def get_user_banner_repository(database: DatabaseDep) -> UserBannerStore | None:
@ -61,13 +61,7 @@ SkillSubmissionRepositoryDep: TypeAlias = Annotated[
async def get_workflow_definition_repository(database: DatabaseDep) -> WorkflowDefinitionStore | None:
return (
None
if database is None
else WorkflowDefinitionRepository(
cast(DefinitionDatabase, database.prisma.writer_db) # cast-ok: Prisma delegates implement the protocol
)
)
return None if database is None else WorkflowDefinitionRepository(database.prisma, use_writer=True)
WorkflowDefinitionRepositoryDep: TypeAlias = Annotated[
@ -76,13 +70,7 @@ WorkflowDefinitionRepositoryDep: TypeAlias = Annotated[
async def get_workflow_execution_repository(database: DatabaseDep) -> WorkflowExecutionStore | None:
return (
None
if database is None
else WorkflowExecutionRepository(
cast(ExecutionDatabase, database.prisma.writer_db) # cast-ok: generated Prisma transaction delegates
)
)
return None if database is None else WorkflowExecutionRepository(database.prisma, use_writer=True)
WorkflowExecutionRepositoryDep: TypeAlias = Annotated[

5
litellm/proxy/management_endpoints/workflow_definition_endpoints.py
View file

@ -1,3 +1,5 @@
"""Gateway lifecycle APIs for workflows authored with https://github.com/BerriAI/liteflow."""
from typing import Annotated, Final, Literal, NoReturn, TypeAlias, TypeVar
from fastapi import APIRouter, Depends, Header, HTTPException, Query
@ -24,8 +26,9 @@ from litellm.proxy.workflows.definitions import (
requested_team,
validate_bundle,
)
from litellm.proxy.workflows.lifespan import require_workflows_enabled
router: Final = APIRouter(tags=["workflow management"])
router: Final = APIRouter(tags=["workflow management"], dependencies=[Depends(require_workflows_enabled)])
_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
_Team: TypeAlias = Annotated[str | None, Query(min_length=1, max_length=256)]
_ETag: TypeAlias = Annotated[str | None, Header(alias="If-Match")]

23
litellm/proxy/management_endpoints/workflow_execution_endpoints.py
View file

@ -28,14 +28,15 @@ from litellm.proxy.management_endpoints.workflow_definition_endpoints import rai
from litellm.proxy.workflows.definitions import (
authorize_team,
authorized_draft,
freeze_draft,
prepare_draft,
published_workflow,
)
from litellm.proxy.workflows.lifespan import require_workflows_enabled
from litellm.proxy.workflows.monty.models import MontyFailure
from litellm.proxy.workflows.monty.runtime import validate_json
from litellm.proxy.workflows.runner import utc_now
router: Final = APIRouter(tags=["workflow management"])
router: Final = APIRouter(tags=["workflow management"], dependencies=[Depends(require_workflows_enabled)])
_Caller: TypeAlias = Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]
_Idempotency: TypeAlias = Annotated[str | None, Header(alias="Idempotency-Key", min_length=1, max_length=256)]
_MESSAGES: Final = TypeAdapter(list[ChatCompletionMessageParam])
@ -153,7 +154,6 @@ async def _admit(
*,
draft: bool,
) -> EngineRun:
_validate_messages(messages)
admitted: Final = await runs.admit(
workflow_id=snapshot.workflow_id,
snapshot_id=snapshot.snapshot_id,
@ -167,6 +167,7 @@ async def _admit(
),
digest=digest,
idempotency_key=key,
draft_snapshot=snapshot if draft else None,
)
if isinstance(admitted, ExecutionError):
_error(admitted)
@ -206,12 +207,13 @@ async def test_workflow(
if isinstance(permission, DefinitionError):
raise_public(permission)
key_hash: Final = _execution_key(caller, draft.team_id)
_validate_messages(data.messages)
digest: Final = _digest(data.messages, caller.end_user_id, if_match)
key: Final = _idempotency_key("test", idempotency_key)
previous: Final = await _recover_admission(runs, workflow_id, draft.team_id, key_hash, key, digest)
if previous is not None:
return _accepted(previous)
snapshot: Final = await freeze_draft(store, caller, workflow_id, if_match)
snapshot: Final = await prepare_draft(store, caller, workflow_id, if_match)
if isinstance(snapshot, DefinitionError):
raise_public(snapshot)
return _accepted(await _admit(runs, snapshot, caller, data.messages, caller.end_user_id, key, digest, draft=True))
@ -223,8 +225,8 @@ async def _read_run(
run: Final = await runs.get(run_id)
if run is None:
_error(ExecutionError("not_found", "Run not found"))
team: Final = await store.get_team(run.team_id)
if team is None or team.blocked:
team: Final = await authorize_team(store, caller, run.team_id, "invoke")
if isinstance(team, DefinitionError):
_error(ExecutionError("not_found", "Run not found"))
owns_run: Final = bool(caller.token) and caller.token == run.created_by and caller.team_id == run.team_id
if not owns_run and not user_api_key_has_admin_view(caller) and not is_team_admin(caller, team):
@ -280,9 +282,10 @@ async def cancel_workflow_run(
"""Request cancellation and return the run ID to poll until execution stops."""
store, runs = _stores(definitions, executions)
run: Final = await _read_run(store, runs, caller, run_id)
permission: Final = await authorize_team(store, caller, run.team_id, "write")
if isinstance(permission, DefinitionError):
raise_public(permission)
if caller.token != run.created_by:
permission: Final = await authorize_team(store, caller, run.team_id, "write")
if isinstance(permission, DefinitionError):
raise_public(permission)
await runs.cancel(run_id)
return _accepted(run)
@ -298,6 +301,7 @@ async def list_engine_runs(
definitions: WorkflowDefinitionStore | None,
executions: WorkflowExecutionStore | None,
) -> RunPage:
require_workflows_enabled()
store, runs = _stores(definitions, executions)
definition: Final = await store.get_draft(workflow_id, team_id)
if definition is None:
@ -377,6 +381,7 @@ async def invoke_workflow(
key: Final = request.headers.get("Idempotency-Key")
if key is not None and not 1 <= len(key) <= 256:
_error(ExecutionError("invalid_idempotency_key", "Idempotency-Key must contain 1 to 256 characters"))
_validate_messages(data.messages)
admitted: Final = await _admit(
runs,
snapshot,

2
litellm/proxy/proxy_server.py
View file

@ -13107,7 +13107,9 @@ async def chat_completion(
get_workflow_execution_repository,
)
from litellm.proxy.management_endpoints.workflow_execution_endpoints import invoke_workflow
from litellm.proxy.workflows.lifespan import require_workflows_enabled
require_workflows_enabled()
workflow_database: Final = await get_database(request)
try:
return await invoke_workflow(

6
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -267,7 +267,11 @@ async def reserve_budget_for_request(
return None
if RouteChecks.is_unbilled_route(route):
return None
if get_model_from_request(request_body, route, llm_router=llm_router) is None:
model: Final = get_model_from_request(request_body, route, llm_router=llm_router)
if model is None:
return None
if isinstance(model, str) and model.startswith("workflow/"):
# Workflow wrappers do not bill; each child reserves its own provider call.
return None
counters: Final = await _get_budget_counters(

87
litellm/proxy/workflows/LITEFLOW.md
View file

@ -1,87 +0,0 @@
# LiteFlow gateway workflows
The gateway accepts LiteFlow compiled contract version `1` and executes sequential `chat_completion` and `python` nodes. Python nodes use the Monty adapter. Conditional, parallel and streaming execution are not supported
Workflow definitions belong to a team. Authors register and edit drafts. A team administrator who did not contribute to the draft publishes an immutable snapshot. Invocation requires a virtual key assigned to that team, with access to the workflow model and each child model
## SDK example
This contract matches LiteFlow PRs [3](https://github.com/BerriAI/liteflow/pull/3) and [5](https://github.com/BerriAI/liteflow/pull/5). Configure `gateway-model` on the proxy and provide two different team members' keys
```python
import os
import time
from liteflow import ChatCompletion, GatewayClient, Input, Workflow
workflow = Workflow(
name="welcome",
nodes=[ChatCompletion(model="gateway-model", messages=Input.messages)],
)
messages = [{"role": "user", "content": "Write one short welcome sentence."}]
with GatewayClient(
base_url=os.environ["PROXY_BASE_URL"],
api_key=os.environ["WORKFLOW_AUTHOR_KEY"],
team_id=os.environ["WORKFLOW_TEAM_ID"],
) as author, GatewayClient(
base_url=os.environ["PROXY_BASE_URL"],
api_key=os.environ["WORKFLOW_PUBLISHER_KEY"],
team_id=os.environ["WORKFLOW_TEAM_ID"],
) as publisher:
draft = author.register(workflow)
author.validate(draft)
publisher.publish(draft)
accepted = author.test(draft, messages=messages, authorize_billing=True)
for _ in range(300):
trace = author.inspect(accepted["run_id"])
print(trace)
if trace.state.value in ("completed", "failed", "cancelled"):
break
time.sleep(1)
result = author.run("welcome", messages=messages)
print(result.run_id, result.output)
```
## API contract
| Operation | Endpoint | Contract |
| --- | --- | --- |
| Register | `POST /v1/workflows` | `{"bundle": <compiled bundle>}`, optional `team_id` query, returns `201` |
| List definitions | `GET /v1/workflows` | `team_id`, `limit`, `cursor`; returns `workflows`, `next_cursor` |
| Retrieve draft | `GET /v1/workflows/{workflow_id}` | Returns the draft and `ETag` header |
| Replace draft | `PUT /v1/workflows/{workflow_id}` | Compiled bundle and exact `If-Match` |
| Validate | `POST /v1/workflows/{workflow_id}/validate` | Exact `If-Match`, no body required |
| Publish | `POST /v1/workflows/{workflow_id}/publish` | Exact `If-Match` and `{"digest": "<draft digest>"}` |
| Test draft | `POST /v1/workflows/{workflow_id}/test` | Exact `If-Match`, `messages`, `authorize_billing: true`; returns `202` |
| Invoke published | `POST /v1/chat/completions` | `model: "workflow/<name>"`, `messages`, `stream: false` |
| List runs | `GET /v1/workflows/{workflow_id}/runs` | `limit`, `status`, `after`; returns `runs`, `count`, `next_cursor` |
| Inspect | `GET /v1/workflows/runs/{run_id}/steps` | Run state and typed step traces |
| Pinned definition | `GET /v1/workflows/runs/{run_id}/definition` | Exact snapshot used by that run |
| Cancel | `POST /v1/workflows/runs/{run_id}/cancel` | Returns `202`; poll until terminal state |
Draft responses contain `name`, `workflow_id`, `team_id`, `etag`, `digest`, `bundle` and `published_snapshot_id`. Missing or stale `If-Match` returns `412`. Validation checks the compiled graph and Python program without executing model or Python nodes
Test runs and chat invocations accept `Idempotency-Key`. Retrying the same operation with the same key and input returns the original run. Reusing that key with different input returns `409`. A test run pins the exact draft revision even if the draft changes later
For live UI progress, send `Prefer: respond-async` on a published chat invocation. The gateway returns `202` with `run_id`, `snapshot_id`, `trace_url` and the `x-litellm-workflow-run-id` header. The ordinary SDK call waits for a Chat Completions response carrying the same run-ID header. A synchronous wait that reaches 60 seconds returns `504` with the run ID so the caller can continue polling
## Trace and spend data
Draw the pinned bundle's `nodes` and resolve each binding's `from` reference. Match `step.node_id` to `node.id`. Run and node states are `pending`, `running`, `completed`, `failed` and `cancelled`
Each step exposes typed `input` and `output`, `attempts`, `started_at`, `finished_at`, `deadline`, sanitized `error`, token counts and `response_cost`. A run-level `error` explains failures before any node starts
Model steps carry a deterministic `call_id` before dispatch. `request_id` and `request_ids` link to actual `LiteLLM_SpendLogs.request_id` values. Spend writes are asynchronous, so receipt IDs can be absent on the first read and appear on later polls. The trace reader sums persisted attempts once and does not duplicate earlier checkpoints
Each child uses the standard gateway authorization, hooks, guardrails, limits and billing path. The workflow wrapper emits no additional bill. Python nodes have no model spend receipt. Known costs remain visible when output validation fails
Definitions and snapshots use separate tables. Execution reuses `LiteLLM_WorkflowRun` and `LiteLLM_WorkflowEvent`. Existing tracking mutations cannot overwrite engine-owned runs
## Execution limits and recovery
The initial runner allows four active runs per worker and 64 pending or running runs per team. Database leases prevent two workers from executing the same run. Completed node checkpoints survive a restart. Recovery fails an interrupted model attempt whose outcome is uncertain instead of replaying a potentially billable request
Cancellation is cooperative. It stops downstream nodes but cannot undo a model request already sent. Trace receipts can still arrive after cancellation
Ordinary virtual keys can read their own runs. Team administrators and proxy administrators can inspect team runs. Definition access does not grant another caller access to private run inputs or outputs

23
litellm/proxy/workflows/definitions.py
View file

@ -6,6 +6,7 @@ from jsonschema.exceptions import SchemaError
from openai.types.chat import ChatCompletionMessageParam
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm.litellm_core_utils.prompt_templates.json_schema_utils import schema_children
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.workflow_definition import (
Binding,
@ -112,7 +113,7 @@ def check_etag(draft: WorkflowDraft, etag: str | None) -> DefinitionError | None
return None
async def freeze_draft(
async def prepare_draft(
store: WorkflowDefinitionStore,
caller: UserAPIKeyAuth,
workflow_id: str,
@ -129,22 +130,18 @@ async def freeze_draft(
return failure
if (invalid := validate_bundle(draft.bundle)) is not None:
return invalid
return await store.freeze_draft(workflow_id, draft.team_id, draft.etag)
return await store.prepare_draft(workflow_id, draft.team_id, draft.etag)
def _local_references(value: JsonValue | Mapping[str, JsonValue], depth: int = 0) -> bool:
def _local_references(schema: Mapping[str, object], depth: int = 0) -> bool:
if depth > 64:
return False
if isinstance(value, Mapping):
if any(keyword in value for keyword in ("$id", "$dynamicRef", "$recursiveRef", "$dynamicAnchor")):
return False
reference: Final = value.get("$ref")
if reference is not None and (not isinstance(reference, str) or not reference.startswith("#/")):
return False
return all(_local_references(child, depth + 1) for child in value.values())
if isinstance(value, list):
return all(_local_references(child, depth + 1) for child in value)
return True
if any(keyword in schema for keyword in ("$id", "$dynamicRef", "$recursiveRef", "$dynamicAnchor")):
return False
reference: Final = schema.get("$ref")
if reference is not None and (not isinstance(reference, str) or not reference.startswith("#/")):
return False
return all(_local_references(child, depth + 1) for child in schema_children(schema))
def _schema_error(schema: Mapping[str, JsonValue], param: str) -> DefinitionError | None:

15
litellm/proxy/workflows/lifespan.py
View file

@ -2,6 +2,8 @@ from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import Final
from fastapi import HTTPException
from litellm.proxy.deps.repositories import (
get_workflow_definition_repository,
get_workflow_execution_repository,
@ -11,9 +13,20 @@ from litellm.proxy.workflows.runner import WorkflowRunner
from litellm.repositories.database import Database
def workflows_enabled() -> bool:
from litellm.proxy.proxy_server import general_settings_view
return general_settings_view().get("enable_workflows") is True
def require_workflows_enabled() -> None:
if not workflows_enabled():
raise HTTPException(status_code=404, detail="Workflows are not enabled")
@asynccontextmanager
async def workflow_lifespan(database: Database | None) -> AsyncGenerator[None]:
if database is None:
if database is None or not workflows_enabled():
yield
return
definitions: Final = await get_workflow_definition_repository(database)

41
litellm/proxy/workflows/runner.py
View file

@ -95,6 +95,14 @@ async def _settle_task(task: asyncio.Task[object]) -> None:
await asyncio.gather(task, return_exceptions=True)
def _terminal_state(result: JsonValue | ExecutionError) -> RunState:
if not isinstance(result, ExecutionError):
return RunState.COMPLETED
if result.code == "cancelled":
return RunState.CANCELLED
return RunState.FAILED
class WorkflowRunner:
def __init__(
self,
@ -145,17 +153,23 @@ class WorkflowRunner:
async def _heartbeat(self, run: EngineRun, token: str) -> ExecutionError:
while True:
await asyncio.sleep(1)
current: EngineRun | None = await self.runs.get(run.run_id)
if current is None or current.lease_token != token:
return ExecutionError("lease_lost", "Execution ownership changed")
if current.cancel_requested:
return ExecutionError("cancelled", "Cancellation requested")
if self._clock() >= run.input.deadline:
return ExecutionError("deadline_exceeded", "Workflow deadline expired")
interrupted: ExecutionError | None = await self._interruption(run, token)
if interrupted is not None:
return interrupted
renewed: bool = await self.runs.renew(run.run_id, token, self._clock() + timedelta(seconds=15))
if not renewed:
return ExecutionError("lease_lost", "Execution ownership expired")
async def _interruption(self, run: EngineRun, token: str) -> ExecutionError | None:
current: Final = await self.runs.get(run.run_id)
if current is None or current.lease_token != token:
return ExecutionError("lease_lost", "Execution ownership changed")
if current.cancel_requested:
return ExecutionError("cancelled", "Cancellation requested")
if self._clock() >= run.input.deadline:
return ExecutionError("deadline_exceeded", "Workflow deadline expired")
return None
async def _run(self, run_id: str) -> None:
token: Final = str(uuid4())
run: Final = await self.runs.claim(run_id, token, self._clock() + timedelta(seconds=15))
@ -173,7 +187,7 @@ class WorkflowRunner:
await self.runs.finish(
run_id,
token,
RunState.FAILED if isinstance(result, ExecutionError) else RunState.COMPLETED,
_terminal_state(result),
None if isinstance(result, ExecutionError) else result,
result.message if isinstance(result, ExecutionError) else None,
)
@ -183,7 +197,7 @@ class WorkflowRunner:
await self.runs.finish(
run_id,
token,
RunState.CANCELLED if reason.code == "cancelled" else RunState.FAILED,
_terminal_state(reason),
None,
reason.message,
)
@ -216,8 +230,9 @@ class WorkflowRunner:
async def _sequence(
self, run: EngineRun, token: str, bundle: WorkflowBundle, steps: tuple[StepTrace, ...], index: int
) -> JsonValue | ExecutionError:
if self._clock() >= run.input.deadline:
return ExecutionError("deadline_exceeded", "Workflow deadline expired")
interrupted: Final = await self._interruption(run, token)
if interrupted is not None:
return interrupted
if index == len(bundle.nodes):
return resolve_binding(bundle.output, run, steps)
node: Final = bundle.nodes[index]
@ -252,7 +267,9 @@ class WorkflowRunner:
self, run: EngineRun, token: str, node: WorkflowNode, step: StepTrace
) -> StepTrace | ExecutionError:
if not await self.runs.checkpoint(run.run_id, token, step):
return ExecutionError("lease_lost", "Execution ownership expired before dispatch")
return await self._interruption(run, token) or ExecutionError(
"lease_lost", "Execution ownership expired before dispatch"
)
result: Final = await self._recorded_node_result(run, token, node, step)
receipt: Final = result.receipt if isinstance(result, ExecutionError) else result
finished: Final = step.model_copy(

20
litellm/repositories/table_repositories.py
View file

@ -10,7 +10,7 @@ methods; richer repositories live in their own modules.
from collections.abc import Mapping, Sequence
from datetime import timedelta
from itertools import groupby
from typing import TYPE_CHECKING, Any, Final, Generic
from typing import TYPE_CHECKING, Any, Final, Generic, cast # noqa: TID251 # generated Prisma writer boundary
from pydantic import TypeAdapter
@ -19,8 +19,11 @@ from litellm.repositories.accounting_delivery_repository import LedgerDatabase,
from litellm.repositories.prisma_protocols import RowT_co, TableActions
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models # noqa: F401 # used by quoted base-class subscripts
from litellm.proxy.utils import PrismaClient
class PrismaTableRepository(Generic[RowT_co]):
"""Base for repositories that expose a single Prisma table."""
@ -37,6 +40,17 @@ class PrismaTableRepository(Generic[RowT_co]):
raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
return self._prisma_client
@property
def writer_db(self) -> "Prisma":
client: Final = cast("PrismaClient", self._prisma_client) # cast-ok: legacy client boundary
return cast("Prisma", client.writer_db) # cast-ok: PrismaWrapper forwards generated client methods
@classmethod
def transaction_table(cls, transaction: "Prisma") -> TableActions[RowT_co]:
return cast( # cast-ok: owner binds delegate row type
"TableActions[RowT_co]", getattr(transaction, cls.table_name)
)
@property
def table(self) -> TableActions[RowT_co]:
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
@ -285,6 +299,10 @@ class DeletedVerificationTokenRepository(PrismaTableRepository["prisma_models.Li
table_name = "litellm_deletedverificationtoken"
class WorkflowSnapshotRepository(PrismaTableRepository["prisma_models.LiteLLM_WorkflowSnapshot"]):
table_name = "litellm_workflowsnapshot"
class WorkflowRunRepository(PrismaTableRepository["prisma_models.LiteLLM_WorkflowRun"]):
table_name = "litellm_workflowrun"

259
litellm/repositories/workflow_definition_repository.py
View file

@ -1,12 +1,9 @@
import hashlib
import json
from contextlib import AbstractAsyncContextManager
from datetime import datetime
from typing import Final, Protocol
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Final
from uuid import uuid4
from pydantic import BaseModel, ConfigDict
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.workflow_definition import (
DefinitionError,
@ -17,78 +14,40 @@ from litellm.models.workflow_definition import (
WorkflowSnapshot,
)
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import PrismaTableRepository, TeamRepository, WorkflowSnapshotRepository
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
class DefinitionTransaction(Protocol):
@property
def litellm_workflowdefinition(self) -> TableActions[object]: ...
@property
def litellm_workflowsnapshot(self) -> TableActions[object]: ...
@property
def litellm_teamtable(self) -> TableActions[object]: ...
async def query_raw(self, query: str, *args: object) -> object: ...
def _definition(row: "prisma_models.LiteLLM_WorkflowDefinition") -> WorkflowDefinition:
return WorkflowDefinition(
draft=WorkflowDraft(
workflow_id=row.workflow_id,
name=row.name,
team_id=row.team_id,
etag=row.etag,
digest=row.digest,
bundle=WorkflowBundle.model_validate_json(row.bundle_json),
published_snapshot_id=row.published_snapshot_id,
),
contributors=tuple(row.contributors),
)
class DefinitionDatabase(DefinitionTransaction, Protocol):
def tx(self) -> AbstractAsyncContextManager[DefinitionTransaction]: ...
class _DefinitionRow(BaseModel):
model_config = ConfigDict(frozen=True, from_attributes=True)
workflow_id: str
name: str
team_id: str
etag: str
digest: str
bundle_json: str
contributors: tuple[str, ...]
published_snapshot_id: str | None
published_etag: str | None
def definition(self) -> WorkflowDefinition:
return WorkflowDefinition(
draft=WorkflowDraft(
workflow_id=self.workflow_id,
name=self.name,
team_id=self.team_id,
etag=self.etag,
digest=self.digest,
bundle=WorkflowBundle.model_validate_json(self.bundle_json),
published_snapshot_id=self.published_snapshot_id,
),
contributors=self.contributors,
)
class _SnapshotRow(BaseModel):
model_config = ConfigDict(frozen=True, from_attributes=True)
snapshot_id: str
workflow_id: str
team_id: str
name: str
digest: str
bundle_json: str
contributors: tuple[str, ...]
published_by: str | None
created_at: datetime
def snapshot(self) -> WorkflowSnapshot:
return WorkflowSnapshot(
snapshot_id=self.snapshot_id,
workflow_id=self.workflow_id,
team_id=self.team_id,
name=self.name,
digest=self.digest,
bundle=WorkflowBundle.model_validate_json(self.bundle_json),
contributors=self.contributors,
published_by=self.published_by,
created_at=self.created_at,
)
def _snapshot(row: "prisma_models.LiteLLM_WorkflowSnapshot") -> WorkflowSnapshot:
return WorkflowSnapshot(
snapshot_id=row.snapshot_id,
workflow_id=row.workflow_id,
team_id=row.team_id,
name=row.name,
digest=row.digest,
bundle=WorkflowBundle.model_validate_json(row.bundle_json),
contributors=tuple(row.contributors),
published_by=row.published_by,
created_at=row.created_at,
)
def bundle_digest(bundle: WorkflowBundle) -> str:
@ -110,88 +69,100 @@ def _etag() -> str:
return f'"{uuid4()}"'
def _prepared_snapshot(row: "prisma_models.LiteLLM_WorkflowDefinition", published_by: str | None) -> WorkflowSnapshot:
return WorkflowSnapshot(
snapshot_id=str(uuid4()),
workflow_id=row.workflow_id,
team_id=row.team_id,
name=row.name,
digest=row.digest,
bundle=WorkflowBundle.model_validate_json(row.bundle_json),
contributors=tuple(row.contributors),
published_by=published_by,
created_at=datetime.now(timezone.utc),
)
async def persist_snapshot(transaction: "Prisma", snapshot: WorkflowSnapshot) -> WorkflowSnapshot:
"""Persist a prepared definition in the caller's publication/admission transaction."""
data: Final = {
"snapshot_id": snapshot.snapshot_id,
"workflow_id": snapshot.workflow_id,
"team_id": snapshot.team_id,
"name": snapshot.name,
"digest": snapshot.digest,
"bundle_json": _bundle_json(snapshot.bundle),
"contributors": list(snapshot.contributors),
"published_by": snapshot.published_by,
"created_at": snapshot.created_at,
}
stored: Final = await WorkflowSnapshotRepository.transaction_table(transaction).create(data=data)
return _snapshot(stored)
async def _locked_definition(
transaction: DefinitionTransaction, workflow_id: str, team_id: str, etag: str
) -> _DefinitionRow | DefinitionError:
await transaction.query_raw(
'SELECT workflow_id FROM "LiteLLM_WorkflowDefinition" WHERE workflow_id = $1 AND team_id = $2 FOR UPDATE',
workflow_id,
team_id,
)
stored: Final = await transaction.litellm_workflowdefinition.find_first(
where={"workflow_id": workflow_id, "team_id": team_id}
)
if stored is None:
return DefinitionError("not_found", "Workflow not found")
row: Final = _DefinitionRow.model_validate(stored)
if row.etag != etag:
return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match")
return row
table: TableActions["prisma_models.LiteLLM_WorkflowDefinition"], workflow_id: str, team_id: str, etag: str
) -> "prisma_models.LiteLLM_WorkflowDefinition | DefinitionError":
# Preserve audit timestamps while taking the row lock. If a concurrent
# publication changed the row, re-read its contributor boundary before editing.
for _ in range(3):
row = await table.find_first(where={"workflow_id": workflow_id, "team_id": team_id})
if row is None:
return DefinitionError("not_found", "Workflow not found")
if row.etag != etag:
break
if await table.update_many(
where={"workflow_id": workflow_id, "team_id": team_id, "etag": etag, "updated_at": row.updated_at},
data={"updated_at": row.updated_at},
):
locked = await table.find_unique(where={"workflow_id": workflow_id})
if locked is not None:
return locked
return DefinitionError("not_found", "Workflow not found")
return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match")
async def _snapshot(
transaction: DefinitionTransaction, row: _DefinitionRow, published_by: str | None
) -> WorkflowSnapshot:
stored: Final = await transaction.litellm_workflowsnapshot.create(
data={
"snapshot_id": str(uuid4()),
"workflow_id": row.workflow_id,
"team_id": row.team_id,
"name": row.name,
"digest": row.digest,
"bundle_json": row.bundle_json,
"contributors": list(row.contributors),
"published_by": published_by,
}
)
return _SnapshotRow.model_validate(stored).snapshot()
class WorkflowDefinitionRepository:
def __init__(self, database: DefinitionDatabase) -> None:
self._database: Final = database
class WorkflowDefinitionRepository(PrismaTableRepository["prisma_models.LiteLLM_WorkflowDefinition"]):
table_name = "litellm_workflowdefinition"
async def get_team(self, team_id: str) -> LiteLLM_TeamTable | None:
row: Final = await self._database.litellm_teamtable.find_unique(where={"team_id": team_id})
row: Final = await TeamRepository(self._prisma_client, use_writer=True).table.find_unique(
where={"team_id": team_id}
)
return None if row is None else LiteLLM_TeamTable.model_validate(row, from_attributes=True)
async def get_draft(self, workflow_id: str, team_id: str | None = None) -> WorkflowDefinition | None:
row: Final = await self._database.litellm_workflowdefinition.find_first(
row: Final = await self.table.find_first(
where={"workflow_id": workflow_id, **({"team_id": team_id} if team_id is not None else {})}
)
return None if row is None else _DefinitionRow.model_validate(row).definition()
return None if row is None else _definition(row)
async def list_drafts(self, team_id: str, limit: int, cursor: str | None) -> DefinitionPage:
rows: Final = await self._database.litellm_workflowdefinition.find_many(
rows: Final = await self.table.find_many(
where={"team_id": team_id, **({"workflow_id": {"gt": cursor}} if cursor is not None else {})},
order={"workflow_id": "asc"},
take=limit + 1,
)
drafts: Final = tuple(_DefinitionRow.model_validate(row).definition().draft for row in rows[:limit])
drafts: Final = tuple(_definition(row).draft for row in rows[:limit])
return DefinitionPage(workflows=drafts, next_cursor=drafts[-1].workflow_id if len(rows) > limit else None)
async def get_snapshot(self, snapshot_id: str, team_id: str | None = None) -> WorkflowSnapshot | None:
row: Final = await self._database.litellm_workflowsnapshot.find_first(
row: Final = await WorkflowSnapshotRepository(self._prisma_client, use_writer=True).table.find_first(
where={"snapshot_id": snapshot_id, **({"team_id": team_id} if team_id is not None else {})}
)
return None if row is None else _SnapshotRow.model_validate(row).snapshot()
return None if row is None else _snapshot(row)
async def get_published(self, name: str, team_id: str) -> WorkflowSnapshot | None:
row: Final = await self._database.litellm_workflowdefinition.find_first(
where={"name": name, "team_id": team_id}
)
if row is None:
row: Final = await self.table.find_first(where={"name": name, "team_id": team_id})
if row is None or row.published_snapshot_id is None:
return None
definition: Final = _DefinitionRow.model_validate(row)
if definition.published_snapshot_id is None:
return None
return await self.get_snapshot(definition.published_snapshot_id, team_id)
return await self.get_snapshot(row.published_snapshot_id, team_id)
async def create(self, bundle: WorkflowBundle, team_id: str, contributor: str) -> WorkflowDraft | DefinitionError:
from prisma.errors import UniqueViolationError
try:
stored: Final = await self._database.litellm_workflowdefinition.create(
stored: Final = await self.table.create(
data={
"workflow_id": str(uuid4()),
"name": bundle.name,
@ -202,15 +173,16 @@ class WorkflowDefinitionRepository:
"contributors": [contributor],
}
)
return _DefinitionRow.model_validate(stored).definition().draft
return _definition(stored).draft
except UniqueViolationError:
return DefinitionError("conflict", "A workflow with this name already exists in the team", "bundle.name")
async def replace(
self, workflow_id: str, team_id: str, etag: str, bundle: WorkflowBundle, contributor: str
) -> WorkflowDraft | DefinitionError:
async with self._database.tx() as transaction:
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
database: Final[Prisma] = self.writer_db
async with database.tx() as transaction:
row: Final = await _locked_definition(self.transaction_table(transaction), workflow_id, team_id, etag)
if isinstance(row, DefinitionError):
return row
if row.name != bundle.name:
@ -222,7 +194,7 @@ class WorkflowDefinitionRepository:
)
new_etag: Final = _etag()
digest: Final = bundle_digest(bundle)
await transaction.litellm_workflowdefinition.update(
await self.transaction_table(transaction).update(
where={"workflow_id": workflow_id},
data={
"etag": new_etag,
@ -244,8 +216,9 @@ class WorkflowDefinitionRepository:
async def publish(
self, workflow_id: str, team_id: str, etag: str, digest: str, publisher: str
) -> WorkflowSnapshot | DefinitionError:
async with self._database.tx() as transaction:
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
database: Final[Prisma] = self.writer_db
async with database.tx() as transaction:
row: Final = await _locked_definition(self.transaction_table(transaction), workflow_id, team_id, etag)
if isinstance(row, DefinitionError):
return row
if row.digest != digest:
@ -253,20 +226,22 @@ class WorkflowDefinitionRepository:
if publisher in row.contributors:
return DefinitionError("forbidden", "Draft contributors cannot publish their own changes")
if row.published_etag == etag and row.published_snapshot_id is not None:
existing: Final = await transaction.litellm_workflowsnapshot.find_unique(
existing: Final = await WorkflowSnapshotRepository.transaction_table(transaction).find_unique(
where={"snapshot_id": row.published_snapshot_id}
)
return _SnapshotRow.model_validate(existing).snapshot()
snapshot: Final = await _snapshot(transaction, row, publisher)
await transaction.litellm_workflowdefinition.update(
if existing is not None:
return _snapshot(existing)
snapshot: Final = await persist_snapshot(transaction, _prepared_snapshot(row, publisher))
await self.transaction_table(transaction).update(
where={"workflow_id": workflow_id},
data={"published_snapshot_id": snapshot.snapshot_id, "published_etag": etag},
)
return snapshot
async def freeze_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError:
async with self._database.tx() as transaction:
row: Final = await _locked_definition(transaction, workflow_id, team_id, etag)
if isinstance(row, DefinitionError):
return row
return await _snapshot(transaction, row, None)
async def prepare_draft(self, workflow_id: str, team_id: str, etag: str) -> WorkflowSnapshot | DefinitionError:
row: Final = await self.table.find_first(where={"workflow_id": workflow_id, "team_id": team_id})
if row is None:
return DefinitionError("not_found", "Workflow not found")
if row.etag != etag:
return DefinitionError("stale_draft", "The draft has changed. Read it again before retrying", "If-Match")
return _prepared_snapshot(row, None)

412
litellm/repositories/workflow_execution_repository.py
View file

@ -1,13 +1,12 @@
import base64
import json
from contextlib import AbstractAsyncContextManager
from datetime import datetime
from collections.abc import Mapping
from datetime import datetime, timezone
from itertools import chain
from typing import Final, Protocol
from uuid import uuid4
from typing import TYPE_CHECKING, Final
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError
from litellm.models.workflow_definition import WorkflowSnapshot
from litellm.models.workflow_execution import (
EngineRun,
ExecutionError,
@ -18,14 +17,18 @@ from litellm.models.workflow_execution import (
RunView,
StepTrace,
)
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
SpendLogsRepository,
TeamRepository,
WorkflowEventRepository,
WorkflowRunRepository,
)
from litellm.repositories.workflow_definition_repository import persist_snapshot
class ExecutionQueries(Protocol):
async def query_raw(self, query: str, *args: object) -> object: ...
class ExecutionDatabase(ExecutionQueries, Protocol):
def tx(self) -> AbstractAsyncContextManager[ExecutionQueries]: ...
if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
class _Cursor(BaseModel):
@ -41,12 +44,6 @@ class _TraceRow(BaseModel):
trace: StepTrace
_RUNS: Final = TypeAdapter(tuple[EngineRun, ...])
_VIEWS: Final = TypeAdapter(tuple[RunView, ...])
_TRACES: Final = TypeAdapter(tuple[_TraceRow, ...])
_ROWS: Final = TypeAdapter(tuple[dict[str, object], ...])
def _cursor(value: str | None) -> _Cursor | ExecutionError | None:
if value is None:
return None
@ -62,11 +59,6 @@ def _encoded_cursor(run: RunView) -> str:
).decode("ascii")
async def _first_run(database: ExecutionQueries, query: str, *args: object) -> EngineRun | None:
rows: Final = _RUNS.validate_python(await database.query_raw(query, *args))
return rows[0] if rows else None
def _aggregate_attempts(attempts: tuple[StepTrace, ...]) -> StepTrace:
latest: Final = max(attempts, key=lambda step: step.attempts)
costs: Final = tuple(step.response_cost for step in attempts if step.response_cost is not None)
@ -83,22 +75,44 @@ def _aggregate_attempts(attempts: tuple[StepTrace, ...]) -> StepTrace:
)
class WorkflowExecutionRepository:
def __init__(self, database: ExecutionDatabase) -> None:
self._database: Final = database
def _owned(run_id: str, token: str) -> Mapping[str, object]:
return {
"run_id": run_id,
"engine_owned": True,
"lease_token": token,
"status": RunState.RUNNING.value,
"lease_expires_at": {"gt": datetime.now(timezone.utc)},
}
async def _lock_team(table: TableActions["prisma_models.LiteLLM_TeamTable"], team_id: str) -> bool:
# A no-op model update serializes team admission until commit. Compare the
# observed timestamp so a concurrent team edit never has its audit time reset.
for _ in range(3):
team = await table.find_unique(where={"team_id": team_id})
if team is None:
return False
if await table.update_many(
where={"team_id": team_id, "updated_at": team.updated_at}, data={"updated_at": team.updated_at}
):
return True
return False
class WorkflowExecutionRepository(WorkflowRunRepository):
async def find_idempotent(
self, workflow_id: str, team_id: str, key_hash: str, idempotency_key: str
) -> EngineRun | None:
return await _first_run(
self._database,
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND workflow_id = $1 '
"AND team_id = $2 AND created_by = $3 AND idempotency_key = $4",
workflow_id,
team_id,
key_hash,
idempotency_key,
row: Final = await self.table.find_first(
where={
"engine_owned": True,
"workflow_id": workflow_id,
"team_id": team_id,
"created_by": key_hash,
"idempotency_key": idempotency_key,
}
)
return None if row is None else EngineRun.model_validate(row, from_attributes=True)
async def admit(
self,
@ -110,172 +124,185 @@ class WorkflowExecutionRepository:
input: RunInput,
digest: str,
idempotency_key: str | None,
draft_snapshot: WorkflowSnapshot | None = None,
) -> RunAdmission | ExecutionError:
async with self._database.tx() as transaction:
await transaction.query_raw(
"SELECT true AS acquired FROM pg_advisory_xact_lock(hashtextextended($1, 0))", team_id
)
existing: Final = await _first_run(
transaction,
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND created_by = $1 '
"AND team_id = $2 AND workflow_id = $3 AND idempotency_key = $4::text",
key_hash,
team_id,
workflow_id,
idempotency_key,
)
if existing is not None:
return (
RunAdmission(existing, False)
if existing.request_digest == digest
else ExecutionError("idempotency_conflict", "Idempotency key was used with different input")
)
capacity: Final = _ROWS.validate_python(
await transaction.query_raw(
'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND team_id = $1 '
"AND status IN ('pending', 'running') LIMIT 64",
team_id,
from prisma import Json
database: Final[Prisma] = self.writer_db
async with database.tx() as transaction:
if not await _lock_team(TeamRepository.transaction_table(transaction), team_id):
return ExecutionError("storage_error", "Workflow team is unavailable. Retry the request")
if idempotency_key is not None:
existing: Final = await self.transaction_table(transaction).find_first(
where={
"engine_owned": True,
"created_by": key_hash,
"team_id": team_id,
"workflow_id": workflow_id,
"idempotency_key": idempotency_key,
}
)
if existing is not None:
return (
RunAdmission(EngineRun.model_validate(existing, from_attributes=True), False)
if existing.request_digest == digest
else ExecutionError("idempotency_conflict", "Idempotency key was used with different input")
)
capacity: Final = await self.transaction_table(transaction).count(
where={"engine_owned": True, "team_id": team_id, "status": {"in": ["pending", "running"]}},
take=64,
)
if len(capacity) >= 64:
if capacity >= 64:
return ExecutionError("capacity_exceeded", "The team's workflow execution queue is full")
created: Final = await _first_run(
transaction,
'INSERT INTO "LiteLLM_WorkflowRun" '
"(run_id, session_id, workflow_type, status, created_by, created_at, updated_at, "
"workflow_id, snapshot_id, team_id, engine_owned, input, request_digest, idempotency_key) "
"VALUES ($1, $2, 'liteflow', 'pending', $3, timezone('UTC', NOW()), timezone('UTC', NOW()), $4, $5, $6, true, $7::jsonb, $8, $9) "
"RETURNING *",
str(uuid4()),
str(uuid4()),
key_hash,
workflow_id,
snapshot_id,
team_id,
input.model_dump_json(),
digest,
idempotency_key,
)
if created is None:
return ExecutionError("storage_error", "Workflow admission did not create a run")
return RunAdmission(created, True)
if draft_snapshot is not None:
await persist_snapshot(transaction, draft_snapshot)
now: Final = datetime.now(timezone.utc)
data: Final = {
"created_at": now,
"updated_at": now,
"workflow_type": "liteflow",
"workflow_id": workflow_id,
"snapshot_id": snapshot_id,
"team_id": team_id,
"created_by": key_hash,
"engine_owned": True,
"input": Json(input.model_dump(mode="json")),
"request_digest": digest,
"idempotency_key": idempotency_key,
}
created: Final = await self.transaction_table(transaction).create(data=data)
return RunAdmission(EngineRun.model_validate(created, from_attributes=True), True)
async def get(self, run_id: str) -> EngineRun | None:
return await _first_run(
self._database,
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND engine_owned = true',
run_id,
)
row: Final = await self.table.find_first(where={"run_id": run_id, "engine_owned": True})
return None if row is None else EngineRun.model_validate(row, from_attributes=True)
async def recover(self, limit: int) -> tuple[EngineRun, ...]:
return _RUNS.validate_python(
await self._database.query_raw(
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true AND '
"(status = 'pending' OR (status = 'running' AND lease_expires_at < timezone('UTC', NOW()))) "
"ORDER BY created_at, run_id LIMIT $1",
limit,
)
rows: Final = await self.table.find_many(
where={
"engine_owned": True,
"OR": [
{"status": RunState.PENDING.value},
{"status": RunState.RUNNING.value, "lease_expires_at": {"lt": datetime.now(timezone.utc)}},
],
},
order=[{"created_at": "asc"}, {"run_id": "asc"}],
take=limit,
)
return tuple(EngineRun.model_validate(row, from_attributes=True) for row in rows)
async def claim(self, run_id: str, token: str, lease_expires_at: datetime) -> EngineRun | None:
return await _first_run(
self._database,
"UPDATE \"LiteLLM_WorkflowRun\" SET status = 'running', lease_token = $2, "
"lease_expires_at = timezone('UTC', $3::text::timestamptz), updated_at = timezone('UTC', NOW()) "
"WHERE run_id = $1 AND engine_owned = true AND "
"(status = 'pending' OR (status = 'running' AND lease_expires_at < timezone('UTC', NOW()))) RETURNING *",
run_id,
token,
lease_expires_at.isoformat(),
)
database: Final[Prisma] = self.writer_db
async with database.tx() as transaction:
claimed: Final = await self.transaction_table(transaction).update_many(
where={
"run_id": run_id,
"engine_owned": True,
"OR": [
{"status": RunState.PENDING.value},
{"status": RunState.RUNNING.value, "lease_expires_at": {"lt": datetime.now(timezone.utc)}},
],
},
data={"status": RunState.RUNNING.value, "lease_token": token, "lease_expires_at": lease_expires_at},
)
if not claimed:
return None
row: Final = await self.transaction_table(transaction).find_unique(where={"run_id": run_id})
return None if row is None else EngineRun.model_validate(row, from_attributes=True)
async def renew(self, run_id: str, token: str, lease_expires_at: datetime) -> bool:
rows: Final = _ROWS.validate_python(
await self._database.query_raw(
"UPDATE \"LiteLLM_WorkflowRun\" SET lease_expires_at = timezone('UTC', $3::text::timestamptz) "
"WHERE run_id = $1 AND lease_token = $2 AND status = 'running' "
"AND lease_expires_at > timezone('UTC', NOW()) RETURNING run_id",
run_id,
token,
lease_expires_at.isoformat(),
)
return bool(
await self.table.update_many(where=_owned(run_id, token), data={"lease_expires_at": lease_expires_at})
)
return bool(rows)
async def checkpoint(self, run_id: str, token: str, step: StepTrace) -> bool:
async with self._database.tx() as transaction:
owned: Final = _ROWS.validate_python(
await transaction.query_raw(
'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE run_id = $1 AND lease_token = $2 '
"AND status = 'running' AND lease_expires_at > timezone('UTC', NOW()) FOR UPDATE",
run_id,
token,
)
)
if not owned:
from prisma import Json
database: Final[Prisma] = self.writer_db
async with database.tx() as transaction:
where: Final = dict(_owned(run_id, token))
if step.state is RunState.RUNNING:
where["cancel_requested"] = False
# The conditional update locks the run while allocating its event
# sequence. Cancellation fences dispatch but still accepts receipts.
if not await self.transaction_table(transaction).update_many(where=where, data={"lease_token": token}):
return False
await transaction.query_raw(
'INSERT INTO "LiteLLM_WorkflowEvent" '
"(event_id, run_id, event_type, step_name, sequence_number, data, created_at) "
"SELECT $1, $2, 'liteflow.step', $3, COALESCE(MAX(sequence_number) + 1, 0), $4::jsonb, timezone('UTC', NOW()) "
'FROM "LiteLLM_WorkflowEvent" WHERE run_id = $2 RETURNING event_id',
str(uuid4()),
run_id,
step.node_id,
json.dumps({"trace": step.model_dump(mode="json")}, allow_nan=False),
last: Final = await WorkflowEventRepository.transaction_table(transaction).find_first(
where={"run_id": run_id}, order={"sequence_number": "desc"}
)
data: Final = {
"created_at": datetime.now(timezone.utc),
"run_id": run_id,
"event_type": "liteflow.step",
"step_name": step.node_id,
"sequence_number": 0 if last is None else last.sequence_number + 1,
"data": Json({"trace": step.model_dump(mode="json")}),
}
await WorkflowEventRepository.transaction_table(transaction).create(data=data)
return True
async def steps(self, run_id: str) -> tuple[StepTrace, ...]:
rows: Final = _TRACES.validate_python(
await self._database.query_raw(
"SELECT trace FROM ("
"SELECT CASE WHEN receipts.ids IS NULL THEN latest.trace ELSE latest.trace || "
"jsonb_build_object('request_id', CASE WHEN cardinality(receipts.ids) = 1 THEN receipts.ids[1] END, "
"'request_ids', to_jsonb(receipts.ids), 'response_cost', receipts.cost, "
"'prompt_tokens', receipts.prompt_tokens, 'completion_tokens', receipts.completion_tokens) END AS trace, "
"latest.step_name FROM ("
"SELECT DISTINCT ON (step_name, data->'trace'->>'attempts') data->'trace' AS trace, step_name "
"FROM \"LiteLLM_WorkflowEvent\" WHERE run_id = $1 AND event_type = 'liteflow.step' "
"ORDER BY step_name, data->'trace'->>'attempts', sequence_number DESC"
') latest JOIN "LiteLLM_WorkflowRun" run ON run.run_id = $1 AND run.engine_owned = true '
"LEFT JOIN LATERAL ("
"SELECT array_agg(log.request_id ORDER BY log.request_id) AS ids, SUM(log.spend) AS cost, "
"SUM(log.prompt_tokens) AS prompt_tokens, SUM(log.completion_tokens) AS completion_tokens "
"FROM \"LiteLLM_SpendLogs\" log WHERE log.litellm_call_id = latest.trace->>'call_id' "
"AND log.api_key = run.created_by AND log.team_id = run.team_id AND log.session_id = run.session_id"
") receipts ON true"
") latest ORDER BY (trace->>'started_at')::timestamptz, step_name",
run_id,
run: Final = await self.get(run_id)
if run is None:
return ()
events: Final = await WorkflowEventRepository(self._prisma_client, use_writer=True).table.find_many(
where={"run_id": run_id, "event_type": "liteflow.step"}, order={"sequence_number": "asc"}
)
latest: Final[dict[tuple[str, int], StepTrace]] = {} # mutable-ok: local attempt index
for event in events:
step = _TraceRow.model_validate(event.data).trace
latest[(step.node_id, step.attempts)] = step
call_ids: Final = [step.call_id for step in latest.values() if step.call_id is not None]
receipts: Final[dict[str, list[prisma_models.LiteLLM_SpendLogs]]] = {} # mutable-ok: local receipt grouping
if call_ids:
logs: Final = await SpendLogsRepository(self._prisma_client, use_writer=True).table.find_many(
where={
"litellm_call_id": {"in": call_ids},
"api_key": run.created_by,
"team_id": run.team_id,
"session_id": run.session_id,
},
order={"request_id": "asc"},
)
)
node_ids: Final = dict.fromkeys(row.trace.node_id for row in rows)
return tuple(
_aggregate_attempts(tuple(row.trace for row in rows if row.trace.node_id == node_id))
for node_id in node_ids
)
for log in logs:
if log.litellm_call_id is not None:
receipts.setdefault(log.litellm_call_id, []).append(log)
attempts: Final[dict[str, list[StepTrace]]] = {} # mutable-ok: local attempt grouping before immutable result
for step in sorted(latest.values(), key=lambda trace: (trace.started_at, trace.node_id)):
billed = receipts.get(step.call_id) if step.call_id is not None else None
if billed:
step = step.model_copy(
update={
"request_id": billed[0].request_id if len(billed) == 1 else None,
"request_ids": tuple(log.request_id for log in billed),
"response_cost": sum(log.spend for log in billed),
"prompt_tokens": sum(log.prompt_tokens for log in billed),
"completion_tokens": sum(log.completion_tokens for log in billed),
}
)
attempts.setdefault(step.node_id, []).append(step)
return tuple(_aggregate_attempts(tuple(node_attempts)) for node_attempts in attempts.values())
async def finish(self, run_id: str, token: str, state: RunState, output: JsonValue, error: str | None) -> bool:
rows: Final = _ROWS.validate_python(
await self._database.query_raw(
'UPDATE "LiteLLM_WorkflowRun" SET status = $3, output = $4::jsonb, '
"metadata = jsonb_build_object('error', $5::text), lease_token = NULL, lease_expires_at = NULL, "
"updated_at = timezone('UTC', NOW()) WHERE run_id = $1 AND lease_token = $2 AND status = 'running' "
"AND lease_expires_at > timezone('UTC', NOW()) RETURNING run_id",
run_id,
token,
state.value,
json.dumps(output, allow_nan=False),
error,
from prisma import Json
return bool(
await self.table.update_many(
where=_owned(run_id, token),
data={
"status": state.value,
"output": Json(output),
"metadata": Json({"error": error}),
"lease_token": None,
"lease_expires_at": None,
},
)
)
return bool(rows)
async def cancel(self, run_id: str) -> None:
await self._database.query_raw(
"UPDATE \"LiteLLM_WorkflowRun\" SET cancel_requested = true, updated_at = timezone('UTC', NOW()) "
"WHERE run_id = $1 AND engine_owned = true AND status IN ('pending', 'running') RETURNING run_id",
run_id,
await self.table.update_many(
where={"run_id": run_id, "engine_owned": True, "status": {"in": ["pending", "running"]}},
data={"cancel_requested": True},
)
async def list_runs(
@ -290,25 +317,22 @@ class WorkflowExecutionRepository:
position: Final = _cursor(cursor)
if isinstance(position, ExecutionError):
return position
rows: Final = _VIEWS.validate_python(
await self._database.query_raw(
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE engine_owned = true '
"AND workflow_id = $1 AND team_id = $2 "
"AND ($3::text IS NULL OR created_by = $3::text) AND ($4::text IS NULL OR status = $4::text) "
"AND ($5::text IS NULL OR (created_at, run_id) < (timezone('UTC', $5::text::timestamptz), $6::text)) "
"ORDER BY created_at DESC, run_id DESC LIMIT $7",
workflow_id,
team_id,
key_hash,
state.value if state is not None else None,
position.created_at.isoformat() if position is not None else None,
position.run_id if position is not None else None,
limit + 1,
)
)
page: Final = rows[:limit]
return RunPage(
runs=page,
count=len(page),
next_cursor=_encoded_cursor(page[-1]) if len(rows) > limit else None,
where: Final[dict[str, object]] = { # mutable-ok: build a local Prisma filter from optional fields
"engine_owned": True,
"workflow_id": workflow_id,
"team_id": team_id,
}
if key_hash is not None:
where["created_by"] = key_hash
if state is not None:
where["status"] = state.value
if position is not None:
where["OR"] = [
{"created_at": {"lt": position.created_at}},
{"created_at": position.created_at, "run_id": {"lt": position.run_id}},
]
rows: Final = await self.table.find_many(
where=where, order=[{"created_at": "desc"}, {"run_id": "desc"}], take=limit + 1
)
page: Final = tuple(RunView.model_validate(row, from_attributes=True) for row in rows[:limit])
return RunPage(runs=page, count=len(page), next_cursor=_encoded_cursor(page[-1]) if len(rows) > limit else None)

193
tests/integration/database/test_workflow_execution.py
View file

@ -4,7 +4,8 @@ import sys
from collections.abc import AsyncIterator, Iterator
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Final, cast
from types import SimpleNamespace
from typing import Final
import psycopg
import pytest
@ -14,8 +15,8 @@ from prisma.types import DatasourceOverride
from psycopg import sql
from pydantic import TypeAdapter
from litellm.models.workflow_execution import RunAdmission, RunInput, RunPage, RunState, StepTrace
from litellm.repositories.workflow_execution_repository import ExecutionDatabase, WorkflowExecutionRepository
from litellm.models.workflow_execution import ExecutionError, RunAdmission, RunInput, RunPage, RunState, StepTrace
from litellm.repositories.workflow_execution_repository import WorkflowExecutionRepository
from tests.integration._support.database import scratch_database
_ROOT: Final = Path(__file__).resolve().parents[3]
@ -24,9 +25,21 @@ _ROOT: Final = Path(__file__).resolve().parents[3]
@pytest.fixture(scope="module")
def schema_sql() -> bytes:
return subprocess.run(
[sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel",
str(_ROOT / "schema.prisma"), "--script"],
check=True, capture_output=True, cwd=_ROOT,
[
sys.executable,
"-I",
"-m",
"prisma",
"migrate",
"diff",
"--from-empty",
"--to-schema-datamodel",
str(_ROOT / "schema.prisma"),
"--script",
],
check=True,
capture_output=True,
cwd=_ROOT,
).stdout
@ -48,15 +61,14 @@ async def ledger(database: str) -> AsyncIterator[Prisma]:
assert TypeAdapter(tuple[dict[str, str], ...]).validate_python(await client.query_raw("SHOW TIME ZONE")) == (
{"TimeZone": "Asia/Singapore"},
)
await client.litellm_teamtable.create(data={"team_id": "team", "admins": [], "members": [], "models": []})
yield client
finally:
await client.disconnect()
def _repository(ledger: Prisma) -> WorkflowExecutionRepository:
return WorkflowExecutionRepository(
cast(ExecutionDatabase, ledger) # cast-ok: generated Prisma transaction delegates satisfy the protocol
)
return WorkflowExecutionRepository(SimpleNamespace(writer_db=ledger, db=ledger), use_writer=True)
def _input(now: datetime) -> RunInput:
@ -64,12 +76,21 @@ def _input(now: datetime) -> RunInput:
async def _admit(
repository: WorkflowExecutionRepository, data: RunInput, key: str | None,
*, snapshot: str = "snapshot", owner: str = "owner",
repository: WorkflowExecutionRepository,
data: RunInput,
key: str | None,
*,
snapshot: str = "snapshot",
owner: str = "owner",
) -> RunAdmission:
result: Final = await repository.admit(
workflow_id="workflow", snapshot_id=snapshot, team_id="team", key_hash=owner,
input=data, digest="same-input", idempotency_key=key,
workflow_id="workflow",
snapshot_id=snapshot,
team_id="team",
key_hash=owner,
input=data,
digest="same-input",
idempotency_key=key,
)
assert isinstance(result, RunAdmission)
return result
@ -97,8 +118,16 @@ async def test_concurrent_admission_pins_snapshot_and_execution_timestamps_are_u
renewed: Final = await repository.get(run_id)
assert renewed is not None and renewed.lease_expires_at == extended
step: Final = StepTrace(
step_id="step", run_id=run_id, node_id="node", state=RunState.COMPLETED, input={"n": 1}, output={"n": 2},
attempts=1, deadline=data.deadline, started_at=before, finished_at=before,
step_id="step",
run_id=run_id,
node_id="node",
state=RunState.COMPLETED,
input={"n": 1},
output={"n": 2},
attempts=1,
deadline=data.deadline,
started_at=before,
finished_at=before,
)
assert await repository.checkpoint(run_id, "other-worker", step) is False
assert await repository.checkpoint(run_id, "worker", step) is True
@ -107,7 +136,10 @@ async def test_concurrent_admission_pins_snapshot_and_execution_timestamps_are_u
finished: Final = await repository.get(run_id)
assert finished is not None and finished.status is RunState.COMPLETED
assert (finished.output, finished.metadata, finished.lease_token, finished.lease_expires_at) == (
{"n": 2}, {"error": None}, None, None,
{"n": 2},
{"error": None},
None,
None,
)
stored: Final = await ledger.litellm_workflowrun.find_unique(where={"run_id": run_id})
events: Final = await ledger.litellm_workflowevent.find_many(where={"run_id": run_id})
@ -137,7 +169,9 @@ async def test_nullable_admission_and_history_filters_round_trip_through_prisma(
scoped: Final = await repository.list_runs("workflow", "team", "owner", RunState.PENDING, 1, None)
assert isinstance(scoped, RunPage) and scoped.next_cursor is not None
assert tuple(run.run_id for run in scoped.runs) == owned[:1]
continuation: Final = await repository.list_runs("workflow", "team", "owner", RunState.PENDING, 1, scoped.next_cursor)
continuation: Final = await repository.list_runs(
"workflow", "team", "owner", RunState.PENDING, 1, scoped.next_cursor
)
assert isinstance(continuation, RunPage) and continuation.count == 1 and continuation.next_cursor is None
assert tuple(run.run_id for run in continuation.runs) == owned[1:]
assert await repository.claim(second.run.run_id, "worker", data.deadline) is not None
@ -147,3 +181,128 @@ async def test_nullable_admission_and_history_filters_round_trip_through_prisma(
filtered: Final = await repository.list_runs("workflow", "team", "owner", RunState.FAILED, 2, None)
assert isinstance(filtered, RunPage) and filtered.count == 1 and filtered.next_cursor is None
assert tuple(run.run_id for run in filtered.runs) == (second.run.run_id,)
@pytest.mark.asyncio(loop_scope="function")
async def test_concurrent_admission_cannot_overfill_team_queue(ledger: Prisma) -> None:
repository: Final = _repository(ledger)
data: Final = _input(datetime.now(timezone.utc))
for index in range(63):
await _admit(repository, data, f"seed-{index}")
outcomes: Final = await asyncio.gather(
*(
repository.admit(
workflow_id=workflow,
snapshot_id="snapshot",
team_id="team",
key_hash="owner",
input=data,
digest="same-input",
idempotency_key=f"last-{workflow}",
)
for workflow in ("workflow", "other-workflow")
)
)
assert sum(isinstance(result, RunAdmission) for result in outcomes) == 1
assert sum(isinstance(result, ExecutionError) and result.code == "capacity_exceeded" for result in outcomes) == 1
assert await ledger.litellm_workflowrun.count() == 64
@pytest.mark.asyncio(loop_scope="function")
async def test_draft_admission_persists_one_snapshot_and_rolls_back_when_run_insert_fails(ledger: Prisma) -> None:
from prisma.errors import UniqueViolationError
from litellm.models.workflow_definition import WorkflowDraft, WorkflowSnapshot
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from tests.unit.proxy.workflow_definition_rows import bundle
client: Final = SimpleNamespace(writer_db=ledger, db=ledger)
definitions: Final = WorkflowDefinitionRepository(client, use_writer=True)
executions: Final = WorkflowExecutionRepository(client, use_writer=True)
draft: Final = await definitions.create(bundle(), "team", "author")
assert isinstance(draft, WorkflowDraft)
prepared: Final = await definitions.prepare_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(prepared, WorkflowSnapshot)
assert await ledger.litellm_workflowsnapshot.count() == 0
data: Final = _input(datetime.now(timezone.utc))
# A legacy row with the same unique key forces the run insert (after snapshot
# creation) to fail, exercising actual database rollback rather than a fake.
await ledger.litellm_workflowrun.create(
data={
"workflow_type": "tracking",
"workflow_id": draft.workflow_id,
"team_id": "team",
"created_by": "owner",
"idempotency_key": "collision",
}
)
with pytest.raises(UniqueViolationError):
await executions.admit(
workflow_id=draft.workflow_id,
snapshot_id=prepared.snapshot_id,
team_id="team",
key_hash="owner",
input=data,
digest="same-input",
idempotency_key="collision",
draft_snapshot=prepared,
)
assert await ledger.litellm_workflowsnapshot.count() == 0
first: Final = await executions.admit(
workflow_id=draft.workflow_id,
snapshot_id=prepared.snapshot_id,
team_id="team",
key_hash="owner",
input=data,
digest="same-input",
idempotency_key="accepted",
draft_snapshot=prepared,
)
assert isinstance(first, RunAdmission) and first.created
again: Final = await definitions.prepare_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(again, WorkflowSnapshot)
replay: Final = await executions.admit(
workflow_id=draft.workflow_id,
snapshot_id=again.snapshot_id,
team_id="team",
key_hash="owner",
input=data,
digest="same-input",
idempotency_key="accepted",
draft_snapshot=again,
)
assert replay == RunAdmission(first.run, False)
assert await ledger.litellm_workflowsnapshot.count() == 1
@pytest.mark.asyncio(loop_scope="function")
async def test_cancelled_run_rejects_dispatch_and_accepts_terminal_receipt(ledger: Prisma) -> None:
repository: Final = _repository(ledger)
data: Final = _input(datetime.now(timezone.utc))
admitted: Final = await _admit(repository, data, "cancel")
run_id: Final = admitted.run.run_id
assert await repository.claim(run_id, "worker", data.deadline) is not None
step: Final = StepTrace(
step_id="paid-step",
run_id=run_id,
node_id="node",
state=RunState.RUNNING,
input=None,
attempts=1,
deadline=data.deadline,
started_at=datetime.now(timezone.utc),
)
assert await repository.checkpoint(run_id, "worker", step)
await repository.cancel(run_id)
assert not await repository.checkpoint(run_id, "worker", step.model_copy(update={"attempts": 2}))
receipt: Final = step.model_copy(
update={
"state": RunState.COMPLETED,
"output": "paid-result",
"response_cost": 0.25,
"finished_at": datetime.now(timezone.utc),
}
)
assert await repository.checkpoint(run_id, "worker", receipt)
(stored,) = await repository.steps(run_id)
assert stored.state is RunState.COMPLETED and stored.output == "paid-result" and stored.response_cost == 0.25

40
tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py
View file

@ -38,6 +38,7 @@ from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.proxy.common_utils.user_api_key_cache import (
NO_TEAM_MEMBERSHIP_SENTINEL,
UserApiKeyCache,
end_user_cache_key,
team_membership_reservation_cache_key,
)
from litellm.proxy.hooks.model_max_budget_limiter import (
@ -188,6 +189,45 @@ async def test_internal_execution_route_keeps_proxy_network_policy(monkeypatch:
assert error.value.status_code == 403
@pytest.mark.asyncio
async def test_internal_execution_keeps_end_user_rate_limits_region_and_permissions(
monkeypatch: pytest.MonkeyPatch,
) -> None:
cache: Final = await _internal_key_cache(monkeypatch)
permissions: Final = LiteLLM_ObjectPermissionTable(object_permission_id="end-user-permissions")
await cache.async_set_cache(
key=end_user_cache_key("restricted-user"),
value=LiteLLM_EndUserTable(
user_id="restricted-user",
blocked=False,
allowed_model_region="eu",
object_permission=permissions,
litellm_budget_table=LiteLLM_BudgetTable(
budget_id="end-user-budget",
rpm_limit=1,
tpm_limit=100,
tpd_limit=1000,
max_budget=100,
),
),
)
auth: Final = await authorize_internal_virtual_key(
"b" * 64,
_internal_management_request(),
{"model": "workflow-model", "user": "restricted-user"},
route="/v1/chat/completions",
expected_team_id="original-team",
)
assert auth.end_user_id == "restricted-user"
assert (
auth.end_user_rpm_limit,
auth.end_user_tpm_limit,
auth.end_user_tpd_limit,
) == (1, 100, 1000)
assert auth.allowed_model_region == "eu"
assert auth.end_user_object_permission == permissions
@pytest.mark.asyncio
async def test_workflow_wrapper_with_model_grant_passes_strict_pricing_but_unpriced_child_fails(
monkeypatch: pytest.MonkeyPatch,

16
tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py
View file

@ -15,6 +15,13 @@ from litellm.repositories.workflow_definition_repository import WorkflowDefiniti
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller
@pytest.fixture(autouse=True)
def enable_workflows(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
monkeypatch.setattr(proxy_server, "general_settings", {"enable_workflows": True})
def _client(database: DefinitionRows | None, identity: UserAPIKeyAuth) -> TestClient:
app: Final = FastAPI()
app.include_router(router)
@ -55,6 +62,15 @@ def test_sdk_draft_contract_round_trip_has_exact_fields_and_etag() -> None:
assert fetched.headers["etag"] == draft.etag
@pytest.mark.parametrize("enabled", [False, True])
def test_definition_routes_require_explicit_opt_in(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
from litellm.proxy import proxy_server
monkeypatch.setattr(proxy_server, "general_settings", {"enable_workflows": True} if enabled else {})
response: Final = _client(DefinitionRows(), caller()).post("/v1/workflows", json=_body())
assert response.status_code == (201 if enabled else 404), response.text
def test_stale_put_and_publish_cannot_replace_or_approve_new_revision() -> None:
database: Final = DefinitionRows()
author: Final = _client(database, caller())

159
tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py
View file

@ -1,5 +1,5 @@
from collections.abc import Mapping
from datetime import timedelta
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
@ -10,7 +10,7 @@ from pydantic import JsonValue, TypeAdapter
from litellm.models.workflow_definition import WorkflowDraft, WorkflowSnapshot
from litellm.models.workflow_execution import RunPage, RunState, StepTrace
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.deps.repositories import get_workflow_definition_repository, get_workflow_execution_repository
from litellm.proxy.management_endpoints.workflow_definition_endpoints import router as definitions_router
@ -26,11 +26,19 @@ _MESSAGES: Final = ({"role": "user", "content": "plan"},)
_ACCEPTED: Final = TypeAdapter(dict[str, str])
def virtual_key(user: str = "author", team: str = "team", key: str = "owner-key") -> UserAPIKeyAuth:
@pytest.fixture(autouse=True)
def enable_workflows(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
monkeypatch.setattr(proxy_server, "general_settings", {"enable_workflows": True})
def virtual_key(user: str | None = "author", team: str = "team", key: str = "owner-key") -> UserAPIKeyAuth:
return caller(user, team).model_copy(update={"token": key, "via_virtual_key": True})
def client_for(definitions: DefinitionRows, executions: ExecutionRows, identity: UserAPIKeyAuth) -> TestClient:
executions.litellm_workflowsnapshot = definitions.litellm_workflowsnapshot
app: Final = FastAPI()
app.include_router(tracking_router)
app.include_router(definitions_router)
@ -74,7 +82,6 @@ def test_paid_test_returns_202_and_persists_pending_admission_without_executing_
assert admitted.input.messages == _MESSAGES and admitted.input.draft
assert (admitted.created_by, admitted.team_id) == ("owner-key", "team")
assert executions.events == ()
assert not any(query.sql.startswith("UPDATE") for query in executions.queries)
trace: Final = client.get(f"/v1/workflows/runs/{admitted.run_id}/steps")
assert trace.status_code == 200, trace.text
assert trace.json() == {
@ -87,6 +94,26 @@ def test_paid_test_returns_202_and_persists_pending_admission_without_executing_
}
@pytest.mark.parametrize("enabled", [False, True])
def test_execution_routes_require_explicit_opt_in(monkeypatch: pytest.MonkeyPatch, enabled: bool) -> None:
from litellm.proxy import proxy_server
definitions: Final = DefinitionRows()
executions: Final = ExecutionRows()
client: Final = client_for(definitions, executions, virtual_key())
draft: Final = register(client)
monkeypatch.setattr(proxy_server, "general_settings", {"enable_workflows": True} if enabled else {})
response: Final = client.post(
f"/v1/workflows/{draft.workflow_id}/test",
json={"messages": _MESSAGES, "authorize_billing": True},
headers={"If-Match": draft.etag},
)
assert response.status_code == (202 if enabled else 404), response.text
assert len(executions.runs) == int(enabled)
history: Final = client.get("/v1/workflows/runs", params={"workflow_id": draft.workflow_id})
assert history.status_code == (200 if enabled else 404), history.text
def test_idempotent_test_after_draft_edit_returns_original_snapshot_and_conflict_for_changed_input() -> None:
definitions: Final = DefinitionRows()
executions: Final = ExecutionRows()
@ -137,6 +164,47 @@ def test_billing_consent_and_virtual_key_are_required_before_creating_snapshot_o
assert executions.runs == () and definitions.litellm_workflowsnapshot.rows == ()
@pytest.mark.parametrize(
"messages",
[
[{"role": "invalid", "content": "plan"}],
[{"role": "user", "content": "x" * (1024 * 1024)}],
],
)
def test_invalid_test_input_does_not_create_snapshot_or_run(messages: list[dict[str, str]]) -> None:
definitions: Final = DefinitionRows()
executions: Final = ExecutionRows()
client: Final = client_for(definitions, executions, virtual_key())
draft: Final = register(client)
response: Final = client.post(
f"/v1/workflows/{draft.workflow_id}/test",
json={"messages": messages, "authorize_billing": True},
headers={"If-Match": draft.etag},
)
assert response.status_code == 400, response.text
assert response.json()["detail"]["error"]["code"] == "invalid_input"
assert definitions.litellm_workflowsnapshot.rows == ()
assert executions.runs == ()
def test_full_team_queue_does_not_create_draft_test_snapshots() -> None:
definitions: Final = DefinitionRows()
existing: Final = tuple(pending_run(f"run-{index}") for index in range(64))
executions: Final = ExecutionRows(existing)
client: Final = client_for(definitions, executions, virtual_key())
draft: Final = register(client)
for _ in range(2):
response: Final = client.post(
f"/v1/workflows/{draft.workflow_id}/test",
json={"messages": _MESSAGES, "authorize_billing": True},
headers={"If-Match": draft.etag},
)
assert response.status_code == 429, response.text
assert response.json()["detail"]["error"]["code"] == "capacity_exceeded"
assert definitions.litellm_workflowsnapshot.rows == ()
assert executions.runs == existing
@pytest.mark.parametrize("identity", [virtual_key("other", key="other-key"), virtual_key("outsider", "other-team")])
def test_foreign_run_denies_trace_definition_and_cancellation_without_side_effects(identity: UserAPIKeyAuth) -> None:
definitions: Final = DefinitionRows()
@ -156,12 +224,43 @@ def test_foreign_run_denies_trace_definition_and_cancellation_without_side_effec
denied: Final = foreign.post(f"{path}/cancel")
assert denied.status_code == 404, denied.text
assert not executions.runs[0].cancel_requested
assert not any(query.sql.startswith("UPDATE") for query in executions.queries)
assert owner.get(f"{path}/steps").status_code == 200
assert owner.post(f"{path}/cancel").status_code == 202
assert executions.runs[0].cancel_requested
@pytest.mark.parametrize(
("identity", "expected_status"),
[
(virtual_key("reviewer", "other-team", "other-team-admin-key"), 404),
(virtual_key("reviewer", "team", "team-admin-key"), 200),
(virtual_key("reviewer", "other-team").model_copy(update={"is_session_token": True}), 200),
(
virtual_key("proxy-viewer", "other-team").model_copy(
update={"user_role": LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY}
),
200,
),
],
)
def test_run_admin_reads_honor_virtual_key_team_scope(identity: UserAPIKeyAuth, expected_status: int) -> None:
definitions: Final = DefinitionRows()
executions: Final = ExecutionRows()
owner: Final = client_for(definitions, executions, virtual_key())
draft: Final = register(owner)
accepted: Final = owner.post(
f"/v1/workflows/{draft.workflow_id}/test",
json={"messages": _MESSAGES, "authorize_billing": True},
headers={"If-Match": draft.etag},
)
assert accepted.status_code == 202, accepted.text
path: Final = f"/v1/workflows/runs/{accepted.json()['run_id']}"
reader: Final = client_for(definitions, executions, identity)
for suffix in ("steps", "definition"):
response: Final = reader.get(f"{path}/{suffix}")
assert response.status_code == expected_status, response.text
def test_history_caps_page_size_and_scopes_ordinary_key_but_allows_team_admin() -> None:
definitions: Final = DefinitionRows()
executions: Final = ExecutionRows()
@ -177,7 +276,6 @@ def test_history_caps_page_size_and_scopes_ordinary_key_but_allows_team_admin()
page: Final = RunPage.model_validate(response.json())
assert page.count == 100 and len(page.runs) == 100 and page.next_cursor is not None
assert tuple(run.run_id for run in page.runs) == tuple(f"run-{index:03}" for index in range(100, 0, -1))
assert executions.queries[-1].args == (draft.workflow_id, "team", "owner-key", None, None, None, 101)
tail_response: Final = owner.get(
f"/v1/workflows/{draft.workflow_id}/runs", params={"limit": 100, "after": page.next_cursor}
)
@ -187,13 +285,55 @@ def test_history_caps_page_size_and_scopes_ordinary_key_but_allows_team_admin()
admin: Final = client_for(definitions, executions, virtual_key("reviewer", key="admin-key"))
team_response: Final = admin.get("/v1/workflows/runs", params={"workflow_id": draft.workflow_id, "limit": 250})
assert team_response.status_code == 200, team_response.text
assert executions.queries[-1].args == (draft.workflow_id, "team", None, None, None, None, 101)
team_page: Final = RunPage.model_validate(team_response.json())
team_tail_response: Final = admin.get(
f"/v1/workflows/{draft.workflow_id}/runs", params={"limit": 100, "after": team_page.next_cursor}
)
assert team_tail_response.status_code == 200, team_tail_response.text
team_tail: Final = RunPage.model_validate(team_tail_response.json())
assert tuple(run.run_id for run in team_tail.runs) == ("run-000", "foreign-key")
assert team_tail.next_cursor is None
outsider: Final = client_for(definitions, executions, virtual_key("outsider", "other-team"))
denied: Final = outsider.get("/v1/workflows/runs", params={"workflow_id": draft.workflow_id, "team_id": "team"})
assert denied.status_code == 404, denied.text
assert owner.get(f"/v1/workflows/{draft.workflow_id}/runs", params={"limit": 101}).status_code == 422
@pytest.mark.asyncio
async def test_application_key_can_cancel_its_own_published_run() -> None:
rows: Final = DefinitionRows()
definitions: Final = WorkflowDefinitionRepository(rows.database)
draft: Final = await definitions.create(bundle(), "team", "author")
assert isinstance(draft, WorkflowDraft)
snapshot: Final = await definitions.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
assert isinstance(snapshot, WorkflowSnapshot)
executions: Final = ExecutionRows()
application: Final = virtual_key(None, key="application-key")
accepted: Final = await invoke_workflow(
Request(
{
"type": "http",
"method": "POST",
"path": "/v1/chat/completions",
"headers": [(b"prefer", b"respond-async")],
}
),
Response(),
application,
{"model": "workflow/example", "messages": _MESSAGES},
definitions,
WorkflowExecutionRepository(executions),
)
assert isinstance(accepted, JSONResponse) and accepted.status_code == 202
run_id: Final = executions.runs[0].run_id
foreign: Final = client_for(rows, executions, virtual_key(None, key="different-application-key"))
assert foreign.post(f"/v1/workflows/runs/{run_id}/cancel").status_code == 404
client: Final = client_for(rows, executions, application)
response: Final = client.post(f"/v1/workflows/runs/{run_id}/cancel")
assert response.status_code == 202, response.text
assert executions.runs[0].cancel_requested
@pytest.mark.asyncio
async def test_prefer_respond_async_returns_published_run_id_without_waiting_for_execution() -> None:
definitions: Final = WorkflowDefinitionRepository(DefinitionRows().database)
@ -224,7 +364,6 @@ async def test_prefer_respond_async_returns_published_run_id_without_waiting_for
assert response.headers["x-litellm-workflow-run-id"] == run.run_id
assert run.status is RunState.PENDING and not run.input.draft and run.output is None
assert database.events == ()
assert not any(query.sql.startswith("UPDATE") for query in database.queries)
@pytest.mark.asyncio
@ -250,7 +389,7 @@ async def test_sdk_retry_after_republish_reuses_pinned_run_and_sums_node_usage()
)
assert isinstance(accepted, JSONResponse) and accepted.status_code == 202
run: Final = database.runs[0]
assert await executions.claim(run.run_id, "worker", NOW + timedelta(seconds=15)) is not None
assert await executions.claim(run.run_id, "worker", datetime.now(timezone.utc) + timedelta(seconds=15)) is not None
for node_id, prompt, completion, cost in (("node_1", 7, 2, 0.1), ("node_2", 11, 3, 0.2)):
assert await executions.checkpoint(
run.run_id,
@ -295,7 +434,7 @@ async def test_sdk_retry_after_republish_reuses_pinned_run_and_sums_node_usage()
assert result == {
"id": f"chatcmpl-workflow-{run.run_id}",
"object": "chat.completion",
"created": int(NOW.timestamp()),
"created": int(run.created_at.timestamp()),
"model": "workflow/example",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": '{"answer": "approved"}'}, "finish_reason": "stop"}

26
tests/unit/proxy/proxy_server/test_routes_chat_completions.py
View file

@ -126,3 +126,29 @@ def test_chat_completion_pipeline_error(client, auth_as, patched_chat_error, pat
response = client.post(path, json=payload)
assert response.status_code == 400
assert "error" in response.json() or response.text != ""
@pytest.mark.parametrize("path", _CHAT_PATHS)
@pytest.mark.parametrize("enabled", [False, True])
def test_workflow_invocations_require_explicit_opt_in(client, auth_as, monkeypatch, path, enabled):
from litellm.proxy.deps import database
from litellm.proxy.management_endpoints import workflow_execution_endpoints
monkeypatch.setattr(proxy_server, "general_settings", {"enable_workflows": True} if enabled else {})
invocation = AsyncMock(return_value=dict(HAPPY_RESPONSE))
connection = AsyncMock(return_value=None)
monkeypatch.setattr(workflow_execution_endpoints, "invoke_workflow", invocation)
monkeypatch.setattr(database, "get_database", connection)
with auth_as():
response = client.post(
path,
json={"model": "workflow/email-processing", "messages": [{"role": "user", "content": "Hello"}]},
)
if enabled:
assert response.status_code == 200
invocation.assert_awaited_once()
connection.assert_awaited_once()
else:
assert response.status_code == 404
invocation.assert_not_awaited()
connection.assert_not_awaited()

37
tests/unit/proxy/spend_tracking/test_budget_reservation.py
View file

@ -150,6 +150,43 @@ async def test_repeated_token_counting_never_touches_a_tiny_budget(
assert spend_counter_store.get_local(counter_key) == pytest.approx(reserved_cost)
@pytest.mark.asyncio
async def test_priced_workflow_wrapper_never_holds_budget_but_child_call_still_reserves(
spend_counter_store: SpendCounterStore,
) -> None:
router: Final = Router(
model_list=[
{
"model_name": "workflow/example",
"litellm_params": {"model": "gpt-4o"},
}
]
)
messages: Final = [{"role": "user", "content": "hello"}]
body: Final = {"model": "workflow/example", "messages": messages}
estimate: Final = estimate_request_max_cost(body, "/v1/chat/completions", router)
assert estimate is not None and estimate > 0
reservation: Final = await reserve_budget_for_request(
request_body=body,
route="/v1/chat/completions",
llm_router=router,
valid_token=UserAPIKeyAuth(token=TINY_BUDGET_KEY_TOKEN, max_budget=0.01, spend=0.0),
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
)
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
assert reservation is None and spend_counter_store.get_local(counter_key) is None
child: Final = await _reserve_for_tiny_budget_key(
"/v1/chat/completions", {"model": "gpt-4o", "messages": messages, "max_tokens": 16}
)
assert child is not None and child["reserved_cost"] > 0
assert spend_counter_store.get_local(counter_key) == pytest.approx(child["reserved_cost"])
await release_unbound_budget_reservation(child)
BEDROCK_SONNET: Final = "us.anthropic.claude-sonnet-4-6"
CONVERSE_BODY: Final = {
"messages": [{"role": "user", "content": [{"text": "Reply with one word: pong"}]}],

77
tests/unit/proxy/workflow_definition_rows.py
View file

@ -2,16 +2,14 @@ import asyncio
from collections.abc import AsyncGenerator, Mapping, Sequence
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from typing import Final, cast # noqa: TID251 # the test database implements the delegates exercised by these routes
from typing import Final
from pydantic import TypeAdapter
from pydantic import BaseModel, ConfigDict
from litellm.models.workflow_definition import WorkflowBundle
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.repositories.workflow_definition_repository import DefinitionDatabase
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
_COMPARISON: Final = TypeAdapter(Mapping[str, str])
def bundle(name: str = "example", model: str = "gateway-model") -> WorkflowBundle:
@ -41,17 +39,34 @@ def caller(user_id: str | None = "author", team_id: str = "team") -> UserAPIKeyA
def _matches(row: Mapping[str, object], where: Mapping[str, object]) -> bool:
return all(
str(row.get(key)) > _COMPARISON.validate_python(value)["gt"]
if isinstance(value, Mapping) and "gt" in value
else row.get(key) == value
for key, value in where.items()
)
for key, value in where.items():
if key == "OR":
if not any(_matches(row, condition) for condition in value):
return False
elif key == "AND":
if not all(_matches(row, condition) for condition in value):
return False
elif isinstance(value, Mapping):
actual = row.get(key)
for operation, expected in value.items():
if operation == "in" and actual not in expected:
return False
if operation == "gt" and (actual is None or actual <= expected):
return False
if operation == "lt" and (actual is None or actual >= expected):
return False
elif row.get(key) != value:
return False
return True
class _Row(BaseModel):
model_config = ConfigDict(extra="allow")
class Rows:
def __init__(self, rows: tuple[Mapping[str, object], ...] = ()) -> None:
self.rows = rows
self.rows = tuple({"updated_at": NOW, **row} for row in rows)
async def find_first(
self,
@ -62,7 +77,8 @@ class Rows:
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
distinct: Sequence[str] | None = None,
) -> object | None:
return next((row for row in self.rows if _matches(row, where or {})), None)
rows = await self.find_many(take=1, where=where, order=order)
return rows[0] if rows else None
async def find_unique(
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
@ -79,8 +95,12 @@ class Rows:
order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None,
distinct: Sequence[str] | None = None,
) -> Sequence[object]:
matches: Final = tuple(row for row in self.rows if _matches(row, where or {}))
return sorted(matches, key=lambda row: str(row.get("workflow_id")))[:take]
matches = [row for row in self.rows if _matches(row, where or {})]
orders = [order] if isinstance(order, Mapping) else list(order or [])
for clause in reversed(orders):
for key, direction in clause.items():
matches.sort(key=lambda row: row[key], reverse=direction == "desc")
return [_Row.model_validate(row) for row in matches[:take]]
async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> object:
if "etag" in data and any(
@ -91,12 +111,21 @@ class Rows:
raise UniqueViolationError({"is_panic": False, "message": "duplicate", "error_code": "P2002"})
row: Final = {
"created_at": NOW,
"updated_at": NOW,
"published_etag": None,
"published_snapshot_id": None,
**data,
}
self.rows = (*self.rows, row)
return row
return _Row.model_validate(row)
async def count(self, where: Mapping[str, object] | None = None, take: int | None = None) -> int:
return len(await self.find_many(where=where, take=take))
async def update_many(self, data: Mapping[str, object], where: Mapping[str, object]) -> int:
count = sum(_matches(row, where) for row in self.rows)
self.rows = tuple({**row, **data} if _matches(row, where) else row for row in self.rows)
return count
async def update(
self, data: Mapping[str, object], where: Mapping[str, object], include: Mapping[str, object] | None = None
@ -123,15 +152,21 @@ class DefinitionRows:
)
)
self.lock = asyncio.Lock()
async def query_raw(self, query: str, *args: object) -> object:
return ()
self.db = self
self.writer_db = self
@property
def database(self) -> DefinitionDatabase:
return cast(DefinitionDatabase, self) # cast-ok: isolated Prisma boundary implements the used table operations
def database(self) -> "DefinitionRows":
return self
@asynccontextmanager
async def tx(self) -> AsyncGenerator["DefinitionRows"]:
async with self.lock:
yield self
definitions = self.litellm_workflowdefinition.rows
snapshots = self.litellm_workflowsnapshot.rows
try:
yield self
except BaseException:
self.litellm_workflowdefinition.rows = definitions
self.litellm_workflowsnapshot.rows = snapshots
raise

20
tests/unit/proxy/workflows/test_definitions.py
View file

@ -4,7 +4,7 @@ from typing import Final
import pytest
from litellm.models.workflow_definition import DefinitionError, WorkflowBundle, WorkflowDraft, WorkflowSnapshot
from litellm.proxy.workflows.definitions import freeze_draft, published_workflow, validate_bundle
from litellm.proxy.workflows.definitions import prepare_draft, published_workflow, validate_bundle
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle, caller
@ -85,14 +85,28 @@ def test_validation_rejects_schema_reference_scope_changes(keyword: str) -> None
)
@pytest.mark.parametrize(
"schema",
[
{"type": "object", "properties": {"$ref": {"type": "string"}, "$id": {"type": "string"}}},
{"const": {"$ref": "https://example.com/literal", "$id": "literal"}},
{"examples": [{"$dynamicRef": "literal"}], "default": {"$ref": "literal"}},
],
)
def test_validation_preserves_reference_named_properties_and_literal_values(schema: dict[str, object]) -> None:
original: Final = bundle().model_dump(mode="json", by_alias=True, exclude_unset=True)
with_literal: Final = WorkflowBundle.model_validate({**original, "input_schema": schema})
assert validate_bundle(with_literal) is None
@pytest.mark.asyncio
async def test_application_identity_can_resolve_published_workflow_but_cannot_freeze_draft() -> None:
async def test_application_identity_can_resolve_published_workflow_but_cannot_prepare_draft() -> None:
store: Final = WorkflowDefinitionRepository(DefinitionRows().database)
draft: Final = await store.create(bundle(), "team", "author")
assert isinstance(draft, WorkflowDraft)
snapshot: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
assert isinstance(snapshot, WorkflowSnapshot)
assert await published_workflow(store, caller(None), "example") == snapshot
denied: Final = await freeze_draft(store, caller(None), draft.workflow_id, draft.etag)
denied: Final = await prepare_draft(store, caller(None), draft.workflow_id, draft.etag)
assert isinstance(denied, DefinitionError)
assert denied.code == "not_found"

47
tests/unit/proxy/workflows/test_lifespan.py Normal file
View file

@ -0,0 +1,47 @@
from unittest.mock import AsyncMock, Mock
import pytest
from fastapi import HTTPException
from litellm.proxy import proxy_server
from litellm.proxy._types import ConfigGeneralSettings
from litellm.proxy.workflows import lifespan
def test_workflows_default_to_disabled_and_require_explicit_enablement(monkeypatch: pytest.MonkeyPatch) -> None:
assert ConfigGeneralSettings().enable_workflows is False
monkeypatch.setattr(proxy_server, "general_settings", {})
assert not lifespan.workflows_enabled()
with pytest.raises(HTTPException) as rejected:
lifespan.require_workflows_enabled()
assert rejected.value.status_code == 404
monkeypatch.setattr(proxy_server, "general_settings", {"enable_workflows": True})
assert ConfigGeneralSettings(enable_workflows=True).enable_workflows is True
assert lifespan.workflows_enabled()
lifespan.require_workflows_enabled()
@pytest.mark.asyncio
@pytest.mark.parametrize("settings", [{}, {"enable_workflows": False}])
async def test_disabled_workflows_do_not_create_repositories_or_start_runner(
monkeypatch: pytest.MonkeyPatch, settings: dict[str, bool]
) -> None:
monkeypatch.setattr(proxy_server, "general_settings", settings)
definitions = AsyncMock(side_effect=AssertionError("Disabled workflows must not initialize storage"))
monkeypatch.setattr(lifespan, "get_workflow_definition_repository", definitions)
async with lifespan.workflow_lifespan(Mock()):
pass
definitions.assert_not_awaited()
@pytest.mark.asyncio
async def test_enabled_workflows_start_runner_and_close_it_on_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(proxy_server, "general_settings", {"enable_workflows": True})
monkeypatch.setattr(lifespan, "get_workflow_definition_repository", AsyncMock(return_value=Mock()))
monkeypatch.setattr(lifespan, "get_workflow_execution_repository", AsyncMock(return_value=Mock()))
runner = Mock(start=AsyncMock(), close=AsyncMock())
monkeypatch.setattr(lifespan, "WorkflowRunner", Mock(return_value=runner))
async with lifespan.workflow_lifespan(Mock()):
runner.start.assert_awaited_once()
runner.close.assert_not_awaited()
runner.close.assert_awaited_once()

60
tests/unit/proxy/workflows/test_runner.py
View file

@ -1,7 +1,8 @@
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta
from datetime import datetime, timedelta, tzinfo
from typing import Final
from unittest.mock import AsyncMock
import pytest
from pydantic import JsonValue
@ -24,12 +25,22 @@ from litellm.models.workflow_execution import (
workflow_call_id,
)
from litellm.proxy.workflows.runner import SchemaValidator, WorkflowRunner
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository, persist_snapshot
from litellm.repositories.workflow_execution_repository import WorkflowExecutionRepository
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle
from tests.unit.repositories.test_workflow_execution_repository import NOW, ExecutionRows, pending_run
@pytest.fixture(autouse=True)
def fixed_repository_clock(monkeypatch: pytest.MonkeyPatch) -> None:
class RepositoryClock(datetime):
@classmethod
def now(cls, tz: tzinfo | None = None) -> datetime:
return NOW
monkeypatch.setattr("litellm.repositories.workflow_execution_repository.datetime", RepositoryClock)
class FixedClock:
def __init__(self) -> None:
self.value = NOW
@ -127,8 +138,9 @@ async def runner_inputs(
definitions: Final = WorkflowDefinitionRepository(DefinitionRows().database)
draft: Final = await definitions.create(definition, "team", "author")
assert isinstance(draft, WorkflowDraft)
snapshot: Final = await definitions.freeze_draft(draft.workflow_id, "team", draft.etag)
snapshot: Final = await definitions.prepare_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(snapshot, WorkflowSnapshot)
await persist_snapshot(definitions.prisma_client.writer_db, snapshot)
run: Final = pending_run(snapshot_id=snapshot.snapshot_id, workflow_id=snapshot.workflow_id, deadline=deadline)
return definitions, ExecutionRows((run,)), run
@ -221,6 +233,48 @@ async def test_first_nonretryable_error_stops_downstream_and_keeps_failure_trace
)
@pytest.mark.asyncio
@pytest.mark.parametrize("retry", [False, True])
async def test_cancellation_stops_downstream_nodes_and_retries_before_the_heartbeat(
retry: bool,
) -> None:
definitions, database, run = await runner_inputs(two_nodes(attempts=2))
repository: Final = WorkflowExecutionRepository(database)
receipt: Final = ChildCallResult([{"role": "user", "content": "next"}], 7, 3, "billed-request", 0.125)
async def cancel_after_first_call(
**_kwargs: object,
) -> ChildCallResult | ExecutionError:
await repository.cancel(run.run_id)
return ExecutionError("provider_error", "Retry later", receipt=receipt) if retry else receipt
chat: Final = AsyncMock(side_effect=cancel_after_first_call)
runner: Final = WorkflowRunner(
repository,
definitions,
chat,
clock=FixedClock(),
validate_input=ValidationReply(),
)
try:
await runner._run(run.run_id) # pyright: ignore[reportPrivateUsage] # control dispatch before the heartbeat
finally:
await runner.close()
assert chat.await_count == 1
finished: Final = await repository.get(run.run_id)
assert finished is not None and finished.status is RunState.CANCELLED
assert finished.output is None and finished.cancel_requested
steps: Final = await repository.steps(run.run_id)
assert len(steps) == 1
assert (
steps[0].node_id,
steps[0].attempts,
steps[0].request_id,
steps[0].response_cost,
) == ("node_1", 1, "billed-request", 0.125)
@pytest.mark.asyncio
async def test_retry_preserves_run_session_step_and_absolute_deadline() -> None:
definitions, database, run = await runner_inputs(

38
tests/unit/repositories/test_workflow_definition_repository.py
View file

@ -4,7 +4,7 @@ from typing import Final
import pytest
from litellm.models.workflow_definition import DefinitionError, WorkflowDraft, WorkflowSnapshot
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository, persist_snapshot
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, bundle
@ -15,8 +15,10 @@ async def test_edit_and_republish_leave_old_snapshot_and_frozen_draft_unchanged(
assert isinstance(draft, WorkflowDraft)
published: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
assert isinstance(published, WorkflowSnapshot)
frozen: Final = await store.freeze_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(frozen, WorkflowSnapshot)
prepared: Final = await store.prepare_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(prepared, WorkflowSnapshot)
assert await store.get_snapshot(prepared.snapshot_id) is None
frozen: Final = await persist_snapshot(store.prisma_client.writer_db, prepared)
assert frozen.published_by is None
updated: Final = await store.replace(draft.workflow_id, "team", draft.etag, bundle(model="second"), "other")
assert isinstance(updated, WorkflowDraft)
@ -69,3 +71,33 @@ async def test_publish_and_edit_race_never_snapshots_unreviewed_content() -> Non
assert published.digest == draft.digest
assert edited.digest != published.digest
assert await store.get_published("example", "team") == published
@pytest.mark.asyncio
async def test_edit_rereads_publication_after_lock_even_when_audit_timestamp_is_unchanged(
monkeypatch: pytest.MonkeyPatch,
) -> None:
database: Final = DefinitionRows()
store: Final = WorkflowDefinitionRepository(database.database)
draft: Final = await store.create(bundle(), "team", "author")
assert isinstance(draft, WorkflowDraft)
table: Final = database.litellm_workflowdefinition
stale: Final = await table.find_first(where={"workflow_id": draft.workflow_id})
published: Final = await store.publish(draft.workflow_id, "team", draft.etag, draft.digest, "reviewer")
assert isinstance(published, WorkflowSnapshot)
original: Final = table.find_first
pending = True
async def read_once_before_publication(*args, **kwargs):
nonlocal pending
if pending:
pending = False
return stale
return await original(*args, **kwargs)
monkeypatch.setattr(table, "find_first", read_once_before_publication)
edited: Final = await store.replace(draft.workflow_id, "team", draft.etag, bundle(model="edited"), "other")
assert isinstance(edited, WorkflowDraft)
assert edited.published_snapshot_id == published.snapshot_id
stored: Final = await store.get_draft(draft.workflow_id)
assert stored is not None and stored.contributors == ("other",)

635
tests/unit/repositories/test_workflow_execution_repository.py
View file

@ -1,13 +1,13 @@
import re
import asyncio
from collections.abc import AsyncGenerator, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Final
from uuid import uuid4
import pytest
from pydantic import JsonValue, TypeAdapter
from litellm.models.workflow_definition import WorkflowSnapshot
from litellm.models.workflow_execution import (
EngineRun,
ExecutionError,
@ -17,13 +17,23 @@ from litellm.models.workflow_execution import (
RunState,
StepTrace,
)
from litellm.repositories.workflow_execution_repository import ExecutionQueries, WorkflowExecutionRepository
from litellm.repositories import workflow_execution_repository
from litellm.repositories.workflow_definition_repository import WorkflowDefinitionRepository
from litellm.repositories.workflow_execution_repository import WorkflowExecutionRepository
from tests.unit.proxy.workflow_definition_rows import DefinitionRows, Rows, bundle
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
_STRING: Final = TypeAdapter(str)
_INTEGER: Final = TypeAdapter(int)
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_TRACE: Final = TypeAdapter(dict[str, StepTrace])
class FixedDatetime(datetime):
@classmethod
def now(cls, tz=None):
return NOW
@pytest.fixture(autouse=True)
def freeze_repository_clock(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(workflow_execution_repository, "datetime", FixedDatetime)
def pending_run(
@ -46,383 +56,304 @@ def pending_run(
created_at=NOW,
updated_at=NOW,
input=RunInput(
messages=({"role": "user", "content": "plan"},),
end_user_id="end-user",
deadline=deadline,
draft=True,
messages=({"role": "user", "content": "plan"},), end_user_id="end-user", deadline=deadline, draft=True
),
request_digest="request-digest",
)
@dataclass(frozen=True, slots=True)
class Query:
sql: str
args: tuple[object, ...]
in_transaction: bool
def _data(data: Mapping[str, object]) -> dict[str, object]:
from prisma import Json
return {key: value.data if isinstance(value, Json) else value for key, value in data.items()}
class RunRows(Rows):
fail_create = False
async def create(self, data: Mapping[str, object], include=None):
if self.fail_create:
raise RuntimeError("Injected run write failure")
return await super().create(
data={
"run_id": str(uuid4()),
"session_id": str(uuid4()),
"status": RunState.PENDING.value,
"output": None,
"metadata": None,
"lease_token": None,
"lease_expires_at": None,
"cancel_requested": False,
**_data(data),
}
)
async def update_many(self, data: Mapping[str, object], where: Mapping[str, object]) -> int:
return await super().update_many(data=_data(data), where=where)
class EventRows(Rows):
async def create(self, data: Mapping[str, object], include=None):
return await super().create(data={"event_id": str(uuid4()), **_data(data)})
class ExecutionRows:
"""Small model-action fake shared by endpoint and runner behavior tests."""
def __init__(self, runs: tuple[EngineRun, ...] = ()) -> None:
self.litellm_workflowrun = RunRows()
self.litellm_workflowevent = EventRows()
self.litellm_workflowsnapshot = Rows()
self.litellm_spendlogs = Rows()
self.litellm_teamtable = Rows(({"team_id": "team", "updated_at": NOW},))
self.db = self
self.writer_db = self
self.lock = asyncio.Lock()
self.runs = runs
self.keys: tuple[tuple[str, str | None], ...] = ()
self.events: tuple[StepTrace, ...] = ()
self.queries: tuple[Query, ...] = ()
self.in_transaction = False
@property
def runs(self) -> tuple[EngineRun, ...]:
return tuple(EngineRun.model_validate(row) for row in self.litellm_workflowrun.rows)
@runs.setter
def runs(self, runs: tuple[EngineRun, ...]) -> None:
self.litellm_workflowrun.rows = tuple({"engine_owned": True, **run.model_dump()} for run in runs)
@property
def events(self) -> tuple[StepTrace, ...]:
return tuple(StepTrace.model_validate(row["data"]["trace"]) for row in self.litellm_workflowevent.rows)
@events.setter
def events(self, events: tuple[StepTrace, ...]) -> None:
self.litellm_workflowevent.rows = tuple(
{
"run_id": step.run_id,
"event_type": "liteflow.step",
"step_name": step.node_id,
"sequence_number": index,
"data": {"trace": step.model_dump(mode="json")},
}
for index, step in enumerate(events)
)
@asynccontextmanager
async def tx(self) -> AsyncGenerator[ExecutionQueries]:
self.in_transaction = True
try:
yield self
finally:
self.in_transaction = False
async def tx(self) -> AsyncGenerator["ExecutionRows"]:
async with self.lock:
runs = self.litellm_workflowrun.rows
events = self.litellm_workflowevent.rows
snapshots = self.litellm_workflowsnapshot.rows
try:
yield self
except BaseException:
self.litellm_workflowrun.rows = runs
self.litellm_workflowevent.rows = events
self.litellm_workflowsnapshot.rows = snapshots
raise
def _rows(self, runs: tuple[EngineRun, ...]) -> tuple[Mapping[str, object], ...]:
return tuple(run.model_dump() for run in runs)
def _replace(self, run: EngineRun) -> tuple[Mapping[str, object], ...]:
self.runs = tuple(run if current.run_id == run.run_id else current for current in self.runs)
return self._rows((run,))
async def admit(repository: WorkflowExecutionRepository, key: str | None = "key", **overrides):
original = pending_run()
return await repository.admit(
**{
"workflow_id": original.workflow_id,
"snapshot_id": original.snapshot_id,
"team_id": original.team_id,
"key_hash": original.created_by,
"input": original.input,
"digest": original.request_digest,
"idempotency_key": key,
**overrides,
}
)
async def query_raw(self, query: str, *args: object) -> object:
self.queries = (*self.queries, Query(query, args, self.in_transaction))
parameters: Final = frozenset(int(match.group(1)) for match in re.finditer(r"\$(\d+)", query))
assert parameters == frozenset(range(1, len(args) + 1)), (query, args)
if "pg_advisory_xact_lock" in query:
assert query.startswith("SELECT true AS acquired FROM "), "Prisma cannot decode a PostgreSQL void column"
return ({"acquired": True},)
if "idempotency_key =" in query:
owner, team, workflow, key = args if "created_by = $1" in query else (args[2], args[1], args[0], args[3])
return self._rows(
tuple(
run
for run in self.runs
if key is not None
and (run.created_by, run.team_id, run.workflow_id) == (owner, team, workflow)
and (run.run_id, key) in self.keys
)
)
if query.startswith('INSERT INTO "LiteLLM_WorkflowRun"'):
run: Final = EngineRun.model_validate(
{
"run_id": args[0],
"session_id": args[1],
"created_by": args[2],
"workflow_id": args[3],
"snapshot_id": args[4],
"team_id": args[5],
"input": RunInput.model_validate_json(_STRING.validate_python(args[6])),
"request_digest": args[7],
"status": RunState.PENDING,
"created_at": NOW,
"updated_at": NOW,
"output": None,
"metadata": None,
"lease_token": None,
"lease_expires_at": None,
}
)
self.runs = (*self.runs, run)
self.keys = (*self.keys, (run.run_id, None if args[8] is None else _STRING.validate_python(args[8])))
return self._rows((run,))
if "LIMIT 64" in query:
return tuple(
{"run_id": run.run_id}
for run in self.runs
if run.team_id == args[0] and run.status in {RunState.PENDING, RunState.RUNNING}
)[:64]
if "ORDER BY created_at DESC, run_id DESC" in query:
boundary: Final = None if args[4] is None else datetime.fromisoformat(_STRING.validate_python(args[4]))
matching: Final = tuple(
run
for run in self.runs
if run.workflow_id == args[0]
and run.team_id == args[1]
and (args[2] is None or run.created_by == args[2])
and (args[3] is None or run.status.value == args[3])
and (boundary is None or (run.created_at, run.run_id) < (boundary, _STRING.validate_python(args[5])))
)
return self._rows(
tuple(sorted(matching, key=lambda run: (run.created_at, run.run_id), reverse=True))[
: _INTEGER.validate_python(args[6])
]
)
if query.startswith("SELECT trace FROM"):
latest: Final = {(step.node_id, step.attempts): step for step in self.events if step.run_id == args[0]}
return tuple(
{"trace": step.model_dump(mode="json")}
for step in sorted(latest.values(), key=lambda step: (step.started_at, step.node_id))
)
if query.startswith('INSERT INTO "LiteLLM_WorkflowEvent"'):
payload: Final = _TRACE.validate_json(_STRING.validate_python(args[3]))
self.events = (*self.events, payload["trace"])
return ({"event_id": args[0]},)
assert query.startswith(
(
'SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id',
'SELECT run_id FROM "LiteLLM_WorkflowRun" WHERE run_id',
'UPDATE "LiteLLM_WorkflowRun"',
)
), query
current: Final = next((run for run in self.runs if run.run_id == args[0]), None)
if current is None:
return ()
if query.startswith('SELECT * FROM "LiteLLM_WorkflowRun" WHERE run_id'):
return self._rows((current,))
if "SET status = 'running'" in query:
if current.status is not RunState.PENDING:
return ()
return self._replace(
current.model_copy(
update={
"status": RunState.RUNNING,
"lease_token": args[1],
"lease_expires_at": datetime.fromisoformat(_STRING.validate_python(args[2])),
}
)
)
if "SET cancel_requested = true" in query:
return self._replace(current.model_copy(update={"cancel_requested": True}))
owned: Final = (
current.status is RunState.RUNNING
and current.lease_token == args[1]
and current.lease_expires_at is not None
and current.lease_expires_at > NOW
)
if "FOR UPDATE" in query:
return ({"run_id": current.run_id},) if owned else ()
if "SET status = $3" in query:
if not owned:
return ()
return self._replace(
current.model_copy(
update={
"status": RunState(args[2]),
"output": _JSON.validate_json(_STRING.validate_python(args[3])),
"metadata": {"error": args[4]},
"lease_token": None,
"lease_expires_at": None,
}
)
)
raise AssertionError(f"Unexpected database operation: {query}")
def trace(state: RunState = RunState.RUNNING, attempt: int = 1) -> StepTrace:
return StepTrace(
step_id="step-1",
run_id="run-1",
node_id="node-1",
state=state,
input=None,
attempts=attempt,
deadline=NOW + timedelta(seconds=30),
started_at=NOW,
)
@pytest.mark.asyncio
async def test_idempotent_admission_pins_snapshot_and_rejects_different_input() -> None:
database: Final = ExecutionRows()
repository: Final = WorkflowExecutionRepository(database)
original: Final = pending_run()
key: Final = "key'; SELECT secret"
first: Final = await repository.admit(
workflow_id=original.workflow_id,
snapshot_id="snapshot-1",
team_id=original.team_id,
key_hash=original.created_by,
input=original.input,
digest=original.request_digest,
idempotency_key=key,
)
replay: Final = await repository.admit(
workflow_id=original.workflow_id,
snapshot_id="snapshot-2",
team_id=original.team_id,
key_hash=original.created_by,
input=original.input,
digest=original.request_digest,
idempotency_key=key,
)
conflict: Final = await repository.admit(
workflow_id=original.workflow_id,
snapshot_id="snapshot-2",
team_id=original.team_id,
key_hash=original.created_by,
input=original.input,
digest="different",
idempotency_key=key,
)
assert isinstance(first, RunAdmission)
assert replay == RunAdmission(first.run, False)
assert first.created and first.run.snapshot_id == "snapshot-1"
assert conflict == ExecutionError("idempotency_conflict", "Idempotency key was used with different input")
async def test_concurrent_idempotent_admission_pins_one_snapshot_and_rejects_different_input() -> None:
database = ExecutionRows()
repository = WorkflowExecutionRepository(database)
first, replay = await asyncio.gather(admit(repository), admit(repository, snapshot_id="edited"))
assert isinstance(first, RunAdmission) and isinstance(replay, RunAdmission)
assert sorted((first.created, replay.created)) == [False, True]
assert first.run == replay.run
assert len(database.runs) == 1
assert all(query.in_transaction for query in database.queries)
assert database.queries[0].args == (original.team_id,)
assert database.queries[0].sql.startswith("SELECT true AS acquired FROM pg_advisory_xact_lock(")
assert all("key'; SELECT secret" not in query.sql for query in database.queries)
insert: Final = next(query for query in database.queries if query.sql.startswith("INSERT"))
assert "timezone('UTC', NOW()), timezone('UTC', NOW())" in insert.sql
assert insert.args[2:6] == (original.created_by, original.workflow_id, "snapshot-1", original.team_id)
assert RunInput.model_validate_json(_STRING.validate_python(insert.args[6])) == original.input
assert insert.args[7:] == (original.request_digest, key)
assert await repository.find_idempotent("workflow-1", "team", "owner-key", "key") == first.run
assert await admit(repository, digest="different") == ExecutionError(
"idempotency_conflict", "Idempotency key was used with different input"
)
@pytest.mark.asyncio
async def test_null_idempotency_and_nullable_database_columns_round_trip_without_coalescing_runs() -> None:
database: Final = ExecutionRows()
repository: Final = WorkflowExecutionRepository(database)
original: Final = pending_run()
first: Final = await repository.admit(
workflow_id=original.workflow_id,
snapshot_id=original.snapshot_id,
team_id=original.team_id,
key_hash=original.created_by,
input=original.input,
digest=original.request_digest,
idempotency_key=None,
)
second: Final = await repository.admit(
workflow_id=original.workflow_id,
snapshot_id=original.snapshot_id,
team_id=original.team_id,
key_hash=original.created_by,
input=original.input,
digest=original.request_digest,
idempotency_key=None,
)
async def test_null_idempotency_keys_create_distinct_runs_and_nullable_output_round_trips() -> None:
repository = WorkflowExecutionRepository(ExecutionRows())
first, second = await admit(repository, None), await admit(repository, None)
assert isinstance(first, RunAdmission) and isinstance(second, RunAdmission)
assert first.created and second.created and first.run.run_id != second.run.run_id
assert (first.run.output, first.run.metadata, first.run.lease_token, first.run.lease_expires_at) == (None,) * 4
lookup: Final = next(query for query in database.queries if "idempotency_key =" in query.sql)
assert lookup.args == (original.created_by, original.team_id, original.workflow_id, None)
assert "idempotency_key = $4::text" in lookup.sql
assert first.run.run_id != second.run.run_id
assert first.run.output is None and first.run.metadata is None
assert await repository.claim(first.run.run_id, "lease", NOW + timedelta(seconds=15)) is not None
claim: Final = database.queries[-1]
assert claim.args == (first.run.run_id, "lease", (NOW + timedelta(seconds=15)).isoformat())
assert "lease_expires_at = timezone('UTC', $3::text::timestamptz)" in claim.sql
assert await repository.finish(first.run.run_id, "lease", RunState.COMPLETED, {"value": None}, None)
finish: Final = database.queries[-1]
assert finish.args == (first.run.run_id, "lease", "completed", '{"value": null}', None)
assert "$4::jsonb" in finish.sql and "$5::text" in finish.sql
finished: Final = await repository.get(first.run.run_id)
assert finished is not None
assert (finished.output, finished.metadata, finished.lease_token) == ({"value": None}, {"error": None}, None)
finished = await repository.get(first.run.run_id)
assert finished is not None and finished.output == {"value": None} and finished.metadata == {"error": None}
@pytest.mark.asyncio
async def test_history_cursor_uses_timestamp_and_id_with_lookahead_and_scoped_parameters() -> None:
database: Final = ExecutionRows(
async def test_concurrent_admission_never_exceeds_team_capacity_across_workflows() -> None:
database = ExecutionRows(tuple(pending_run(f"run-{index}") for index in range(63)))
repository = WorkflowExecutionRepository(database)
outcomes = await asyncio.gather(admit(repository, "one"), admit(repository, "two", workflow_id="other"))
assert sum(isinstance(outcome, RunAdmission) for outcome in outcomes) == 1
assert sum(isinstance(outcome, ExecutionError) and outcome.code == "capacity_exceeded" for outcome in outcomes) == 1
assert len(database.runs) == 64
@pytest.mark.asyncio
async def test_draft_snapshot_commits_with_admitted_run_only_and_rolls_back_failed_write() -> None:
definitions = DefinitionRows()
store = WorkflowDefinitionRepository(definitions.database)
draft = await store.create(bundle(), "team", "author")
snapshot = await store.prepare_draft(draft.workflow_id, "team", draft.etag)
assert isinstance(snapshot, WorkflowSnapshot)
database = ExecutionRows()
database.litellm_workflowsnapshot = definitions.litellm_workflowsnapshot
repository = WorkflowExecutionRepository(database)
assert database.litellm_workflowsnapshot.rows == ()
database.litellm_workflowrun.fail_create = True
with pytest.raises(RuntimeError, match="Injected run write failure"):
await admit(
repository, draft_snapshot=snapshot, workflow_id=draft.workflow_id, snapshot_id=snapshot.snapshot_id
)
assert database.litellm_workflowsnapshot.rows == () and database.runs == ()
database.litellm_workflowrun.fail_create = False
admitted = await admit(
repository, draft_snapshot=snapshot, workflow_id=draft.workflow_id, snapshot_id=snapshot.snapshot_id
)
assert isinstance(admitted, RunAdmission) and len(database.litellm_workflowsnapshot.rows) == 1
next_snapshot = await store.prepare_draft(draft.workflow_id, "team", draft.etag)
replay = await admit(
repository, draft_snapshot=next_snapshot, workflow_id=draft.workflow_id, snapshot_id=next_snapshot.snapshot_id
)
assert replay == RunAdmission(admitted.run, False)
assert len(database.litellm_workflowsnapshot.rows) == 1
database.runs = tuple(pending_run(f"full-{index}") for index in range(64))
refused = await admit(
repository,
"new-key",
draft_snapshot=next_snapshot,
workflow_id=draft.workflow_id,
snapshot_id=next_snapshot.snapshot_id,
)
assert isinstance(refused, ExecutionError) and refused.code == "capacity_exceeded"
assert len(database.litellm_workflowsnapshot.rows) == 1
@pytest.mark.asyncio
async def test_stale_worker_cannot_checkpoint_renew_or_finish_after_reclaim() -> None:
expired = pending_run().model_copy(
update={
"status": RunState.RUNNING,
"lease_token": "old",
"lease_expires_at": NOW - timedelta(seconds=1),
}
)
database = ExecutionRows((expired,))
repository = WorkflowExecutionRepository(database)
assert await repository.recover(1) == (expired,)
assert await repository.claim("run-1", "new", NOW + timedelta(seconds=15)) is not None
assert await repository.claim("run-1", "other", NOW + timedelta(seconds=15)) is None
assert not await repository.renew("run-1", "old", NOW + timedelta(seconds=20))
assert not await repository.checkpoint("run-1", "old", trace())
assert not await repository.finish("run-1", "old", RunState.COMPLETED, "forged", None)
assert database.events == ()
@pytest.mark.asyncio
async def test_cancellation_fences_new_dispatch_but_preserves_completed_receipts() -> None:
database = ExecutionRows((pending_run(),))
repository = WorkflowExecutionRepository(database)
assert await repository.claim("run-1", "lease", NOW + timedelta(seconds=15)) is not None
assert await repository.checkpoint("run-1", "lease", trace())
await repository.cancel("run-1")
assert not await repository.checkpoint("run-1", "lease", trace(attempt=2))
completed = trace(RunState.COMPLETED).model_copy(update={"output": "paid-result", "response_cost": 0.25})
assert await repository.checkpoint("run-1", "lease", completed)
assert await repository.steps("run-1") == (completed.model_copy(update={"request_ids": ()}),)
@pytest.mark.asyncio
async def test_attempt_receipts_aggregate_once_and_remain_scoped_to_run_principal() -> None:
database = ExecutionRows((pending_run(),))
repository = WorkflowExecutionRepository(database)
assert await repository.claim("run-1", "lease", NOW + timedelta(seconds=15)) is not None
first = trace(RunState.FAILED).model_copy(
update={
"call_id": "call-1",
"request_id": "failed",
"response_cost": 0.125,
"prompt_tokens": 4,
"completion_tokens": 2,
}
)
second = trace(RunState.COMPLETED, 2).model_copy(update={"call_id": "call-2", "output": "done"})
await asyncio.gather(*(repository.checkpoint("run-1", "lease", step) for step in (first, first, second, second)))
database.litellm_spendlogs.rows = (
{
"request_id": "paid",
"litellm_call_id": "call-2",
"api_key": "owner-key",
"team_id": "team",
"session_id": "session-run-1",
"spend": 0.25,
"prompt_tokens": 6,
"completion_tokens": 3,
},
{
"request_id": "foreign",
"litellm_call_id": "call-2",
"api_key": "other-key",
"team_id": "team",
"session_id": "session-run-1",
"spend": 99,
"prompt_tokens": 99,
"completion_tokens": 99,
},
)
(step,) = await repository.steps("run-1")
assert step.attempts == 2 and step.output == "done"
assert (step.prompt_tokens, step.completion_tokens, step.response_cost) == (10, 5, 0.375)
assert step.request_ids == ("failed", "paid")
assert len(database.events) == 4
assert tuple(row["sequence_number"] for row in database.litellm_workflowevent.rows) == (0, 1, 2, 3)
@pytest.mark.asyncio
async def test_history_cursor_uses_timestamp_tiebreaker_and_principal_scope() -> None:
database = ExecutionRows(
(
pending_run("run-c"),
pending_run("run-a"),
pending_run("run-b"),
pending_run("foreign-team", team_id="other-team"),
pending_run("foreign-workflow", workflow_id="other-workflow"),
pending_run("foreign-key", key_hash="other-key"),
pending_run("foreign", team_id="other"),
pending_run("foreign-key", key_hash="other"),
)
)
repository: Final = WorkflowExecutionRepository(database)
first: Final = await repository.list_runs("workflow-1", "team", "owner-key", RunState.PENDING, 2, None)
assert isinstance(first, RunPage)
repository = WorkflowExecutionRepository(database)
first = await repository.list_runs("workflow-1", "team", "owner-key", RunState.PENDING, 2, None)
assert isinstance(first, RunPage) and first.next_cursor is not None
assert tuple(run.run_id for run in first.runs) == ("run-c", "run-b")
assert first.count == 2 and first.next_cursor is not None
assert database.queries[-1].args == ("workflow-1", "team", "owner-key", "pending", None, None, 3)
database.runs = (*database.runs, pending_run("run-d"))
tail: Final = await repository.list_runs("workflow-1", "team", "owner-key", RunState.PENDING, 2, first.next_cursor)
assert isinstance(tail, RunPage)
tail = await repository.list_runs("workflow-1", "team", "owner-key", RunState.PENDING, 2, first.next_cursor)
assert isinstance(tail, RunPage) and tail.next_cursor is None
assert tuple(run.run_id for run in tail.runs) == ("run-a",)
assert tail.count == 1 and tail.next_cursor is None
query: Final = database.queries[-1]
assert query.args == ("workflow-1", "team", "owner-key", "pending", NOW.isoformat(), "run-b", 3)
assert "(created_at, run_id) < (timezone('UTC', $5::text::timestamptz), $6::text)" in query.sql
assert "ORDER BY created_at DESC, run_id DESC LIMIT $7" in query.sql
assert "workflow_id = $1 AND team_id = $2" in query.sql
assert "created_by = $3::text" in query.sql and "status = $4::text" in query.sql
assert "engine_owned = true" in query.sql
@pytest.mark.asyncio
async def test_invalid_cursor_never_queries_and_unfiltered_history_passes_sql_nulls() -> None:
database: Final = ExecutionRows()
repository: Final = WorkflowExecutionRepository(database)
assert await repository.list_runs("workflow-1", "team", None, None, 1, "not-a-cursor") == ExecutionError(
"invalid_cursor", "Invalid run cursor"
)
assert database.queries == ()
assert await repository.list_runs("workflow-1", "team", None, None, 1, None) == RunPage(
runs=(), count=0, next_cursor=None
)
query: Final = database.queries[-1]
assert query.args == ("workflow-1", "team", None, None, None, None, 2)
assert "($3::text IS NULL" in query.sql and "($4::text IS NULL" in query.sql and "($5::text IS NULL" in query.sql
@pytest.mark.asyncio
async def test_lost_lease_cannot_write_trace_or_terminal_result() -> None:
database: Final = ExecutionRows((pending_run(),))
repository: Final = WorkflowExecutionRepository(database)
assert await repository.claim("run-1", "owner-lease", NOW + timedelta(seconds=15)) is not None
step: Final = StepTrace(
step_id="step-1",
run_id="run-1",
node_id="node-1",
state=RunState.RUNNING,
input=None,
attempts=1,
deadline=NOW + timedelta(seconds=10),
started_at=NOW,
)
assert not await repository.checkpoint("run-1", "stale-lease", step)
assert not await repository.finish("run-1", "stale-lease", RunState.COMPLETED, "forged", None)
assert database.events == ()
current: Final = await repository.get("run-1")
assert current is not None and current.status is RunState.RUNNING and current.output is None
ownership: Final = database.queries[1]
assert ownership.in_transaction and ownership.args == ("run-1", "stale-lease")
assert "lease_token = $2" in ownership.sql
assert "lease_expires_at > timezone('UTC', NOW()) FOR UPDATE" in ownership.sql
@pytest.mark.asyncio
async def test_attempt_receipts_aggregate_without_counting_duplicate_checkpoints_twice() -> None:
database: Final = ExecutionRows((pending_run(),))
repository: Final = WorkflowExecutionRepository(database)
assert await repository.claim("run-1", "lease", NOW + timedelta(seconds=15)) is not None
first: Final = StepTrace(
step_id="step-1",
run_id="run-1",
node_id="node-1",
state=RunState.FAILED,
input=None,
attempts=1,
deadline=NOW + timedelta(seconds=10),
started_at=NOW,
finished_at=NOW,
prompt_tokens=4,
completion_tokens=2,
response_cost=0.125,
request_id="failed-receipt",
call_id="call-1",
error="Provider failed after billing",
)
second: Final = first.model_copy(
update={
"state": RunState.COMPLETED,
"attempts": 2,
"output": {"done": True},
"error": None,
"prompt_tokens": 6,
"completion_tokens": 3,
"response_cost": 0.25,
"request_id": "successful-receipt",
"call_id": "call-2",
}
)
for checkpoint in (first, first, second, second):
assert await repository.checkpoint("run-1", "lease", checkpoint)
traces: Final = await repository.steps("run-1")
assert traces == (
second.model_copy(
update={
"prompt_tokens": 10,
"completion_tokens": 5,
"response_cost": 0.375,
"request_ids": ("failed-receipt", "successful-receipt"),
}
),
)
assert len(database.events) == 4
query: Final = database.queries[-1]
assert query.args == ("run-1",)
assert "DISTINCT ON (step_name, data->'trace'->>'attempts')" in query.sql
invalid = await repository.list_runs("workflow-1", "team", None, None, 2, "not-a-cursor")
assert isinstance(invalid, ExecutionError) and invalid.code == "invalid_cursor"

6
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -31198,6 +31198,12 @@ export interface components {
* @default false
*/
enable_public_model_hub: boolean;
/**
* Enable Workflows
* @description Enable LiteFlow workflow authoring and execution. Off by default; restart the proxy after changing this setting.
* @default false
*/
enable_workflows: boolean;
/**
* Enforce Fallback Model Access
* @description If True, router fallbacks configured in router_settings are only attempted when the calling key (and its team and project) is allowed to call the fallback model; unauthorized fallback targets are skipped and the primary model's error is returned. Default is True. Set False to disable fallback model access enforcement.