diff --git a/litellm/litellm_core_utils/prompt_templates/json_schema_utils.py b/litellm/litellm_core_utils/prompt_templates/json_schema_utils.py index 1d6d28c2680..73e0733abbe 100644 --- a/litellm/litellm_core_utils/prompt_templates/json_schema_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/json_schema_utils.py @@ -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)) diff --git a/litellm/models/workflow_definition.py b/litellm/models/workflow_definition.py index 8a0a7907a76..61cfaf2f23a 100644 --- a/litellm/models/workflow_definition.py +++ b/litellm/models/workflow_definition.py @@ -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: ... diff --git a/litellm/models/workflow_execution.py b/litellm/models/workflow_execution.py index 28cce6dfb1b..978e8232a8e 100644 --- a/litellm/models/workflow_execution.py +++ b/litellm/models/workflow_execution.py @@ -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: ... diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6cb6e61864d..c4ef7e3a518 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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", diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2f9a1300244..aec2e59ea78 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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( diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index 908d0fa8969..430f8cf693c 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -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[ diff --git a/litellm/proxy/management_endpoints/workflow_definition_endpoints.py b/litellm/proxy/management_endpoints/workflow_definition_endpoints.py index f134db68f64..faf51f857c7 100644 --- a/litellm/proxy/management_endpoints/workflow_definition_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_definition_endpoints.py @@ -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")] diff --git a/litellm/proxy/management_endpoints/workflow_execution_endpoints.py b/litellm/proxy/management_endpoints/workflow_execution_endpoints.py index a4ec49a58bf..b1d42b31087 100644 --- a/litellm/proxy/management_endpoints/workflow_execution_endpoints.py +++ b/litellm/proxy/management_endpoints/workflow_execution_endpoints.py @@ -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, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 13fa99aa117..d9362e61840 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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( diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 90acac978ed..eba81604bf4 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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( diff --git a/litellm/proxy/workflows/LITEFLOW.md b/litellm/proxy/workflows/LITEFLOW.md deleted file mode 100644 index a6b459161b8..00000000000 --- a/litellm/proxy/workflows/LITEFLOW.md +++ /dev/null @@ -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": }`, 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": ""}` | -| 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/"`, `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 diff --git a/litellm/proxy/workflows/definitions.py b/litellm/proxy/workflows/definitions.py index c956574963d..50b791b56ab 100644 --- a/litellm/proxy/workflows/definitions.py +++ b/litellm/proxy/workflows/definitions.py @@ -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: diff --git a/litellm/proxy/workflows/lifespan.py b/litellm/proxy/workflows/lifespan.py index 7db910be0bf..bd559b910c5 100644 --- a/litellm/proxy/workflows/lifespan.py +++ b/litellm/proxy/workflows/lifespan.py @@ -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) diff --git a/litellm/proxy/workflows/runner.py b/litellm/proxy/workflows/runner.py index 9d5782ed1db..0992d84869b 100644 --- a/litellm/proxy/workflows/runner.py +++ b/litellm/proxy/workflows/runner.py @@ -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( diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 22d525faad5..1004682fc46 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -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" diff --git a/litellm/repositories/workflow_definition_repository.py b/litellm/repositories/workflow_definition_repository.py index e4bbd9c884a..de603b245ac 100644 --- a/litellm/repositories/workflow_definition_repository.py +++ b/litellm/repositories/workflow_definition_repository.py @@ -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) diff --git a/litellm/repositories/workflow_execution_repository.py b/litellm/repositories/workflow_execution_repository.py index 257df3c1603..f0298053e35 100644 --- a/litellm/repositories/workflow_execution_repository.py +++ b/litellm/repositories/workflow_execution_repository.py @@ -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) diff --git a/tests/integration/database/test_workflow_execution.py b/tests/integration/database/test_workflow_execution.py index 0295a5659fc..bbe8b3fe8a2 100644 --- a/tests/integration/database/test_workflow_execution.py +++ b/tests/integration/database/test_workflow_execution.py @@ -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 diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index b4ea51a55f8..4cdd96522ca 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -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, diff --git a/tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py b/tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py index 87a0cad7450..be74705e275 100644 --- a/tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_workflow_definition_endpoints.py @@ -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()) diff --git a/tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py b/tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py index cb8459ede01..8b4d52f6786 100644 --- a/tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_workflow_execution_endpoints.py @@ -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"} diff --git a/tests/unit/proxy/proxy_server/test_routes_chat_completions.py b/tests/unit/proxy/proxy_server/test_routes_chat_completions.py index e124bde8120..840dbb9c073 100644 --- a/tests/unit/proxy/proxy_server/test_routes_chat_completions.py +++ b/tests/unit/proxy/proxy_server/test_routes_chat_completions.py @@ -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() diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index 72927fa0197..4e9601277f2 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -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"}]}], diff --git a/tests/unit/proxy/workflow_definition_rows.py b/tests/unit/proxy/workflow_definition_rows.py index d03e1d5bef7..cabea4f541f 100644 --- a/tests/unit/proxy/workflow_definition_rows.py +++ b/tests/unit/proxy/workflow_definition_rows.py @@ -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 diff --git a/tests/unit/proxy/workflows/test_definitions.py b/tests/unit/proxy/workflows/test_definitions.py index ac5ccdb89e4..ae3076421af 100644 --- a/tests/unit/proxy/workflows/test_definitions.py +++ b/tests/unit/proxy/workflows/test_definitions.py @@ -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" diff --git a/tests/unit/proxy/workflows/test_lifespan.py b/tests/unit/proxy/workflows/test_lifespan.py new file mode 100644 index 00000000000..2f4fab01760 --- /dev/null +++ b/tests/unit/proxy/workflows/test_lifespan.py @@ -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() diff --git a/tests/unit/proxy/workflows/test_runner.py b/tests/unit/proxy/workflows/test_runner.py index 12989251219..51f47d4c81b 100644 --- a/tests/unit/proxy/workflows/test_runner.py +++ b/tests/unit/proxy/workflows/test_runner.py @@ -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( diff --git a/tests/unit/repositories/test_workflow_definition_repository.py b/tests/unit/repositories/test_workflow_definition_repository.py index 79606e12d13..e8532576460 100644 --- a/tests/unit/repositories/test_workflow_definition_repository.py +++ b/tests/unit/repositories/test_workflow_definition_repository.py @@ -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",) diff --git a/tests/unit/repositories/test_workflow_execution_repository.py b/tests/unit/repositories/test_workflow_execution_repository.py index 4665913ebc3..6f5b7393457 100644 --- a/tests/unit/repositories/test_workflow_execution_repository.py +++ b/tests/unit/repositories/test_workflow_execution_repository.py @@ -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" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0555cc95b1f..013c7b135ac 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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.