mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/pensive-bartik-e24048
# Conflicts: # ui/litellm-dashboard/src/components/leftnav.tsx
This commit is contained in:
Yuneng Jiang
commit
cfda5e17ac
26 changed files with 3138 additions and 146 deletions
|
|
@ -0,0 +1,75 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowRun" (
|
||||
"run_id" TEXT NOT NULL,
|
||||
"session_id" TEXT NOT NULL,
|
||||
"workflow_type" TEXT NOT NULL,
|
||||
"status" TEXT NOT NULL DEFAULT 'pending',
|
||||
"created_by" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
"input" JSONB,
|
||||
"output" JSONB,
|
||||
"metadata" JSONB,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowRun_pkey" PRIMARY KEY ("run_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowEvent" (
|
||||
"event_id" TEXT NOT NULL,
|
||||
"run_id" TEXT NOT NULL,
|
||||
"event_type" TEXT NOT NULL,
|
||||
"step_name" TEXT NOT NULL,
|
||||
"sequence_number" INTEGER NOT NULL,
|
||||
"data" JSONB,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowEvent_pkey" PRIMARY KEY ("event_id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_WorkflowMessage" (
|
||||
"message_id" TEXT NOT NULL,
|
||||
"run_id" TEXT NOT NULL,
|
||||
"role" TEXT NOT NULL,
|
||||
"content" TEXT NOT NULL,
|
||||
"sequence_number" INTEGER NOT NULL,
|
||||
"session_id" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_WorkflowMessage_pkey" PRIMARY KEY ("message_id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowRun_session_id_key" ON "LiteLLM_WorkflowRun"("session_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_workflow_type_status_idx" ON "LiteLLM_WorkflowRun"("workflow_type", "status");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_session_id_idx" ON "LiteLLM_WorkflowRun"("session_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_created_at_idx" ON "LiteLLM_WorkflowRun"("created_at");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowRun_created_by_idx" ON "LiteLLM_WorkflowRun"("created_by");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowEvent_run_id_idx" ON "LiteLLM_WorkflowEvent"("run_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowEvent_run_id_sequence_number_key" ON "LiteLLM_WorkflowEvent"("run_id", "sequence_number");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_WorkflowMessage_run_id_idx" ON "LiteLLM_WorkflowMessage"("run_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX "LiteLLM_WorkflowMessage_run_id_sequence_number_key" ON "LiteLLM_WorkflowMessage"("run_id", "sequence_number");
|
||||
|
||||
-- AddForeignKey
|
||||
ALTER TABLE "LiteLLM_WorkflowEvent" ADD CONSTRAINT "LiteLLM_WorkflowEvent_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
|
||||
-- AddForeignKey
|
||||
ALTER TABLE "LiteLLM_WorkflowMessage" ADD CONSTRAINT "LiteLLM_WorkflowMessage_run_id_fkey" FOREIGN KEY ("run_id") REFERENCES "LiteLLM_WorkflowRun"("run_id") ON DELETE RESTRICT ON UPDATE CASCADE;
|
||||
|
||||
|
|
@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession {
|
|||
@@id([session_id, router_name, model_name])
|
||||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Workflow Run Tracking
|
||||
//
|
||||
// Generic durable state tracking for any agent or automated workflow.
|
||||
// Design: three tables — run (header + materialized status), event (append-only
|
||||
// source of truth for state transitions), message (conversation inbox/outbox).
|
||||
//
|
||||
// Usage:
|
||||
// - Set `workflow_type` to identify the owning system (e.g. "shin-builder").
|
||||
// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.).
|
||||
// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to
|
||||
// the proxy — all spend logs for this run are automatically tagged.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// One instance of work being done. `status` is a materialized cache of the
|
||||
// latest event; the event log is the authoritative source of truth.
|
||||
model LiteLLM_WorkflowRun {
|
||||
run_id String @id @default(uuid())
|
||||
session_id String @unique @default(uuid())
|
||||
workflow_type String
|
||||
status String @default("pending")
|
||||
created_by String? // user_id of the key that created this run; null = created by master key
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
input Json?
|
||||
output Json?
|
||||
metadata Json?
|
||||
|
||||
events LiteLLM_WorkflowEvent[]
|
||||
messages LiteLLM_WorkflowMessage[]
|
||||
|
||||
@@index([workflow_type, status])
|
||||
@@index([session_id])
|
||||
@@index([created_at])
|
||||
@@index([created_by])
|
||||
}
|
||||
|
||||
// Append-only log of state transitions. Never mutate rows here.
|
||||
// `step_name` and `event_type` are caller-defined strings — no hardcoded enums.
|
||||
// Status auto-update rules (applied by the append endpoint):
|
||||
// step.started → run.status = running
|
||||
// step.failed → run.status = failed
|
||||
// hook.waiting → run.status = paused
|
||||
// hook.received → run.status = running
|
||||
model LiteLLM_WorkflowEvent {
|
||||
event_id String @id @default(uuid())
|
||||
run_id String
|
||||
event_type String
|
||||
step_name String
|
||||
sequence_number Int
|
||||
data Json?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
|
||||
|
||||
@@unique([run_id, sequence_number])
|
||||
@@index([run_id])
|
||||
}
|
||||
|
||||
// Conversation inbox/outbox — full message content, separate from the durable
|
||||
// event log. Spend logs truncate messages; this table stores them in full.
|
||||
// `session_id` here is the Claude --resume session ID (or similar).
|
||||
model LiteLLM_WorkflowMessage {
|
||||
message_id String @id @default(uuid())
|
||||
run_id String
|
||||
role String
|
||||
content String
|
||||
sequence_number Int
|
||||
session_id String?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
|
||||
|
||||
@@unique([run_id, sequence_number])
|
||||
@@index([run_id])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -87,9 +87,7 @@ class PromptManagementBase(ABC):
|
|||
try:
|
||||
messages = compiled_prompt_client["prompt_template"] + client_messages
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}"
|
||||
)
|
||||
raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}")
|
||||
|
||||
compiled_prompt_client["completed_messages"] = messages
|
||||
return compiled_prompt_client
|
||||
|
|
@ -116,9 +114,7 @@ class PromptManagementBase(ABC):
|
|||
try:
|
||||
messages = compiled_prompt_client["prompt_template"] + client_messages
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}"
|
||||
)
|
||||
raise ValueError(f"Error compiling prompt: {e}. Prompt id={prompt_id}")
|
||||
|
||||
compiled_prompt_client["completed_messages"] = messages
|
||||
return compiled_prompt_client
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from typing import Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import litellm
|
||||
from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH
|
||||
|
|
@ -8,6 +9,43 @@ from litellm.secret_managers.main import get_secret, get_secret_str
|
|||
from ..types.router import LiteLLM_Params
|
||||
|
||||
|
||||
def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool:
|
||||
"""
|
||||
Match a registered openai-compatible endpoint against a caller-supplied
|
||||
``api_base`` using parsed-URL semantics, not unanchored substring search.
|
||||
|
||||
Both inputs may be a bare hostname (``api.perplexity.ai``), host+path
|
||||
(``api.deepinfra.com/v1/openai``), or a full URL
|
||||
(``https://api.cerebras.ai/v1``). Hostnames must match exactly
|
||||
(case-insensitive); if the registered endpoint has a non-trivial path,
|
||||
the api_base path must start with it on a segment boundary.
|
||||
|
||||
The naive ``endpoint in api_base`` shape lets a caller pass
|
||||
``https://attacker.com/api.groq.com/openai/v1`` to coerce the proxy
|
||||
into reading the server's GROQ_API_KEY from the environment and
|
||||
forwarding it to the attacker's host as a Bearer credential.
|
||||
"""
|
||||
|
||||
def _parse(value: str):
|
||||
# Ensure urlparse sees a scheme so it populates hostname / path.
|
||||
normalized = value if "://" in value else f"https://{value}"
|
||||
return urlparse(normalized)
|
||||
|
||||
parsed_endpoint = _parse(endpoint)
|
||||
parsed_url = _parse(api_base)
|
||||
|
||||
endpoint_host = (parsed_endpoint.hostname or "").lower()
|
||||
url_host = (parsed_url.hostname or "").lower()
|
||||
if not endpoint_host or endpoint_host != url_host:
|
||||
return False
|
||||
|
||||
endpoint_path = parsed_endpoint.path.rstrip("/")
|
||||
if not endpoint_path:
|
||||
return True
|
||||
url_path = parsed_url.path.rstrip("/")
|
||||
return url_path == endpoint_path or url_path.startswith(endpoint_path + "/")
|
||||
|
||||
|
||||
def _is_non_openai_azure_model(model: str) -> bool:
|
||||
try:
|
||||
model_name = model.split("/", 1)[1]
|
||||
|
|
@ -210,7 +248,7 @@ def get_llm_provider( # noqa: PLR0915
|
|||
# check if api base is a known openai compatible endpoint
|
||||
if api_base:
|
||||
for endpoint in litellm.openai_compatible_endpoints:
|
||||
if endpoint in api_base:
|
||||
if _endpoint_matches_api_base(endpoint, api_base):
|
||||
if endpoint == "api.perplexity.ai":
|
||||
custom_llm_provider = "perplexity"
|
||||
dynamic_api_key = get_secret_str("PERPLEXITYAI_API_KEY")
|
||||
|
|
|
|||
|
|
@ -824,8 +824,6 @@ def convert_to_model_response_object( # noqa: PLR0915
|
|||
stream=stream,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
hidden_params=hidden_params,
|
||||
_response_headers=_response_headers,
|
||||
convert_tool_call_to_json_mode=convert_tool_call_to_json_mode,
|
||||
)
|
||||
raise Exception(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from typing import Any, Coroutine, Dict, Optional, Union
|
|||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
|
|
@ -224,8 +225,14 @@ class VertexAIBatchPrediction(VertexLLM):
|
|||
},
|
||||
)
|
||||
|
||||
response = sync_handler.get(
|
||||
url=api_base,
|
||||
# ``api_base`` here can come from caller-supplied request kwargs
|
||||
# (clientside override). Wrap the fetch in ``safe_get`` so DNS
|
||||
# rebind / private / cloud-metadata targets are rejected; the
|
||||
# proxy auth gate already blocks malicious clientside ``api_base``
|
||||
# at the boundary — this is defense-in-depth for SDK callers.
|
||||
response = safe_get(
|
||||
sync_handler,
|
||||
api_base,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
|
@ -270,8 +277,13 @@ class VertexAIBatchPrediction(VertexLLM):
|
|||
},
|
||||
)
|
||||
|
||||
response = await client.get(
|
||||
url=api_base,
|
||||
# Mirror the sync path: ``api_base`` may come from caller-supplied
|
||||
# request kwargs, so wrap the fetch in ``async_safe_get`` to reject
|
||||
# DNS-rebind / private / cloud-metadata targets. Defense-in-depth
|
||||
# behind the proxy auth gate's clientside ``api_base`` check.
|
||||
response = await async_safe_get(
|
||||
client,
|
||||
api_base,
|
||||
headers=headers,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
|
|
|
|||
|
|
@ -6,9 +6,11 @@ from typing import Any, List, Optional, Tuple
|
|||
|
||||
from fastapi import HTTPException, Request, status
|
||||
|
||||
import litellm
|
||||
from litellm import Router, provider_list
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import STANDARD_CUSTOMER_ID_HEADERS
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
|
||||
from litellm.proxy._types import *
|
||||
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
|
||||
|
||||
|
|
@ -53,6 +55,12 @@ def _check_valid_ip(
|
|||
def check_complete_credentials(request_body: dict) -> bool:
|
||||
"""
|
||||
if 'api_base' in request body. Check if complete credentials given. Prevent malicious attacks.
|
||||
|
||||
Supplying an ``api_key`` is necessary but not sufficient: even with
|
||||
credentials supplied, an ``api_base`` / ``base_url`` that resolves to a
|
||||
private/internal/cloud-metadata address would still allow the proxy to
|
||||
be used as an SSRF pivot. Validate any URL fields here so the gate
|
||||
can't be bypassed with ``api_key=anything`` plus a malicious target.
|
||||
"""
|
||||
given_model: Optional[str] = None
|
||||
|
||||
|
|
@ -70,10 +78,27 @@ def check_complete_credentials(request_body: dict) -> bool:
|
|||
return False
|
||||
|
||||
api_key_value = request_body.get("api_key")
|
||||
if api_key_value and isinstance(api_key_value, str) and api_key_value.strip():
|
||||
return True
|
||||
if not (api_key_value and isinstance(api_key_value, str) and api_key_value.strip()):
|
||||
return False
|
||||
|
||||
return False
|
||||
# ``validate_url`` itself doesn't consult the toggle; ``safe_get`` /
|
||||
# ``async_safe_get`` do. Mirror that here so admins who explicitly
|
||||
# disabled URL validation (e.g. for an internal Ollama endpoint they
|
||||
# accept the SSRF risk for) aren't blocked at the proxy boundary.
|
||||
if getattr(litellm, "user_url_validation", False):
|
||||
for url_field in ("api_base", "base_url"):
|
||||
url_value = request_body.get(url_field)
|
||||
if not url_value or not isinstance(url_value, str):
|
||||
continue
|
||||
try:
|
||||
validate_url(url_value)
|
||||
except SSRFError as e:
|
||||
raise ValueError(
|
||||
f"Rejected request: client-side {url_field}={url_value!r} "
|
||||
f"is rejected by the SSRF guard ({e})."
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def check_regex_or_str_match(request_body_value: Any, regex_str: str) -> bool:
|
||||
|
|
@ -159,15 +184,42 @@ def is_request_body_safe(
|
|||
"aws_web_identity_token",
|
||||
"aws_role_name",
|
||||
"vertex_credentials",
|
||||
# Endpoint-targeting fields that retarget the outbound request or
|
||||
# an observability callback. An attacker-controlled value either
|
||||
# exfiltrates the request payload (incl. messages + admin-set
|
||||
# tokens) to the attacker's host, or coerces the proxy into
|
||||
# authenticating against the attacker's host with admin secrets.
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"langsmith_base_url",
|
||||
"langfuse_host",
|
||||
"posthog_host",
|
||||
"braintrust_host",
|
||||
"slack_webhook_url",
|
||||
# Provider-specific endpoint overrides that flow into the outbound
|
||||
# request via ``optional_params``. Same threat as ``api_base``:
|
||||
# ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
|
||||
# S3; ``sagemaker_base_url`` redirects all SageMaker traffic;
|
||||
# ``deployment_url`` redirects SAP deployments.
|
||||
"s3_endpoint_url",
|
||||
"sagemaker_base_url",
|
||||
"deployment_url",
|
||||
]
|
||||
|
||||
# The blocklist is enforced unconditionally. Legitimate clientside
|
||||
# credential / endpoint passthrough goes through one of the two
|
||||
# explicit admin opt-ins (``general_settings.allow_client_side_credentials``
|
||||
# proxy-wide or ``configurable_clientside_auth_params`` per deployment).
|
||||
# Historically there was a third, *implicit*, *caller-controlled* path:
|
||||
# ``check_complete_credentials`` returned True when the caller supplied
|
||||
# any non-empty ``api_key``, which made the entire blocklist a no-op.
|
||||
# That bypass turned every missing entry on the blocklist into an
|
||||
# exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp,
|
||||
# GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l,
|
||||
# b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now
|
||||
# has a single, predictable failure mode for missing entries (a 400),
|
||||
# not a credential leak.
|
||||
for param in banned_params:
|
||||
if (
|
||||
param in request_body
|
||||
and not check_complete_credentials( # allow client-credentials to be passed to proxy
|
||||
request_body=request_body
|
||||
)
|
||||
):
|
||||
if param in request_body:
|
||||
if general_settings.get("allow_client_side_credentials") is True:
|
||||
return True
|
||||
elif (
|
||||
|
|
@ -182,7 +234,10 @@ def is_request_body_safe(
|
|||
return True
|
||||
raise ValueError(
|
||||
f"Rejected Request: {param} is not allowed in request body. "
|
||||
"Enable with `general_settings::allow_client_side_credentials` on proxy config.yaml. "
|
||||
"Clientside passthrough requires explicit admin opt-in via "
|
||||
"either `general_settings.allow_client_side_credentials = true` "
|
||||
"(proxy-wide) or `configurable_clientside_auth_params` on the "
|
||||
"deployment in your proxy config.yaml. "
|
||||
"Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997",
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -52,18 +52,25 @@ class PrismaWrapper:
|
|||
engine = self._original_prisma._engine
|
||||
process = getattr(engine, "process", None) if engine is not None else None
|
||||
if process is not None:
|
||||
return process.pid
|
||||
pid = process.pid
|
||||
if isinstance(pid, int):
|
||||
return pid
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
async def _kill_engine_process(pid: int) -> None:
|
||||
"""Force-kill an orphaned engine subprocess to prevent DB connection pool leaks.
|
||||
"""Force-kill the engine subprocess to prevent DB connection pool leaks.
|
||||
|
||||
Called when disconnect() fails and the old engine process may still be
|
||||
holding open connections. Sends SIGTERM for graceful shutdown, waits
|
||||
briefly, then SIGKILL as a backstop.
|
||||
Called on every reconnect (in `recreate_prisma_client`) to retire the
|
||||
old query-engine subprocess without invoking prisma-client-py's
|
||||
synchronous `disconnect()` — which blocks the asyncio event loop on
|
||||
`subprocess.Popen.wait()` for 30-120+ seconds when the engine is
|
||||
stuck on TCP close.
|
||||
|
||||
Sends SIGTERM for graceful shutdown, waits briefly, then SIGKILL as
|
||||
a backstop.
|
||||
"""
|
||||
if pid <= 0:
|
||||
return
|
||||
|
|
@ -72,7 +79,7 @@ class PrismaWrapper:
|
|||
except (ProcessLookupError, PermissionError, OSError):
|
||||
return # Already dead or inaccessible
|
||||
verbose_proxy_logger.warning(
|
||||
"Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.",
|
||||
"Sent SIGTERM to prisma-query-engine PID %s during reconnect.",
|
||||
pid,
|
||||
)
|
||||
# Brief wait for graceful shutdown, then force-kill
|
||||
|
|
@ -217,15 +224,18 @@ class PrismaWrapper:
|
|||
async def recreate_prisma_client(
|
||||
self, new_db_url: str, http_client: Optional[Any] = None
|
||||
):
|
||||
"""Disconnect and reconnect the Prisma client with a new database URL."""
|
||||
"""Disconnect and reconnect the Prisma client with a new database URL.
|
||||
|
||||
Kills the old engine subprocess directly (SIGTERM → SIGKILL) rather than
|
||||
calling `disconnect()`. prisma-client-py's `disconnect()` calls a
|
||||
synchronous `subprocess.Popen.wait()` that can freeze the asyncio event
|
||||
loop for 30-120+ seconds when the engine is stuck on TCP close,
|
||||
breaking `/health/liveliness` and causing Kubernetes pod restarts.
|
||||
"""
|
||||
from prisma import Prisma # type: ignore
|
||||
|
||||
old_engine_pid = self._get_engine_pid()
|
||||
|
||||
try:
|
||||
await self._original_prisma.disconnect()
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}")
|
||||
if old_engine_pid > 0:
|
||||
await self._kill_engine_process(old_engine_pid)
|
||||
|
||||
if http_client is not None:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,492 @@
|
|||
"""
|
||||
WORKFLOW RUN MANAGEMENT
|
||||
|
||||
Generic durable state tracking for agents and automated workflows.
|
||||
|
||||
POST /v1/workflows/runs - Create a workflow run
|
||||
GET /v1/workflows/runs - List runs (filter by type, status)
|
||||
GET /v1/workflows/runs/{run_id} - Get run with latest event
|
||||
PATCH /v1/workflows/runs/{run_id} - Update status, metadata, output
|
||||
POST /v1/workflows/runs/{run_id}/events - Append event (updates run status)
|
||||
GET /v1/workflows/runs/{run_id}/events - Full event log
|
||||
POST /v1/workflows/runs/{run_id}/messages - Append conversation message
|
||||
GET /v1/workflows/runs/{run_id}/messages - Fetch conversation history
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, Literal, Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
||||
try:
|
||||
from prisma.errors import UniqueViolationError
|
||||
except ImportError:
|
||||
UniqueViolationError = None # type: ignore
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
_MAX_SEQUENCE_RETRIES = 5
|
||||
|
||||
|
||||
def _json(value: Any) -> str:
|
||||
"""Serialize a Python value for prisma-client-py Json fields (must be a string)."""
|
||||
return json.dumps(value)
|
||||
|
||||
|
||||
def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
|
||||
|
||||
def _caller_key(user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
|
||||
"""Return the hashed key token that identifies this caller, or None for master key."""
|
||||
return user_api_key_dict.token
|
||||
|
||||
|
||||
# Status transitions driven by event_type
|
||||
_EVENT_STATUS_MAP: Dict[str, str] = {
|
||||
"step.started": "running",
|
||||
"step.failed": "failed",
|
||||
"hook.waiting": "paused",
|
||||
"hook.received": "running",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Request / Response models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WorkflowRunCreateRequest(BaseModel):
|
||||
workflow_type: str
|
||||
input: Optional[Dict[str, Any]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
WorkflowRunStatus = Literal["pending", "running", "paused", "completed", "failed"]
|
||||
|
||||
|
||||
class WorkflowRunUpdateRequest(BaseModel):
|
||||
status: Optional[WorkflowRunStatus] = None
|
||||
output: Optional[Dict[str, Any]] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class WorkflowEventCreateRequest(BaseModel):
|
||||
event_type: str
|
||||
step_name: str
|
||||
data: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
class WorkflowMessageCreateRequest(BaseModel):
|
||||
role: str
|
||||
content: str
|
||||
session_id: Optional[str] = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _get_next_sequence_number(prisma_client: Any, run_id: str, table: str) -> int:
|
||||
"""Return MAX(sequence_number) + 1 for the given run, for either events or messages."""
|
||||
if table == "events":
|
||||
rows = await prisma_client.db.litellm_workflowevent.find_many(
|
||||
where={"run_id": run_id},
|
||||
order={"sequence_number": "desc"},
|
||||
take=1,
|
||||
)
|
||||
else:
|
||||
rows = await prisma_client.db.litellm_workflowmessage.find_many(
|
||||
where={"run_id": run_id},
|
||||
order={"sequence_number": "desc"},
|
||||
take=1,
|
||||
)
|
||||
return (rows[0].sequence_number + 1) if rows else 0
|
||||
|
||||
|
||||
async def _require_run(
|
||||
prisma_client: Any,
|
||||
run_id: str,
|
||||
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
|
||||
) -> Any:
|
||||
"""Return the run or raise 404. For non-admin callers, also enforce key ownership."""
|
||||
run = await prisma_client.db.litellm_workflowrun.find_unique(
|
||||
where={"run_id": run_id}
|
||||
)
|
||||
if run is None:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
if user_api_key_dict is not None and not _is_admin(user_api_key_dict):
|
||||
caller = _caller_key(user_api_key_dict)
|
||||
if not caller or run.created_by != caller:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
return run
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Endpoints
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/workflows/runs",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def create_workflow_run(
|
||||
data: WorkflowRunCreateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Create a new workflow run. Returns run_id and session_id.
|
||||
|
||||
The caller's API key token is stored as created_by so that non-admin keys
|
||||
can only see and modify their own runs.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
try:
|
||||
create_data: Dict[str, Any] = {
|
||||
"workflow_type": data.workflow_type,
|
||||
"created_by": _caller_key(user_api_key_dict),
|
||||
}
|
||||
if data.input is not None:
|
||||
create_data["input"] = _json(data.input)
|
||||
if data.metadata is not None:
|
||||
create_data["metadata"] = _json(data.metadata)
|
||||
run = await prisma_client.db.litellm_workflowrun.create(data=create_data)
|
||||
return run
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error creating workflow run: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/workflows/runs",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def list_workflow_runs(
|
||||
workflow_type: Optional[str] = Query(None),
|
||||
status: Optional[str] = Query(None),
|
||||
limit: int = Query(50, ge=1, le=250),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""List workflow runs. Filter by workflow_type and/or status.
|
||||
|
||||
Non-admin callers only see runs created by their own API key.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
where: Dict[str, Any] = {}
|
||||
if workflow_type:
|
||||
where["workflow_type"] = workflow_type
|
||||
if status:
|
||||
statuses = [s.strip() for s in status.split(",")]
|
||||
where["status"] = {"in": statuses} if len(statuses) > 1 else statuses[0]
|
||||
|
||||
# Non-admin callers are scoped to their own key.
|
||||
if not _is_admin(user_api_key_dict):
|
||||
caller = _caller_key(user_api_key_dict)
|
||||
if caller:
|
||||
where["created_by"] = caller
|
||||
|
||||
try:
|
||||
runs = await prisma_client.db.litellm_workflowrun.find_many(
|
||||
where=where,
|
||||
order={"created_at": "desc"},
|
||||
take=limit,
|
||||
)
|
||||
return {"runs": runs, "count": len(runs)}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error listing workflow runs: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/workflows/runs/{run_id}",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def get_workflow_run(
|
||||
run_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Get a workflow run with its most recent event."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
try:
|
||||
run = await prisma_client.db.litellm_workflowrun.find_unique(
|
||||
where={"run_id": run_id},
|
||||
include={"events": {"order_by": {"sequence_number": "desc"}, "take": 1}},
|
||||
)
|
||||
if run is None:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
if not _is_admin(user_api_key_dict):
|
||||
caller = _caller_key(user_api_key_dict)
|
||||
if not caller or run.created_by != caller:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
return run
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error getting workflow run: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.patch(
|
||||
"/v1/workflows/runs/{run_id}",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def update_workflow_run(
|
||||
run_id: str,
|
||||
data: WorkflowRunUpdateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Update status, metadata, or output on a workflow run."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
update: Dict[str, Any] = {}
|
||||
if data.status is not None:
|
||||
update["status"] = data.status
|
||||
if data.output is not None:
|
||||
update["output"] = _json(data.output)
|
||||
if data.metadata is not None:
|
||||
update["metadata"] = _json(data.metadata)
|
||||
|
||||
if not update:
|
||||
raise HTTPException(status_code=400, detail="No fields to update")
|
||||
|
||||
# Enforce ownership before writing.
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
|
||||
try:
|
||||
run = await prisma_client.db.litellm_workflowrun.update(
|
||||
where={"run_id": run_id},
|
||||
data=update,
|
||||
)
|
||||
if run is None:
|
||||
raise HTTPException(status_code=404, detail=f"Run '{run_id}' not found")
|
||||
return run
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error updating workflow run: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/workflows/runs/{run_id}/events",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def append_workflow_event(
|
||||
run_id: str,
|
||||
data: WorkflowEventCreateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Append an event to the run's event log. Also updates run.status if event_type maps to a status.
|
||||
|
||||
Sequence numbers use optimistic concurrency: on a unique-constraint collision
|
||||
(concurrent append), retries up to _MAX_SEQUENCE_RETRIES times with a fresh MAX+1.
|
||||
The event+status update is atomic in a single DB transaction.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
|
||||
new_status = _EVENT_STATUS_MAP.get(data.event_type)
|
||||
|
||||
for attempt in range(_MAX_SEQUENCE_RETRIES):
|
||||
try:
|
||||
seq = await _get_next_sequence_number(prisma_client, run_id, "events")
|
||||
event_data: Dict[str, Any] = {
|
||||
"run_id": run_id,
|
||||
"event_type": data.event_type,
|
||||
"step_name": data.step_name,
|
||||
"sequence_number": seq,
|
||||
}
|
||||
if data.data is not None:
|
||||
event_data["data"] = _json(data.data)
|
||||
|
||||
async with prisma_client.db.tx() as tx:
|
||||
event = await tx.litellm_workflowevent.create(data=event_data)
|
||||
if new_status:
|
||||
await tx.litellm_workflowrun.update(
|
||||
where={"run_id": run_id},
|
||||
data={"status": new_status},
|
||||
)
|
||||
|
||||
return event
|
||||
|
||||
except Exception as e:
|
||||
if UniqueViolationError is not None and isinstance(e, UniqueViolationError):
|
||||
if attempt == _MAX_SEQUENCE_RETRIES - 1:
|
||||
verbose_proxy_logger.exception(
|
||||
"Sequence number collision after %d retries for run %s",
|
||||
_MAX_SEQUENCE_RETRIES,
|
||||
run_id,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Concurrent write conflict — please retry",
|
||||
)
|
||||
continue
|
||||
verbose_proxy_logger.exception("Error appending workflow event: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
raise HTTPException(
|
||||
status_code=500, detail="Failed to append event"
|
||||
) # pragma: no cover
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/workflows/runs/{run_id}/events",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def list_workflow_events(
|
||||
run_id: str,
|
||||
limit: int = Query(100, ge=1, le=500),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Fetch event log for a run, ordered by sequence_number. Default limit 100, max 500."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
|
||||
try:
|
||||
events = await prisma_client.db.litellm_workflowevent.find_many(
|
||||
where={"run_id": run_id},
|
||||
order={"sequence_number": "asc"},
|
||||
take=limit,
|
||||
)
|
||||
return {"events": events, "count": len(events)}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error listing workflow events: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/workflows/runs/{run_id}/messages",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def append_workflow_message(
|
||||
run_id: str,
|
||||
data: WorkflowMessageCreateRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Append a conversation message. Stores full content (not truncated).
|
||||
|
||||
Uses optimistic concurrency for sequence numbers.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
|
||||
for attempt in range(_MAX_SEQUENCE_RETRIES):
|
||||
try:
|
||||
seq = await _get_next_sequence_number(prisma_client, run_id, "messages")
|
||||
msg_data: Dict[str, Any] = {
|
||||
"run_id": run_id,
|
||||
"role": data.role,
|
||||
"content": data.content,
|
||||
"sequence_number": seq,
|
||||
}
|
||||
if data.session_id is not None:
|
||||
msg_data["session_id"] = data.session_id
|
||||
msg = await prisma_client.db.litellm_workflowmessage.create(data=msg_data)
|
||||
return msg
|
||||
|
||||
except Exception as e:
|
||||
if UniqueViolationError is not None and isinstance(e, UniqueViolationError):
|
||||
if attempt == _MAX_SEQUENCE_RETRIES - 1:
|
||||
verbose_proxy_logger.exception(
|
||||
"Sequence number collision after %d retries for run %s",
|
||||
_MAX_SEQUENCE_RETRIES,
|
||||
run_id,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Concurrent write conflict — please retry",
|
||||
)
|
||||
continue
|
||||
verbose_proxy_logger.exception("Error appending workflow message: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
raise HTTPException(
|
||||
status_code=500, detail="Failed to append message"
|
||||
) # pragma: no cover
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/workflows/runs/{run_id}/messages",
|
||||
tags=["workflow management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def list_workflow_messages(
|
||||
run_id: str,
|
||||
limit: int = Query(100, ge=1, le=500),
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Fetch conversation history for a run, ordered by sequence_number. Default limit 100, max 500."""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(
|
||||
status_code=500, detail=CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
|
||||
await _require_run(prisma_client, run_id, user_api_key_dict)
|
||||
|
||||
try:
|
||||
messages = await prisma_client.db.litellm_workflowmessage.find_many(
|
||||
where={"run_id": run_id},
|
||||
order={"sequence_number": "asc"},
|
||||
take=limit,
|
||||
)
|
||||
return {"messages": messages, "count": len(messages)}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error listing workflow messages: %s", e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
|
@ -427,6 +427,9 @@ from litellm.proxy.management_endpoints.team_endpoints import (
|
|||
from litellm.proxy.management_endpoints.tool_management_endpoints import (
|
||||
router as tool_management_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.workflow_management_endpoints import (
|
||||
router as workflow_management_router,
|
||||
)
|
||||
from litellm.proxy.memory.memory_endpoints import router as memory_router
|
||||
from litellm.proxy.management_endpoints.ui_sso import (
|
||||
get_disabled_non_admin_personal_key_creation,
|
||||
|
|
@ -14285,6 +14288,7 @@ app.include_router(model_management_router)
|
|||
app.include_router(model_access_group_management_router)
|
||||
app.include_router(tag_management_router)
|
||||
app.include_router(tool_management_router)
|
||||
app.include_router(workflow_management_router)
|
||||
app.include_router(memory_router)
|
||||
app.include_router(cost_tracking_settings_router)
|
||||
app.include_router(router_settings_router)
|
||||
|
|
|
|||
|
|
@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession {
|
|||
@@id([session_id, router_name, model_name])
|
||||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Workflow Run Tracking
|
||||
//
|
||||
// Generic durable state tracking for any agent or automated workflow.
|
||||
// Design: three tables — run (header + materialized status), event (append-only
|
||||
// source of truth for state transitions), message (conversation inbox/outbox).
|
||||
//
|
||||
// Usage:
|
||||
// - Set `workflow_type` to identify the owning system (e.g. "shin-builder").
|
||||
// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.).
|
||||
// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to
|
||||
// the proxy — all spend logs for this run are automatically tagged.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// One instance of work being done. `status` is a materialized cache of the
|
||||
// latest event; the event log is the authoritative source of truth.
|
||||
model LiteLLM_WorkflowRun {
|
||||
run_id String @id @default(uuid())
|
||||
session_id String @unique @default(uuid())
|
||||
workflow_type String
|
||||
status String @default("pending")
|
||||
created_by String? // user_id of the key that created this run; null = created by master key
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
input Json?
|
||||
output Json?
|
||||
metadata Json?
|
||||
|
||||
events LiteLLM_WorkflowEvent[]
|
||||
messages LiteLLM_WorkflowMessage[]
|
||||
|
||||
@@index([workflow_type, status])
|
||||
@@index([session_id])
|
||||
@@index([created_at])
|
||||
@@index([created_by])
|
||||
}
|
||||
|
||||
// Append-only log of state transitions. Never mutate rows here.
|
||||
// `step_name` and `event_type` are caller-defined strings — no hardcoded enums.
|
||||
// Status auto-update rules (applied by the append endpoint):
|
||||
// step.started → run.status = running
|
||||
// step.failed → run.status = failed
|
||||
// hook.waiting → run.status = paused
|
||||
// hook.received → run.status = running
|
||||
model LiteLLM_WorkflowEvent {
|
||||
event_id String @id @default(uuid())
|
||||
run_id String
|
||||
event_type String
|
||||
step_name String
|
||||
sequence_number Int
|
||||
data Json?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
|
||||
|
||||
@@unique([run_id, sequence_number])
|
||||
@@index([run_id])
|
||||
}
|
||||
|
||||
// Conversation inbox/outbox — full message content, separate from the durable
|
||||
// event log. Spend logs truncate messages; this table stores them in full.
|
||||
// `session_id` here is the Claude --resume session ID (or similar).
|
||||
model LiteLLM_WorkflowMessage {
|
||||
message_id String @id @default(uuid())
|
||||
run_id String
|
||||
role String
|
||||
content String
|
||||
sequence_number Int
|
||||
session_id String?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
|
||||
|
||||
@@unique([run_id, sequence_number])
|
||||
@@index([run_id])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4198,8 +4198,11 @@ class PrismaClient:
|
|||
|
||||
Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll
|
||||
handlers) to choose between heavy reconnect (engine dead -- recreate
|
||||
Prisma client, re-arm watcher) and lightweight reconnect (network
|
||||
blip -- disconnect, connect, SELECT 1).
|
||||
Prisma client, re-arm watcher) and direct reconnect (network blip --
|
||||
recreate Prisma client, re-arm watcher, SELECT 1). Both paths recreate
|
||||
the client via the non-blocking kill-then-construct flow rather than
|
||||
calling disconnect(), which blocks the event loop on the synchronous
|
||||
subprocess.Popen.wait() inside prisma-client-py (see issue #26191).
|
||||
"""
|
||||
effective_timeout = (
|
||||
timeout_seconds
|
||||
|
|
@ -4243,17 +4246,20 @@ class PrismaClient:
|
|||
)
|
||||
|
||||
async def _do_direct_reconnect() -> None:
|
||||
old_pid = self._get_engine_pid()
|
||||
try:
|
||||
await self.db.disconnect()
|
||||
except Exception as disconnect_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Prisma DB disconnect before reconnect failed: %s",
|
||||
disconnect_err,
|
||||
db_url = os.getenv("DATABASE_URL", "")
|
||||
if not db_url:
|
||||
verbose_proxy_logger.error(
|
||||
"DATABASE_URL not set; cannot reconnect Prisma client."
|
||||
)
|
||||
await PrismaWrapper._kill_engine_process(old_pid)
|
||||
|
||||
await self.db.connect()
|
||||
raise RuntimeError("DATABASE_URL not set")
|
||||
# Fresh Prisma client + new engine subprocess. The previous
|
||||
# "lightweight" path called `disconnect()` which blocks the
|
||||
# event loop on `subprocess.Popen.wait()`; since that call
|
||||
# ends up killing the engine anyway, we do it non-blockingly
|
||||
# via `_kill_engine_process` inside `recreate_prisma_client`.
|
||||
self._cleanup_engine_watcher()
|
||||
await self.db.recreate_prisma_client(db_url)
|
||||
await self._start_engine_watcher()
|
||||
await self.db.query_raw("SELECT 1")
|
||||
|
||||
await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout)
|
||||
|
|
|
|||
150
litellm/proxy/workflows/README.md
Normal file
150
litellm/proxy/workflows/README.md
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
# Workflow Run Tracking
|
||||
|
||||
Generic durable state tracking for agents and automated workflows built on the LiteLLM proxy.
|
||||
|
||||
## The Problem
|
||||
|
||||
Agents like [shin-builder](https://github.com/BerriAI/shin-builder) run multi-stage pipelines (triage → plan → implement → PR). Their task state and conversation history lived in memory — a process restart lost everything.
|
||||
|
||||
## Three-Table Design
|
||||
|
||||
```
|
||||
WorkflowRun one instance of work (header + materialized status)
|
||||
WorkflowEvent append-only state transitions (source of truth for replay)
|
||||
WorkflowMessage conversation inbox/outbox (full content, not truncated)
|
||||
```
|
||||
|
||||
**WorkflowEvent is the source of truth.** `WorkflowRun.status` is a materialized cache updated automatically when events are appended. If you need to debug a run, replay its events.
|
||||
|
||||
## API
|
||||
|
||||
All endpoints require a valid LiteLLM API key (`Authorization: Bearer sk-...`).
|
||||
|
||||
### Runs
|
||||
|
||||
```
|
||||
POST /v1/workflows/runs Create a run
|
||||
GET /v1/workflows/runs List runs (?workflow_type=&status=)
|
||||
GET /v1/workflows/runs/{run_id} Get run + latest event
|
||||
PATCH /v1/workflows/runs/{run_id} Update status / metadata / output
|
||||
```
|
||||
|
||||
### Events
|
||||
|
||||
```
|
||||
POST /v1/workflows/runs/{run_id}/events Append event (auto-updates run status)
|
||||
GET /v1/workflows/runs/{run_id}/events Full event log (ordered by sequence)
|
||||
```
|
||||
|
||||
### Messages
|
||||
|
||||
```
|
||||
POST /v1/workflows/runs/{run_id}/messages Append message
|
||||
GET /v1/workflows/runs/{run_id}/messages Conversation history (ordered by sequence)
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Create a run
|
||||
curl -X POST http://localhost:4000/v1/workflows/runs \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"workflow_type": "shin-builder", "metadata": {"title": "Fix login bug"}}'
|
||||
|
||||
# {"run_id": "abc-123", "session_id": "xyz-456", "status": "pending", ...}
|
||||
|
||||
# Mark step started (sets status → running)
|
||||
curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/events \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"event_type": "step.started", "step_name": "grill", "data": {"claude_session_id": "sess-789"}}'
|
||||
|
||||
# Store a conversation message
|
||||
curl -X POST http://localhost:4000/v1/workflows/runs/abc-123/messages \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"role": "user", "content": "What is the expected behavior?", "session_id": "sess-789"}'
|
||||
|
||||
# Restart recovery: fetch active runs and resume from last event's data.claude_session_id
|
||||
curl "http://localhost:4000/v1/workflows/runs?status=running,paused&workflow_type=shin-builder" \
|
||||
-H "Authorization: Bearer sk-1234"
|
||||
```
|
||||
|
||||
## Status Auto-Update Rules
|
||||
|
||||
When you append an event, the run's status is updated automatically:
|
||||
|
||||
| event_type | run.status |
|
||||
|-----------------|------------|
|
||||
| `step.started` | `running` |
|
||||
| `step.failed` | `failed` |
|
||||
| `hook.waiting` | `paused` |
|
||||
| `hook.received` | `running` |
|
||||
|
||||
Set `status = completed` explicitly via PATCH when the workflow finishes.
|
||||
|
||||
## Linking to Spend Logs
|
||||
|
||||
`WorkflowRun.session_id` is generated automatically (UUID). Pass it as the `x-litellm-session-id` header when making completions through the proxy:
|
||||
|
||||
```python
|
||||
headers = {"x-litellm-session-id": run.session_id}
|
||||
```
|
||||
|
||||
All spend log entries for this run are then tagged automatically. Query cost per run:
|
||||
|
||||
```
|
||||
POST /ui/spend_logs/view_session_spend_logs?session_id={run.session_id}
|
||||
```
|
||||
|
||||
## Sequence Numbers
|
||||
|
||||
Sequence numbers on events and messages are assigned server-side (`MAX + 1` per run). Callers never supply them. This guarantees ordering even under concurrent writes.
|
||||
|
||||
## Using from shin-builder
|
||||
|
||||
Replace the in-memory `tasks.py` dict with calls to these endpoints:
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
class WorkflowRunClient:
|
||||
def __init__(self, base_url: str, api_key: str):
|
||||
self._client = httpx.AsyncClient(
|
||||
base_url=base_url,
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
)
|
||||
|
||||
async def create_task(self, title: str, **metadata) -> dict:
|
||||
r = await self._client.post("/v1/workflows/runs", json={
|
||||
"workflow_type": "shin-builder",
|
||||
"metadata": {"title": title, **metadata},
|
||||
})
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
async def list_active_tasks(self) -> list:
|
||||
r = await self._client.get(
|
||||
"/v1/workflows/runs",
|
||||
params={"workflow_type": "shin-builder", "status": "running,paused"},
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()["runs"]
|
||||
|
||||
async def transition(self, run_id: str, step_name: str, event_type: str, data: dict = None):
|
||||
r = await self._client.post(f"/v1/workflows/runs/{run_id}/events", json={
|
||||
"event_type": event_type,
|
||||
"step_name": step_name,
|
||||
"data": data or {},
|
||||
})
|
||||
r.raise_for_status()
|
||||
|
||||
async def append_message(self, run_id: str, role: str, content: str, session_id: str = None):
|
||||
r = await self._client.post(f"/v1/workflows/runs/{run_id}/messages", json={
|
||||
"role": role, "content": content, "session_id": session_id,
|
||||
})
|
||||
r.raise_for_status()
|
||||
```
|
||||
|
||||
On startup, call `list_active_tasks()` to restore in-flight runs. The last `step.started` event's `data.claude_session_id` gives you the `--resume` ID.
|
||||
|
|
@ -11,9 +11,60 @@ If given, generate a unique model_id for the deployment.
|
|||
Ensures cooldowns are applied correctly.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
clientside_credential_keys = ["api_key", "api_base", "base_url"]
|
||||
|
||||
|
||||
def _admin_config_fields_to_clear_on_base_override() -> List[str]:
|
||||
"""
|
||||
Provider-specific credential / endpoint-targeting fields that must NOT
|
||||
flow through to a client-redirected upstream.
|
||||
|
||||
Built dynamically from ``CredentialLiteLLMParams.model_fields`` so any
|
||||
new provider field added there (Bedrock endpoint, Watsonx region, etc.)
|
||||
is gated automatically — plus a fixed list of kwargs-only fields that
|
||||
aren't declared on the typed model.
|
||||
"""
|
||||
from litellm.types.router import CredentialLiteLLMParams
|
||||
|
||||
typed_fields = [
|
||||
f
|
||||
for f in CredentialLiteLLMParams.model_fields
|
||||
if f not in clientside_credential_keys
|
||||
]
|
||||
kwargs_only_fields = [
|
||||
# Caller-supplied via **kwargs, not declared on CredentialLiteLLMParams.
|
||||
"organization",
|
||||
"extra_body",
|
||||
"extra_headers",
|
||||
"default_headers",
|
||||
"api_type",
|
||||
"azure_ad_token",
|
||||
"azure_ad_token_provider",
|
||||
"aws_session_token",
|
||||
"aws_sts_endpoint",
|
||||
"aws_web_identity_token",
|
||||
"aws_role_name",
|
||||
# OCI provider — consumed by litellm/llms/oci/* via optional_params
|
||||
# and not declared on CredentialLiteLLMParams. Without these here,
|
||||
# an admin's OCI signing key / tenancy / fingerprint would flow
|
||||
# through to an attacker-redirected upstream.
|
||||
"oci_signer",
|
||||
"oci_user",
|
||||
"oci_fingerprint",
|
||||
"oci_tenancy",
|
||||
"oci_key",
|
||||
"oci_key_file",
|
||||
]
|
||||
return typed_fields + kwargs_only_fields
|
||||
|
||||
|
||||
_ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE = (
|
||||
_admin_config_fields_to_clear_on_base_override()
|
||||
)
|
||||
|
||||
|
||||
def is_clientside_credential(request_kwargs: dict) -> bool:
|
||||
"""
|
||||
Check if the credential is a clientside credential.
|
||||
|
|
@ -34,4 +85,20 @@ def get_dynamic_litellm_params(litellm_params: dict, request_kwargs: dict) -> di
|
|||
for key in clientside_credential_keys:
|
||||
if key in request_kwargs:
|
||||
litellm_params[key] = request_kwargs[key]
|
||||
|
||||
# If the caller redirected api_base/base_url to a client-controlled value,
|
||||
# don't forward the admin's organization / extra_body / region / token /
|
||||
# vertex / aws fields — those were meant for the original upstream.
|
||||
# Always drop the admin's value first, then write the caller's value back
|
||||
# if they resupplied the field. The naive
|
||||
# ``if field not in request_kwargs: pop`` shape lets a caller *echo* a
|
||||
# field name (with any value, including an empty string) to keep the
|
||||
# admin's value in ``litellm_params`` and have it forwarded to the
|
||||
# redirected upstream.
|
||||
if "api_base" in request_kwargs or "base_url" in request_kwargs:
|
||||
for field in _ADMIN_CONFIG_FIELDS_TO_CLEAR_ON_BASE_OVERRIDE:
|
||||
litellm_params.pop(field, None)
|
||||
if field in request_kwargs:
|
||||
litellm_params[field] = request_kwargs[field]
|
||||
|
||||
return litellm_params
|
||||
|
|
|
|||
|
|
@ -1290,3 +1290,80 @@ model LiteLLM_AdaptiveRouterSession {
|
|||
@@id([session_id, router_name, model_name])
|
||||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Workflow Run Tracking
|
||||
//
|
||||
// Generic durable state tracking for any agent or automated workflow.
|
||||
// Design: three tables — run (header + materialized status), event (append-only
|
||||
// source of truth for state transitions), message (conversation inbox/outbox).
|
||||
//
|
||||
// Usage:
|
||||
// - Set `workflow_type` to identify the owning system (e.g. "shin-builder").
|
||||
// - Store domain-specific fields in `metadata` (worktree_path, pr_url, etc.).
|
||||
// - `session_id` on WorkflowRun matches `x-litellm-session-id` header sent to
|
||||
// the proxy — all spend logs for this run are automatically tagged.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// One instance of work being done. `status` is a materialized cache of the
|
||||
// latest event; the event log is the authoritative source of truth.
|
||||
model LiteLLM_WorkflowRun {
|
||||
run_id String @id @default(uuid())
|
||||
session_id String @unique @default(uuid())
|
||||
workflow_type String
|
||||
status String @default("pending")
|
||||
created_by String? // user_id of the key that created this run; null = created by master key
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
input Json?
|
||||
output Json?
|
||||
metadata Json?
|
||||
|
||||
events LiteLLM_WorkflowEvent[]
|
||||
messages LiteLLM_WorkflowMessage[]
|
||||
|
||||
@@index([workflow_type, status])
|
||||
@@index([session_id])
|
||||
@@index([created_at])
|
||||
@@index([created_by])
|
||||
}
|
||||
|
||||
// Append-only log of state transitions. Never mutate rows here.
|
||||
// `step_name` and `event_type` are caller-defined strings — no hardcoded enums.
|
||||
// Status auto-update rules (applied by the append endpoint):
|
||||
// step.started → run.status = running
|
||||
// step.failed → run.status = failed
|
||||
// hook.waiting → run.status = paused
|
||||
// hook.received → run.status = running
|
||||
model LiteLLM_WorkflowEvent {
|
||||
event_id String @id @default(uuid())
|
||||
run_id String
|
||||
event_type String
|
||||
step_name String
|
||||
sequence_number Int
|
||||
data Json?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
|
||||
|
||||
@@unique([run_id, sequence_number])
|
||||
@@index([run_id])
|
||||
}
|
||||
|
||||
// Conversation inbox/outbox — full message content, separate from the durable
|
||||
// event log. Spend logs truncate messages; this table stores them in full.
|
||||
// `session_id` here is the Claude --resume session ID (or similar).
|
||||
model LiteLLM_WorkflowMessage {
|
||||
message_id String @id @default(uuid())
|
||||
run_id String
|
||||
role String
|
||||
content String
|
||||
sequence_number Int
|
||||
session_id String?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
run LiteLLM_WorkflowRun @relation(fields: [run_id], references: [run_id])
|
||||
|
||||
@@unique([run_id, sequence_number])
|
||||
@@index([run_id])
|
||||
}
|
||||
|
|
|
|||
|
|
@ -557,6 +557,7 @@ async def test_avertex_batch_prediction(monkeypatch):
|
|||
mock_get_response = MagicMock()
|
||||
mock_get_response.json.return_value = mock_vertex_batch_response
|
||||
mock_get_response.status_code = 200
|
||||
mock_get_response.is_redirect = False
|
||||
mock_get_response.raise_for_status.return_value = None
|
||||
mock_get.return_value = mock_get_response
|
||||
|
||||
|
|
|
|||
|
|
@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive(
|
||||
async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
|
||||
engine_client,
|
||||
) -> None:
|
||||
"""_run_reconnect_cycle uses disconnect/connect when engine is alive."""
|
||||
engine_client._engine_pid = 1234
|
||||
"""Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1.
|
||||
|
||||
with patch.object(engine_client, "_is_engine_alive", return_value=True):
|
||||
The old "lightweight" path called `disconnect()` + `connect()`, which
|
||||
blocks the event loop on the sync `process.wait()` inside aclose().
|
||||
The fix routes both engine-alive and engine-dead paths through
|
||||
`recreate_prisma_client`, which non-blockingly kills the old engine.
|
||||
"""
|
||||
engine_client._engine_pid = 1234
|
||||
engine_client._start_engine_watcher = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(engine_client, "_is_engine_alive", return_value=True),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
engine_client.db.connect.assert_awaited_once()
|
||||
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
|
||||
"postgresql://test"
|
||||
)
|
||||
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
engine_client.db.recreate_prisma_client.assert_not_awaited()
|
||||
engine_client.db.disconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown(
|
||||
async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown(
|
||||
engine_client,
|
||||
) -> None:
|
||||
"""_run_reconnect_cycle uses lightweight path when engine PID is not tracked."""
|
||||
"""When the engine PID is not tracked, direct reconnect still runs."""
|
||||
engine_client._engine_pid = 0
|
||||
engine_client._start_engine_watcher = AsyncMock()
|
||||
|
||||
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
engine_client.db.connect.assert_awaited_once()
|
||||
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
|
||||
"postgresql://test"
|
||||
)
|
||||
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
engine_client.db.recreate_prisma_client.assert_not_awaited()
|
||||
engine_client.db.disconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_escalation_after_consecutive_lightweight_failures(engine_client):
|
||||
"""After N consecutive lightweight reconnect failures, _engine_confirmed_dead
|
||||
async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client):
|
||||
"""After N consecutive direct reconnect failures, _engine_confirmed_dead
|
||||
is set to True so _run_reconnect_cycle takes the heavy reconnect path."""
|
||||
engine_client._reconnect_escalation_threshold = 3
|
||||
engine_client._consecutive_reconnect_failures = 0
|
||||
engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test
|
||||
engine_client._start_engine_watcher = AsyncMock(return_value=None)
|
||||
|
||||
# Make lightweight reconnect fail every time
|
||||
engine_client.db.disconnect = AsyncMock(return_value=None)
|
||||
engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed"))
|
||||
# Make direct reconnect fail every time
|
||||
engine_client.db.recreate_prisma_client = AsyncMock(
|
||||
side_effect=Exception("recreate failed")
|
||||
)
|
||||
|
||||
# Run 3 failed reconnect attempts
|
||||
for i in range(3):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
assert result is False
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
for _ in range(3):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
assert result is False
|
||||
|
||||
assert engine_client._consecutive_reconnect_failures == 3
|
||||
|
||||
# Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle
|
||||
# Next attempt should escalate to the heavy path (recreate_prisma_client still
|
||||
# the call, but via the _engine_confirmed_dead branch that also re-arms the watcher).
|
||||
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
engine_client._start_engine_watcher = AsyncMock(return_value=None)
|
||||
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test_escalation", timeout_seconds=5.0
|
||||
)
|
||||
|
||||
# Heavy reconnect should have been attempted (recreate_prisma_client called)
|
||||
engine_client.db.recreate_prisma_client.assert_awaited_once()
|
||||
|
||||
|
||||
|
|
@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client):
|
|||
"""A successful reconnect resets _consecutive_reconnect_failures to 0."""
|
||||
engine_client._consecutive_reconnect_failures = 2
|
||||
engine_client._db_reconnect_cooldown_seconds = 0
|
||||
engine_client._start_engine_watcher = AsyncMock()
|
||||
|
||||
# Make reconnect succeed
|
||||
engine_client.db.disconnect = AsyncMock(return_value=None)
|
||||
engine_client.db.connect = AsyncMock(return_value=None)
|
||||
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
result = await engine_client._attempt_reconnect_inside_lock(
|
||||
force=True, reason="test", timeout_seconds=5.0
|
||||
)
|
||||
|
||||
assert result is True
|
||||
assert engine_client._consecutive_reconnect_failures == 0
|
||||
|
|
|
|||
|
|
@ -0,0 +1,138 @@
|
|||
"""
|
||||
Regression tests for the parsed-URL hostname match used to identify a
|
||||
caller-supplied ``api_base`` as a known openai-compatible provider.
|
||||
|
||||
The previous shape (``if endpoint in api_base:``) used unanchored
|
||||
substring search, which let a caller pass
|
||||
``https://attacker.com/api.groq.com/openai/v1`` and have the proxy
|
||||
return ``GROQ_API_KEY`` as the dynamic credential — exfiltrating the
|
||||
server's real provider key to an attacker-controlled host on the
|
||||
outbound request.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
_endpoint_matches_api_base,
|
||||
get_llm_provider,
|
||||
)
|
||||
|
||||
|
||||
class TestEndpointMatchesApiBase:
|
||||
"""Direct unit tests on the parsed-URL matcher."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint, api_base",
|
||||
[
|
||||
# Bare hostname endpoint, exact host match.
|
||||
("api.perplexity.ai", "https://api.perplexity.ai/v1"),
|
||||
# Endpoint includes a path; api_base path starts with it.
|
||||
("api.groq.com/openai/v1", "https://api.groq.com/openai/v1"),
|
||||
# Endpoint with full URL scheme.
|
||||
("https://api.cerebras.ai/v1", "https://api.cerebras.ai/v1/chat"),
|
||||
# Trailing-slash on registered endpoint must not break match.
|
||||
("https://llm.chutes.ai/v1/", "https://llm.chutes.ai/v1/chat"),
|
||||
# Case-insensitive on hostname.
|
||||
("api.groq.com/openai/v1", "https://API.GROQ.COM/openai/v1"),
|
||||
],
|
||||
)
|
||||
def test_legitimate_provider_urls_match(self, endpoint, api_base):
|
||||
assert _endpoint_matches_api_base(endpoint, api_base) is True
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"endpoint, api_base",
|
||||
[
|
||||
# Attacker host, registered endpoint smuggled into path.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://attacker.com/api.groq.com/openai/v1",
|
||||
),
|
||||
# Attacker host, registered endpoint smuggled into a path segment.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://attacker.com/foo/api.groq.com/openai/v1",
|
||||
),
|
||||
# Lookalike host that contains the registered host as a suffix label.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://api.groq.com.attacker.com/openai/v1",
|
||||
),
|
||||
# Lookalike host with the registered host as a prefix.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://api.groq.com.evil.example/openai/v1",
|
||||
),
|
||||
# Right host, wrong path — endpoint requires ``/openai/v1`` prefix.
|
||||
("api.groq.com/openai/v1", "https://api.groq.com/v1"),
|
||||
# Path-segment lookalike: ``/openai/v10`` must not match ``/openai/v1``.
|
||||
("api.groq.com/openai/v1", "https://api.groq.com/openai/v10"),
|
||||
# Userinfo / @-injection trick — the ``hostname`` after ``@`` is
|
||||
# what httpx connects to.
|
||||
(
|
||||
"api.groq.com/openai/v1",
|
||||
"https://api.groq.com@attacker.com/openai/v1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_attacker_smuggling_does_not_match(self, endpoint, api_base):
|
||||
assert _endpoint_matches_api_base(endpoint, api_base) is False
|
||||
|
||||
|
||||
class TestGetLlmProviderRejectsAttackerSmuggledApiBase:
|
||||
"""
|
||||
End-to-end: ``get_llm_provider`` must NOT return the server's stored
|
||||
secret (e.g. ``GROQ_API_KEY``) for an api_base whose hostname is
|
||||
attacker-controlled, even when the registered endpoint string appears
|
||||
elsewhere in the URL.
|
||||
"""
|
||||
|
||||
def test_attacker_host_does_not_yield_groq_secret(self):
|
||||
# The function may either fall through (different provider) or
|
||||
# raise BadRequestError because the model can't be identified.
|
||||
# The invariant under test is that ``GROQ_API_KEY`` is never
|
||||
# looked up against an attacker-controlled hostname.
|
||||
import litellm
|
||||
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
|
||||
return_value="server-real-groq-key",
|
||||
) as mocked_secret:
|
||||
try:
|
||||
_, _, dynamic_api_key, _ = get_llm_provider(
|
||||
model="some-model",
|
||||
api_base="https://attacker.com/api.groq.com/openai/v1",
|
||||
)
|
||||
# If it returned, the dynamic key must not be the secret.
|
||||
assert dynamic_api_key != "server-real-groq-key"
|
||||
except litellm.exceptions.BadRequestError:
|
||||
# Acceptable outcome: provider unidentifiable, no secret
|
||||
# was returned.
|
||||
pass
|
||||
|
||||
# Regardless of return / raise, the secret must never have been
|
||||
# read against this attacker-controlled api_base.
|
||||
groq_lookups = [
|
||||
call
|
||||
for call in mocked_secret.call_args_list
|
||||
if call.args and call.args[0] == "GROQ_API_KEY"
|
||||
]
|
||||
assert groq_lookups == []
|
||||
|
||||
def test_legitimate_groq_api_base_still_resolves(self):
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_secret_str",
|
||||
return_value="server-real-groq-key",
|
||||
):
|
||||
_, provider, dynamic_api_key, _ = get_llm_provider(
|
||||
model="some-model",
|
||||
api_base="https://api.groq.com/openai/v1",
|
||||
)
|
||||
|
||||
assert provider == "groq"
|
||||
assert dynamic_api_key == "server-real-groq-key"
|
||||
|
|
@ -5,6 +5,8 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext
|
|||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_get_customer_id_from_standard_headers,
|
||||
|
|
@ -15,6 +17,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
get_key_model_tpm_limit,
|
||||
get_project_model_rpm_limit,
|
||||
get_project_model_tpm_limit,
|
||||
is_request_body_safe,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -660,3 +663,304 @@ class TestCheckCompleteCredentials:
|
|||
def test_returns_true_when_api_key_is_valid(self):
|
||||
result = check_complete_credentials({"model": "gpt-4", "api_key": "sk-valid"})
|
||||
assert result is True
|
||||
|
||||
|
||||
class TestCheckCompleteCredentialsBlocksSSRF:
|
||||
"""
|
||||
Even with credentials supplied, ``api_base`` / ``base_url`` must not
|
||||
point at private / internal / cloud-metadata addresses. Without this
|
||||
the gate accepts ``api_key=anything`` plus a malicious target and the
|
||||
proxy is used as an SSRF pivot.
|
||||
|
||||
The check only runs when ``litellm.user_url_validation`` is True, so
|
||||
every test in this class flips the toggle. Tests stay mock-only — no
|
||||
real DNS is performed.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _enable_url_validation(self, monkeypatch):
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "user_url_validation", True, raising=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url_field",
|
||||
["api_base", "base_url"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"blocked_url",
|
||||
[
|
||||
"http://169.254.169.254/latest/meta-data/iam/security-credentials/",
|
||||
"http://metadata.google.internal/computeMetadata/v1/",
|
||||
"http://127.0.0.1:8080/admin",
|
||||
"http://10.0.0.1/",
|
||||
"http://192.168.1.1/",
|
||||
],
|
||||
)
|
||||
def test_rejects_private_or_metadata_targets(self, url_field, blocked_url):
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_utils.validate_url",
|
||||
side_effect=SSRFError(f"blocked: {blocked_url}"),
|
||||
):
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
check_complete_credentials(
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-some-clientside-key",
|
||||
url_field: blocked_url,
|
||||
}
|
||||
)
|
||||
assert url_field in str(exc_info.value)
|
||||
assert "SSRF" in str(exc_info.value)
|
||||
|
||||
def test_allows_public_target_when_validate_url_passes(self):
|
||||
# ``validate_url`` is mocked so no real DNS is performed.
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_utils.validate_url",
|
||||
return_value=("https://api.openai.com/v1", "api.openai.com"),
|
||||
):
|
||||
result = check_complete_credentials(
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-some-clientside-key",
|
||||
"api_base": "https://api.openai.com/v1",
|
||||
}
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_skips_url_validation_when_toggle_is_off(self, monkeypatch):
|
||||
# Admins who disable ``user_url_validation`` (default) should not
|
||||
# have requests rejected at the proxy boundary even if the URL
|
||||
# would fail the SSRF guard.
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
|
||||
with patch(
|
||||
"litellm.proxy.auth.auth_utils.validate_url",
|
||||
) as mocked:
|
||||
result = check_complete_credentials(
|
||||
{
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-some-clientside-key",
|
||||
"api_base": "http://127.0.0.1:8080/admin",
|
||||
}
|
||||
)
|
||||
assert result is True
|
||||
mocked.assert_not_called()
|
||||
|
||||
|
||||
class TestGetDynamicLitellmParamsClearsAdminConfigOnBaseOverride:
|
||||
"""
|
||||
When the caller redirects ``api_base`` / ``base_url`` to their own
|
||||
server, admin-set fields like ``OpenAI-Organization``, ``extra_body``,
|
||||
AWS / Vertex / Azure tokens, and per-deployment ``api_version`` must
|
||||
NOT flow through to that destination.
|
||||
"""
|
||||
|
||||
def test_clears_admin_organization_and_extra_body_on_base_override(self):
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
admin_params = {
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-admin-key",
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin-corp",
|
||||
"extra_body": {"x-admin-secret": "super-secret"},
|
||||
"api_version": "2026-04-01",
|
||||
}
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params=dict(admin_params),
|
||||
request_kwargs={
|
||||
"api_key": "sk-attacker",
|
||||
"api_base": "https://attacker.example",
|
||||
},
|
||||
)
|
||||
assert out["api_base"] == "https://attacker.example"
|
||||
assert out["api_key"] == "sk-attacker"
|
||||
assert "organization" not in out
|
||||
assert "extra_body" not in out
|
||||
assert "api_version" not in out
|
||||
|
||||
def test_clears_aws_and_vertex_secrets_on_base_override(self):
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
admin_params = {
|
||||
"model": "bedrock/claude-3",
|
||||
"aws_access_key_id": "AKIA-EXAMPLE",
|
||||
"aws_secret_access_key": "secret-example",
|
||||
"aws_session_token": "session-example",
|
||||
"vertex_credentials": '{"private_key":"-----BEGIN..."}',
|
||||
"vertex_project": "admin-gcp-project",
|
||||
}
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params=dict(admin_params),
|
||||
request_kwargs={"base_url": "https://attacker.example"},
|
||||
)
|
||||
assert "aws_access_key_id" not in out
|
||||
assert "aws_secret_access_key" not in out
|
||||
assert "aws_session_token" not in out
|
||||
assert "vertex_credentials" not in out
|
||||
assert "vertex_project" not in out
|
||||
|
||||
def test_caller_resupplied_value_overrides_admin_value_on_base_override(self):
|
||||
# When the caller redirects ``api_base`` and *also* supplies their
|
||||
# own value for one of the admin fields (e.g. ``organization``),
|
||||
# the caller's value must win — never the admin's. The naive
|
||||
# ``if field not in request_kwargs: pop`` shape lets a caller echo
|
||||
# the field name with any value (or empty string) to keep the
|
||||
# admin's value forwarded, which is the exfiltration vector this
|
||||
# test guards against.
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin",
|
||||
"extra_body": {"admin": "value"},
|
||||
},
|
||||
request_kwargs={
|
||||
"api_base": "https://attacker.example",
|
||||
"organization": "org-attacker",
|
||||
"extra_body": {"attacker": "value"},
|
||||
},
|
||||
)
|
||||
assert out["organization"] == "org-attacker"
|
||||
assert out["extra_body"] == {"attacker": "value"}
|
||||
|
||||
def test_field_echo_does_not_preserve_admin_value(self):
|
||||
# Regression: a caller that echoes an admin-config field name with
|
||||
# an *empty* value (or any value) must not be able to keep the
|
||||
# admin's value in ``litellm_params``.
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin-secret",
|
||||
"extra_body": {"x-admin-only": "secret"},
|
||||
},
|
||||
request_kwargs={
|
||||
"api_base": "https://attacker.example",
|
||||
"organization": "",
|
||||
"extra_body": "",
|
||||
},
|
||||
)
|
||||
assert out["organization"] == ""
|
||||
assert out["extra_body"] == ""
|
||||
assert "org-admin-secret" not in str(out)
|
||||
|
||||
def test_no_clearing_when_only_api_key_overridden(self):
|
||||
from litellm.router_utils.clientside_credential_handler import (
|
||||
get_dynamic_litellm_params,
|
||||
)
|
||||
|
||||
# Caller only overrides api_key (BYOK pattern); admin's organization /
|
||||
# extra_body / region still apply because the destination is unchanged.
|
||||
out = get_dynamic_litellm_params(
|
||||
litellm_params={
|
||||
"api_base": "https://admin.upstream/v1",
|
||||
"organization": "org-admin",
|
||||
"api_version": "2026-04-01",
|
||||
},
|
||||
request_kwargs={"api_key": "sk-byok"},
|
||||
)
|
||||
assert out["organization"] == "org-admin"
|
||||
assert out["api_version"] == "2026-04-01"
|
||||
assert out["api_base"] == "https://admin.upstream/v1"
|
||||
|
||||
|
||||
class TestIsRequestBodySafeBlocksEndpointTargetingFields:
|
||||
"""
|
||||
``is_request_body_safe`` rejects request-body fields that retarget the
|
||||
outbound request to a caller-controlled host. Beyond the original
|
||||
``api_base`` / ``base_url``, the same protection must apply to:
|
||||
|
||||
* ``aws_bedrock_runtime_endpoint`` — Bedrock endpoint redirect; an
|
||||
attacker-controlled value coerces the proxy to authenticate against
|
||||
their host with the admin's AWS creds.
|
||||
* ``langsmith_base_url`` — Langsmith callback host; attacker-controlled
|
||||
values exfiltrate the entire request payload (incl. message content)
|
||||
via the observability hook.
|
||||
* ``langfuse_host`` — same exfil vector via the Langfuse hook.
|
||||
"""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _disable_url_validation(self, monkeypatch):
|
||||
# The new banned-params entries should be rejected even when
|
||||
# ``user_url_validation`` is off — the gate isn't the URL guard,
|
||||
# it's the banned-params list.
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False, raising=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field",
|
||||
[
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"langsmith_base_url",
|
||||
"langfuse_host",
|
||||
"posthog_host",
|
||||
"braintrust_host",
|
||||
"slack_webhook_url",
|
||||
"s3_endpoint_url",
|
||||
"sagemaker_base_url",
|
||||
"deployment_url",
|
||||
],
|
||||
)
|
||||
def test_endpoint_targeting_field_in_request_body_is_rejected(self, field):
|
||||
with pytest.raises(ValueError) as exc:
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", field: "https://attacker.example"},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
# The function lists the offending param name in the error.
|
||||
assert field in str(exc.value)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field",
|
||||
["api_base", "base_url", "user_config", "langfuse_host", "slack_webhook_url"],
|
||||
)
|
||||
def test_api_key_does_not_bypass_blocklist(self, field):
|
||||
# Regression: the historical ``check_complete_credentials`` clause
|
||||
# made the entire blocklist a no-op for any caller that supplied
|
||||
# a non-empty ``api_key``. That bypass turned every missing entry
|
||||
# on the blocklist into an SSRF / credential-exfil hole. Verify
|
||||
# that supplying an api_key (alongside the banned param) does NOT
|
||||
# bypass the gate — it can only be opened by an admin opt-in.
|
||||
with pytest.raises(ValueError) as exc:
|
||||
is_request_body_safe(
|
||||
request_body={
|
||||
"model": "gpt-4",
|
||||
"api_key": "sk-anything",
|
||||
field: "https://attacker.example",
|
||||
},
|
||||
general_settings={},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
assert field in str(exc.value)
|
||||
|
||||
def test_admin_opt_in_proxy_wide_still_allows(self):
|
||||
# ``general_settings.allow_client_side_credentials = True`` remains
|
||||
# the documented proxy-wide BYOK opt-in.
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body={"model": "gpt-4", "api_base": "https://my-byok.example"},
|
||||
general_settings={"allow_client_side_credentials": True},
|
||||
llm_router=None,
|
||||
model="gpt-4",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
|
|
|||
|
|
@ -35,18 +35,18 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging):
|
|||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
result = await client.attempt_db_reconnect(
|
||||
reason="unit_test_reconnect_success",
|
||||
force=True,
|
||||
)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
result = await client.attempt_db_reconnect(
|
||||
reason="unit_test_reconnect_success",
|
||||
force=True,
|
||||
)
|
||||
|
||||
assert result is True
|
||||
client.db.disconnect.assert_awaited_once()
|
||||
client.db.connect.assert_awaited_once()
|
||||
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
|
||||
client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
|
||||
|
||||
|
|
@ -141,15 +141,19 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
|
|||
)
|
||||
client._db_last_reconnect_attempt_ts = 0.0
|
||||
client._db_reconnect_cooldown_seconds = 10
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
# Use a counter-based mock to avoid StopIteration when time.time() is called
|
||||
# more times than expected (varies by Python version / internal code paths).
|
||||
fake_clock = iter(range(100, 10000))
|
||||
with patch(
|
||||
"litellm.proxy.utils.time.time", side_effect=lambda: float(next(fake_clock))
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.utils.time.time",
|
||||
side_effect=lambda: float(next(fake_clock)),
|
||||
),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
result = await client.attempt_db_reconnect(
|
||||
reason="unit_test_cooldown_timestamp_after_attempt",
|
||||
|
|
@ -163,23 +167,28 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_reconnect_cycle_watchdog_should_use_direct_db_ops(
|
||||
async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Direct reconnect goes through recreate_prisma_client (which non-blockingly
|
||||
kills the old engine) instead of calling disconnect() — see issue #26191.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.disconnect = AsyncMock(side_effect=AssertionError("wrapper disconnect used"))
|
||||
client.connect = AsyncMock(side_effect=AssertionError("wrapper connect used"))
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.disconnect = AsyncMock(
|
||||
side_effect=AssertionError("disconnect must not be called")
|
||||
)
|
||||
client.db.recreate_prisma_client = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
await client._run_reconnect_cycle(timeout_seconds=None)
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
await client._run_reconnect_cycle(timeout_seconds=None)
|
||||
|
||||
client.db.disconnect.assert_awaited_once()
|
||||
client.db.connect.assert_awaited_once()
|
||||
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
|
||||
client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
client.db.disconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -190,19 +199,22 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget(
|
|||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client._db_watchdog_reconnect_timeout_seconds = 0.1
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
async def _slow_connect():
|
||||
async def _slow_recreate(_db_url):
|
||||
await asyncio.sleep(0.08)
|
||||
|
||||
async def _slow_query(_query: str):
|
||||
await asyncio.sleep(0.08)
|
||||
return [{"result": 1}]
|
||||
|
||||
client.db.connect = AsyncMock(side_effect=_slow_connect)
|
||||
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
|
||||
client.db.query_raw = AsyncMock(side_effect=_slow_query)
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
with (
|
||||
pytest.raises(asyncio.TimeoutError),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
await client._run_reconnect_cycle(timeout_seconds=None)
|
||||
|
||||
|
||||
|
|
@ -213,19 +225,22 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget(
|
|||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
|
||||
async def _slow_connect():
|
||||
async def _slow_recreate(_db_url):
|
||||
await asyncio.sleep(0.08)
|
||||
|
||||
async def _slow_query(_query: str):
|
||||
await asyncio.sleep(0.08)
|
||||
return [{"result": 1}]
|
||||
|
||||
client.db.connect = AsyncMock(side_effect=_slow_connect)
|
||||
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
|
||||
client.db.query_raw = AsyncMock(side_effect=_slow_query)
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
with (
|
||||
pytest.raises(asyncio.TimeoutError),
|
||||
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
|
||||
):
|
||||
await client._run_reconnect_cycle(timeout_seconds=0.1)
|
||||
|
||||
|
||||
|
|
@ -320,45 +335,35 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lightweight_reconnect_kills_engine_on_disconnect_failure(
|
||||
async def test_recreate_prisma_client_kills_old_engine_without_disconnect(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Lightweight reconnect must kill the old engine PID when disconnect() fails."""
|
||||
"""recreate_prisma_client SIGTERMs the old engine PID directly rather than
|
||||
calling `disconnect()`, which blocks the asyncio event loop on the sync
|
||||
`subprocess.Popen.wait()` inside prisma-client-py — see issue #26191.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(side_effect=Exception("disconnect failed"))
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
disconnect_mock = AsyncMock(
|
||||
side_effect=AssertionError("disconnect must not be called on reconnect path")
|
||||
)
|
||||
client.db._original_prisma.disconnect = disconnect_mock
|
||||
|
||||
with (
|
||||
patch.object(client, "_get_engine_pid", return_value=9999),
|
||||
patch("os.kill") as mock_kill,
|
||||
patch("asyncio.sleep", new_callable=AsyncMock),
|
||||
patch.object(client.db, "_get_engine_pid", return_value=9999),
|
||||
patch("litellm.proxy.db.prisma_client.os.kill") as mock_kill,
|
||||
patch("litellm.proxy.db.prisma_client.asyncio.sleep", new_callable=AsyncMock),
|
||||
):
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
# Return a Prisma instance whose connect() is awaitable.
|
||||
fake_new_prisma = MagicMock()
|
||||
fake_new_prisma.connect = AsyncMock(return_value=None)
|
||||
with patch("prisma.Prisma", return_value=fake_new_prisma):
|
||||
await client.db.recreate_prisma_client("postgresql://test")
|
||||
|
||||
mock_kill.assert_any_call(9999, signal.SIGTERM)
|
||||
client.db.connect.assert_awaited_once()
|
||||
client.db.query_raw.assert_awaited_once_with("SELECT 1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Lightweight reconnect must NOT kill when disconnect() succeeds."""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client.db.disconnect = AsyncMock(return_value=None)
|
||||
client.db.connect = AsyncMock(return_value=None)
|
||||
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
|
||||
|
||||
with patch("os.kill") as mock_kill:
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
mock_kill.assert_not_called()
|
||||
disconnect_mock.assert_not_awaited()
|
||||
fake_new_prisma.connect.assert_awaited_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -0,0 +1,611 @@
|
|||
"""
|
||||
Unit tests for workflow management endpoints (/v1/workflows/runs/*).
|
||||
Uses FastAPI TestClient with a mocked prisma_client.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.management_endpoints.workflow_management_endpoints import router
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_run(
|
||||
run_id: str = "run-1",
|
||||
session_id: str = "sess-1",
|
||||
workflow_type: str = "shin-builder",
|
||||
status: str = "pending",
|
||||
created_by: Any = "tok-test",
|
||||
) -> MagicMock:
|
||||
obj = MagicMock()
|
||||
obj.run_id = run_id
|
||||
obj.session_id = session_id
|
||||
obj.workflow_type = workflow_type
|
||||
obj.status = status
|
||||
obj.created_by = created_by
|
||||
obj.created_at = datetime.now(timezone.utc)
|
||||
obj.updated_at = datetime.now(timezone.utc)
|
||||
obj.input = None
|
||||
obj.output = None
|
||||
obj.metadata = None
|
||||
return obj
|
||||
|
||||
|
||||
def _make_event(
|
||||
event_id: str = "evt-1",
|
||||
run_id: str = "run-1",
|
||||
event_type: str = "step.started",
|
||||
step_name: str = "grill",
|
||||
sequence_number: int = 0,
|
||||
) -> MagicMock:
|
||||
obj = MagicMock()
|
||||
obj.event_id = event_id
|
||||
obj.run_id = run_id
|
||||
obj.event_type = event_type
|
||||
obj.step_name = step_name
|
||||
obj.sequence_number = sequence_number
|
||||
obj.data = None
|
||||
obj.created_at = datetime.now(timezone.utc)
|
||||
return obj
|
||||
|
||||
|
||||
def _make_message(
|
||||
message_id: str = "msg-1",
|
||||
run_id: str = "run-1",
|
||||
role: str = "user",
|
||||
content: str = "hello",
|
||||
sequence_number: int = 0,
|
||||
) -> MagicMock:
|
||||
obj = MagicMock()
|
||||
obj.message_id = message_id
|
||||
obj.run_id = run_id
|
||||
obj.role = role
|
||||
obj.content = content
|
||||
obj.sequence_number = sequence_number
|
||||
obj.session_id = None
|
||||
obj.created_at = datetime.now(timezone.utc)
|
||||
return obj
|
||||
|
||||
|
||||
def _make_tx(event_return=None, run_return=None, msg_return=None) -> MagicMock:
|
||||
"""Build an async context-manager mock for prisma_client.db.tx()."""
|
||||
tx = MagicMock()
|
||||
tx.litellm_workflowevent = MagicMock()
|
||||
tx.litellm_workflowevent.create = AsyncMock(
|
||||
return_value=event_return or _make_event()
|
||||
)
|
||||
tx.litellm_workflowrun = MagicMock()
|
||||
tx.litellm_workflowrun.update = AsyncMock(return_value=run_return or _make_run())
|
||||
tx.litellm_workflowmessage = MagicMock()
|
||||
tx.litellm_workflowmessage.create = AsyncMock(
|
||||
return_value=msg_return or _make_message()
|
||||
)
|
||||
tx.__aenter__ = AsyncMock(return_value=tx)
|
||||
tx.__aexit__ = AsyncMock(return_value=False)
|
||||
return tx
|
||||
|
||||
|
||||
def _make_prisma_client() -> MagicMock:
|
||||
client = MagicMock()
|
||||
client.db = MagicMock()
|
||||
client.db.litellm_workflowrun = MagicMock()
|
||||
client.db.litellm_workflowevent = MagicMock()
|
||||
client.db.litellm_workflowmessage = MagicMock()
|
||||
# default tx() returns a no-op transaction
|
||||
client.db.tx = MagicMock(return_value=_make_tx())
|
||||
return client
|
||||
|
||||
|
||||
def _make_app() -> FastAPI:
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
return app
|
||||
|
||||
|
||||
def _override_auth() -> Any:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
auth = UserAPIKeyAuth(api_key="sk-test", user_id="admin")
|
||||
auth.token = "tok-test"
|
||||
return auth
|
||||
|
||||
|
||||
def _override_auth_admin() -> Any:
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
auth = UserAPIKeyAuth(api_key="sk-master")
|
||||
auth.user_role = LitellmUserRoles.PROXY_ADMIN # type: ignore[assignment]
|
||||
return auth
|
||||
|
||||
|
||||
def _override_auth_user_with_token(token: str = "tok-abc") -> Any:
|
||||
"""Return a non-admin caller whose hashed token equals `token`."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
auth = UserAPIKeyAuth(api_key="sk-user", user_id="user-1")
|
||||
auth.token = token # override the computed hash with a predictable value
|
||||
return auth
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCreateWorkflowRun:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_create_returns_run(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.create = AsyncMock(return_value=_make_run())
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs",
|
||||
json={"workflow_type": "shin-builder"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
self._prisma.db.litellm_workflowrun.create.assert_awaited_once()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client", None)
|
||||
def test_create_500_when_no_db(self):
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs",
|
||||
json={"workflow_type": "shin-builder"},
|
||||
)
|
||||
assert resp.status_code == 500
|
||||
|
||||
|
||||
class TestListWorkflowRuns:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_returns_runs(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(
|
||||
return_value=[_make_run()]
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["count"] == 1
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_filters_by_status(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs?status=running")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert call_kwargs["where"]["status"] == "running"
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_filters_by_multiple_statuses(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs?status=running,paused")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert call_kwargs["where"]["status"] == {"in": ["running", "paused"]}
|
||||
|
||||
|
||||
class TestGetWorkflowRun:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_get_existing_run(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 200
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_get_missing_run_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/nonexistent")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestUpdateWorkflowRun:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_update_status(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
updated = _make_run(status="completed")
|
||||
self._prisma.db.litellm_workflowrun.update = AsyncMock(return_value=updated)
|
||||
|
||||
resp = self.client.patch(
|
||||
"/v1/workflows/runs/run-1", json={"status": "completed"}
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
self._prisma.db.litellm_workflowrun.update.assert_awaited_once()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_update_no_fields_returns_400(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
resp = self.client.patch("/v1/workflows/runs/run-1", json={})
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
class TestAppendWorkflowEvent:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_event_updates_run_status(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
# _require_run check
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
tx = _make_tx(
|
||||
event_return=_make_event(), run_return=_make_run(status="running")
|
||||
)
|
||||
self._prisma.db.tx = MagicMock(return_value=tx)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "step.started", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
# run status updated inside tx
|
||||
tx.litellm_workflowrun.update.assert_awaited_once()
|
||||
update_call = tx.litellm_workflowrun.update.call_args[1]
|
||||
assert update_call["data"]["status"] == "running"
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_event_no_status_update_for_unknown_type(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
tx = _make_tx(event_return=_make_event(event_type="custom.event"))
|
||||
self._prisma.db.tx = MagicMock(return_value=tx)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "custom.event", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
# no status update inside tx for unknown event_type
|
||||
tx.litellm_workflowrun.update.assert_not_awaited()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_sequence_number_increments(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
existing = _make_event(sequence_number=4)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(
|
||||
return_value=[existing]
|
||||
)
|
||||
tx = _make_tx(event_return=_make_event(sequence_number=5))
|
||||
self._prisma.db.tx = MagicMock(return_value=tx)
|
||||
|
||||
self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "step.started", "step_name": "plan"},
|
||||
)
|
||||
create_call = tx.litellm_workflowevent.create.call_args[1]
|
||||
assert create_call["data"]["sequence_number"] == 5
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_unknown_run_id_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/nonexistent/events",
|
||||
json={"event_type": "step.started", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_sequence_collision_retries_and_succeeds(self, mock_pc):
|
||||
"""UniqueViolationError on first attempt triggers retry; second attempt succeeds."""
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# First tx raises UniqueViolationError; second succeeds.
|
||||
tx_fail = _make_tx()
|
||||
tx_fail.__aenter__ = AsyncMock(return_value=tx_fail)
|
||||
tx_fail.litellm_workflowevent.create = AsyncMock(
|
||||
side_effect=UniqueViolationError(
|
||||
{"user_facing_error": {"message": "unique"}}
|
||||
)
|
||||
)
|
||||
tx_fail.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
tx_ok = _make_tx(event_return=_make_event(sequence_number=1))
|
||||
|
||||
self._prisma.db.tx = MagicMock(side_effect=[tx_fail, tx_ok])
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/events",
|
||||
json={"event_type": "step.started", "step_name": "grill"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
|
||||
class TestWorkflowMessages:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_message(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
|
||||
self._prisma.db.litellm_workflowmessage.create = AsyncMock(
|
||||
return_value=_make_message()
|
||||
)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/run-1/messages",
|
||||
json={"role": "user", "content": "fix the bug"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_append_message_unknown_run_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.post(
|
||||
"/v1/workflows/runs/nonexistent/messages",
|
||||
json={"role": "user", "content": "hello"},
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_messages_ordered(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_make_message(sequence_number=0),
|
||||
_make_message(sequence_number=1, role="assistant"),
|
||||
]
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/messages")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["count"] == 2
|
||||
call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1]
|
||||
assert call_kwargs["order"] == {"sequence_number": "asc"}
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_messages_respects_limit(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowmessage.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/messages?limit=25")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowmessage.find_many.call_args[1]
|
||||
assert call_kwargs["take"] == 25
|
||||
|
||||
|
||||
class TestListWorkflowEvents:
|
||||
def setup_method(self):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = _override_auth
|
||||
self.client = TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_events_ordered(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_make_event(sequence_number=0),
|
||||
_make_event(sequence_number=1),
|
||||
]
|
||||
)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/events")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["count"] == 2
|
||||
call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1]
|
||||
assert call_kwargs["order"] == {"sequence_number": "asc"}
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_events_respects_limit(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run()
|
||||
)
|
||||
self._prisma.db.litellm_workflowevent.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/run-1/events?limit=10")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowevent.find_many.call_args[1]
|
||||
assert call_kwargs["take"] == 10
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_list_events_unknown_run_returns_404(self, mock_pc):
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
resp = self.client.get("/v1/workflows/runs/nonexistent/events")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestTenantIsolation:
|
||||
"""Ownership enforcement: non-admin callers only see their own runs."""
|
||||
|
||||
def _make_app_with_auth(self, auth_fn):
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
self._prisma = _make_prisma_client()
|
||||
app = _make_app()
|
||||
app.dependency_overrides[user_api_key_auth] = auth_fn
|
||||
return TestClient(app, raise_server_exceptions=True)
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_create_stores_caller_token(self, mock_pc):
|
||||
token = "tok-owner"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.create = AsyncMock(
|
||||
return_value=_make_run(created_by=token)
|
||||
)
|
||||
|
||||
resp = client.post("/v1/workflows/runs", json={"workflow_type": "test"})
|
||||
assert resp.status_code == 200
|
||||
create_call = self._prisma.db.litellm_workflowrun.create.call_args[1]
|
||||
assert create_call["data"]["created_by"] == token
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_list_scoped_to_caller_token(self, mock_pc):
|
||||
token = "tok-owner"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = client.get("/v1/workflows/runs")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert call_kwargs["where"].get("created_by") == token
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_admin_list_not_scoped(self, mock_pc):
|
||||
client = self._make_app_with_auth(_override_auth_admin)
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_many = AsyncMock(return_value=[])
|
||||
|
||||
resp = client.get("/v1/workflows/runs")
|
||||
assert resp.status_code == 200
|
||||
call_kwargs = self._prisma.db.litellm_workflowrun.find_many.call_args[1]
|
||||
assert "created_by" not in call_kwargs["where"]
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_get_other_users_run_returns_404(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
# Run owned by a different key
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by="tok-other-owner")
|
||||
)
|
||||
|
||||
resp = client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_get_null_owner_run_returns_404(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by=None)
|
||||
)
|
||||
|
||||
resp = client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 404
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_update_null_owner_run_returns_404(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by=None)
|
||||
)
|
||||
self._prisma.db.litellm_workflowrun.update = AsyncMock(
|
||||
return_value=_make_run(status="completed")
|
||||
)
|
||||
|
||||
resp = client.patch("/v1/workflows/runs/run-1", json={"status": "completed"})
|
||||
assert resp.status_code == 404
|
||||
self._prisma.db.litellm_workflowrun.update.assert_not_awaited()
|
||||
|
||||
@patch("litellm.proxy.proxy_server.prisma_client")
|
||||
def test_non_admin_get_own_run_succeeds(self, mock_pc):
|
||||
token = "tok-caller"
|
||||
client = self._make_app_with_auth(lambda: _override_auth_user_with_token(token))
|
||||
mock_pc.db = self._prisma.db
|
||||
self._prisma.db.litellm_workflowrun.find_unique = AsyncMock(
|
||||
return_value=_make_run(created_by=token)
|
||||
)
|
||||
|
||||
resp = client.get("/v1/workflows/runs/run-1")
|
||||
assert resp.status_code == 200
|
||||
|
|
@ -40,6 +40,7 @@ import { ProjectsPage } from "@/components/Projects/ProjectsPage";
|
|||
import VectorStoreManagement from "@/components/vector_store_management";
|
||||
import ToolPoliciesView from "@/components/ToolPoliciesView";
|
||||
import { MemoryView } from "@/components/MemoryView";
|
||||
import WorkflowRuns from "@/components/workflow_runs";
|
||||
import SpendLogsTable from "@/components/view_logs";
|
||||
import ViewUserDashboard from "@/components/view_users";
|
||||
import { ThemeProvider } from "@/contexts/ThemeContext";
|
||||
|
|
@ -631,6 +632,8 @@ function CreateKeyPageContent() {
|
|||
<VectorStoreManagement accessToken={accessToken} userRole={userRole} userID={userID} />
|
||||
) : page == "tool-policies" ? (
|
||||
<ToolPoliciesView accessToken={accessToken} userRole={userRole} />
|
||||
) : page == "workflows" ? (
|
||||
<WorkflowRuns accessToken={accessToken} />
|
||||
) : page == "memory" ? (
|
||||
<MemoryView
|
||||
accessToken={accessToken}
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ describe("Sidebar (leftnav)", () => {
|
|||
"Virtual Keys",
|
||||
"Playground",
|
||||
"Models + Endpoints",
|
||||
"Agents",
|
||||
"Agentic",
|
||||
"MCP Servers",
|
||||
"Guardrails",
|
||||
"Policies",
|
||||
|
|
@ -163,10 +163,17 @@ describe("Sidebar (leftnav)", () => {
|
|||
expect(screen.getByText("Models + Endpoints")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows Agents to Admin Viewer (read-only)", () => {
|
||||
it("shows Agents (under Agentic) to Admin Viewer (read-only)", async () => {
|
||||
mockUseAuthorized.mockReturnValueOnce(adminViewerAuth);
|
||||
renderWithProviders(<Sidebar {...defaultProps} />);
|
||||
expect(screen.getByText("Agents")).toBeInTheDocument();
|
||||
// Agents is now nested under the "Agentic" submenu — expand parent
|
||||
// first to render the children, then assert Agents is visible.
|
||||
act(() => {
|
||||
fireEvent.click(screen.getByText("Agentic"));
|
||||
});
|
||||
await waitFor(() => {
|
||||
expect(screen.getByText("Agents")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("shows Logs to Admin Viewer", () => {
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
|
|||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import {
|
||||
ApiOutlined,
|
||||
ApartmentOutlined,
|
||||
AppstoreOutlined,
|
||||
AuditOutlined,
|
||||
BankOutlined,
|
||||
|
|
@ -129,11 +130,33 @@ const menuGroups: MenuGroup[] = [
|
|||
roles: rolesAllowedToViewWriteScopedPages,
|
||||
},
|
||||
{
|
||||
key: "agents",
|
||||
page: "agents",
|
||||
label: "Agents",
|
||||
key: "agentic",
|
||||
page: "agentic",
|
||||
label: "Agentic",
|
||||
icon: <RobotOutlined />,
|
||||
roles: rolesAllowedToViewWriteScopedPages,
|
||||
children: [
|
||||
{
|
||||
key: "agents",
|
||||
page: "agents",
|
||||
label: "Agents",
|
||||
icon: <RobotOutlined />,
|
||||
// Admin Viewer can view agents read-only (write actions are
|
||||
// hidden inside the page); Playground above stays write-only.
|
||||
roles: rolesAllowedToViewWriteScopedPages,
|
||||
},
|
||||
{
|
||||
key: "workflows",
|
||||
page: "workflows",
|
||||
label: "Workflow Runs",
|
||||
icon: <ApartmentOutlined />,
|
||||
},
|
||||
{
|
||||
key: "memory",
|
||||
page: "memory",
|
||||
label: "Memory",
|
||||
icon: <BookOutlined />,
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
key: "mcp-servers",
|
||||
|
|
@ -148,12 +171,6 @@ const menuGroups: MenuGroup[] = [
|
|||
icon: <ApiOutlined />,
|
||||
roles: all_admin_roles,
|
||||
},
|
||||
{
|
||||
key: "memory",
|
||||
page: "memory",
|
||||
label: "Memory",
|
||||
icon: <BookOutlined />,
|
||||
},
|
||||
{
|
||||
key: "guardrails",
|
||||
page: "guardrails",
|
||||
|
|
|
|||
|
|
@ -9,6 +9,8 @@ export const pageDescriptions: Record<string, string> = {
|
|||
"llm-playground": "Interactive playground for testing LLM requests",
|
||||
models: "Configure and manage LLM models and endpoints",
|
||||
agents: "Create and manage AI agents",
|
||||
agentic: "Manage agentic resources: agents, workflow runs, and memory",
|
||||
workflows: "Track and inspect durable workflow run history",
|
||||
"mcp-servers": "Configure Model Context Protocol servers",
|
||||
memory: "Inspect and manage agent memory entries stored under /v1/memory",
|
||||
guardrails: "Set up content moderation and safety guardrails",
|
||||
|
|
|
|||
751
ui/litellm-dashboard/src/components/workflow_runs/index.tsx
Normal file
751
ui/litellm-dashboard/src/components/workflow_runs/index.tsx
Normal file
|
|
@ -0,0 +1,751 @@
|
|||
import React, { useState, useEffect, useCallback } from "react";
|
||||
import { Button, Collapse, Drawer, Empty, Spin, Table, Tooltip, Typography } from "antd";
|
||||
import { ReloadOutlined } from "@ant-design/icons";
|
||||
import { proxyBaseUrl } from "@/components/networking";
|
||||
|
||||
const { Text } = Typography;
|
||||
|
||||
interface WorkflowRunsProps {
|
||||
accessToken: string | null;
|
||||
}
|
||||
|
||||
type RunStatus = "pending" | "running" | "paused" | "completed" | "failed";
|
||||
|
||||
interface RunMetadata {
|
||||
title?: string;
|
||||
state?: string;
|
||||
pr_url?: string;
|
||||
worktree_path?: string;
|
||||
plan_text?: string;
|
||||
grill_session_id?: string;
|
||||
session_id?: string;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
interface WorkflowRun {
|
||||
run_id: string;
|
||||
status: RunStatus;
|
||||
workflow_type: string;
|
||||
created_at: string;
|
||||
metadata?: RunMetadata | null;
|
||||
}
|
||||
|
||||
interface WorkflowRunEvent {
|
||||
event_id: string;
|
||||
event_type: string;
|
||||
step_name: string;
|
||||
sequence_number: number;
|
||||
created_at: string;
|
||||
data?: Record<string, unknown> | null;
|
||||
}
|
||||
|
||||
interface WorkflowRunMessage {
|
||||
message_id: string;
|
||||
role: string;
|
||||
content: string;
|
||||
sequence_number: number;
|
||||
created_at: string;
|
||||
}
|
||||
|
||||
// ── design tokens ─────────────────────────────────────────────────────────────
|
||||
|
||||
const STATUS_DOT: Record<RunStatus, string> = {
|
||||
pending: "#a1a1aa",
|
||||
running: "#3b82f6",
|
||||
paused: "#f59e0b",
|
||||
completed: "#22c55e",
|
||||
failed: "#ef4444",
|
||||
};
|
||||
|
||||
const EVENT_COLOR: Record<string, { bar: string; border: string; text: string }> = {
|
||||
"step.started": { bar: "#f0fdf4", border: "#86efac", text: "#16a34a" },
|
||||
"step.failed": { bar: "#fef2f2", border: "#fca5a5", text: "#dc2626" },
|
||||
"hook.waiting": { bar: "#fffbeb", border: "#fcd34d", text: "#d97706" },
|
||||
"hook.received": { bar: "#eff6ff", border: "#93c5fd", text: "#2563eb" },
|
||||
};
|
||||
|
||||
function eventStyle(type: string) {
|
||||
return EVENT_COLOR[type] ?? { bar: "#f4f4f5", border: "#d4d4d8", text: "#52525b" };
|
||||
}
|
||||
|
||||
// ── helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
function timeAgo(iso: string): string {
|
||||
const diff = Date.now() - new Date(iso).getTime();
|
||||
if (isNaN(diff)) return iso;
|
||||
const s = Math.floor(diff / 1000);
|
||||
if (s < 60) return `${s}s ago`;
|
||||
const m = Math.floor(s / 60);
|
||||
if (m < 60) return `${m}m ago`;
|
||||
const h = Math.floor(m / 60);
|
||||
if (h < 24) return `${h}h ago`;
|
||||
return `${Math.floor(h / 24)}d ago`;
|
||||
}
|
||||
|
||||
function fmtDuration(ms: number): string {
|
||||
if (ms < 0) return "";
|
||||
if (ms < 1000) return `${ms}ms`;
|
||||
return `${(ms / 1000).toFixed(1)}s`;
|
||||
}
|
||||
|
||||
function runTitle(run: WorkflowRun): string {
|
||||
const t = run.metadata?.title;
|
||||
if (t) return String(t);
|
||||
return run.workflow_type ?? run.run_id.slice(0, 8);
|
||||
}
|
||||
|
||||
function shortId(id: string): string {
|
||||
return id.slice(0, 8);
|
||||
}
|
||||
|
||||
// ── status dot ────────────────────────────────────────────────────────────────
|
||||
|
||||
const StatusDot: React.FC<{ status: RunStatus; size?: number }> = ({ status, size = 8 }) => (
|
||||
<span
|
||||
style={{
|
||||
display: "inline-block",
|
||||
width: size,
|
||||
height: size,
|
||||
borderRadius: "50%",
|
||||
background: STATUS_DOT[status] ?? "#a1a1aa",
|
||||
flexShrink: 0,
|
||||
}}
|
||||
/>
|
||||
);
|
||||
|
||||
// ── truncated text value ──────────────────────────────────────────────────────
|
||||
|
||||
const TRUNCATE_AT = 120;
|
||||
|
||||
const TruncatedValue: React.FC<{ value: string }> = ({ value }) => {
|
||||
const [expanded, setExpanded] = useState(false);
|
||||
if (value.length <= TRUNCATE_AT) {
|
||||
return <span style={{ color: "#27272a", wordBreak: "break-all" }}>{value}</span>;
|
||||
}
|
||||
return (
|
||||
<span style={{ color: "#27272a", wordBreak: "break-all" }}>
|
||||
{expanded ? value : value.slice(0, TRUNCATE_AT) + "…"}
|
||||
<button
|
||||
onClick={() => setExpanded((e) => !e)}
|
||||
style={{
|
||||
background: "none",
|
||||
border: "none",
|
||||
padding: "0 4px",
|
||||
cursor: "pointer",
|
||||
color: "#2563eb",
|
||||
fontSize: 11,
|
||||
flexShrink: 0,
|
||||
}}
|
||||
>
|
||||
{expanded ? "less" : "more"}
|
||||
</button>
|
||||
</span>
|
||||
);
|
||||
};
|
||||
|
||||
// ── metadata card ─────────────────────────────────────────────────────────────
|
||||
|
||||
const MetadataCard: React.FC<{ run: WorkflowRun }> = ({ run }) => {
|
||||
const meta = run.metadata ?? {};
|
||||
|
||||
const primaryFields: { key: string; label: string }[] = [
|
||||
{ key: "state", label: "state" },
|
||||
{ key: "worktree_path", label: "worktree" },
|
||||
{ key: "grill_session_id", label: "grill session" },
|
||||
{ key: "session_id", label: "session" },
|
||||
];
|
||||
|
||||
const primaryKeys = new Set(["title", ...primaryFields.map((f) => f.key)]);
|
||||
const extraEntries = Object.entries(meta).filter(
|
||||
([k, v]) => !primaryKeys.has(k) && v !== null && v !== undefined && v !== ""
|
||||
);
|
||||
|
||||
return (
|
||||
<div
|
||||
style={{
|
||||
borderRadius: 8,
|
||||
border: "1px solid #e4e4e7",
|
||||
marginBottom: 16,
|
||||
overflow: "hidden",
|
||||
}}
|
||||
>
|
||||
{/* title bar */}
|
||||
<div
|
||||
style={{
|
||||
padding: "14px 20px",
|
||||
borderBottom: "1px solid #f4f4f5",
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
gap: 10,
|
||||
}}
|
||||
>
|
||||
<StatusDot status={run.status} size={10} />
|
||||
<span style={{ fontSize: 14, fontWeight: 600, color: "#18181b", flex: 1 }}>
|
||||
{runTitle(run)}
|
||||
</span>
|
||||
<span
|
||||
style={{
|
||||
fontFamily: "monospace",
|
||||
fontSize: 11,
|
||||
color: "#a1a1aa",
|
||||
background: "#f4f4f5",
|
||||
padding: "2px 8px",
|
||||
borderRadius: 4,
|
||||
}}
|
||||
>
|
||||
{shortId(run.run_id)}
|
||||
</span>
|
||||
<span
|
||||
style={{
|
||||
fontSize: 11,
|
||||
color: "#a1a1aa",
|
||||
background: "#f4f4f5",
|
||||
padding: "2px 8px",
|
||||
borderRadius: 4,
|
||||
}}
|
||||
>
|
||||
{run.workflow_type}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
{/* key fields grid */}
|
||||
<div
|
||||
style={{
|
||||
padding: "12px 20px",
|
||||
display: "grid",
|
||||
gridTemplateColumns: "repeat(auto-fill, minmax(220px, 1fr))",
|
||||
gap: "8px 24px",
|
||||
fontFamily: "monospace",
|
||||
fontSize: 12,
|
||||
}}
|
||||
>
|
||||
<FieldPair label="status">
|
||||
<span style={{ textTransform: "capitalize", color: "#27272a" }}>{run.status}</span>
|
||||
</FieldPair>
|
||||
<FieldPair label="created">
|
||||
<span style={{ color: "#27272a" }}>{timeAgo(run.created_at)}</span>
|
||||
</FieldPair>
|
||||
|
||||
{meta.pr_url && (
|
||||
<FieldPair label="pr">
|
||||
<a
|
||||
href={String(meta.pr_url)}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
style={{ color: "#2563eb", textDecoration: "none", wordBreak: "break-all" }}
|
||||
>
|
||||
{String(meta.pr_url)}
|
||||
</a>
|
||||
</FieldPair>
|
||||
)}
|
||||
|
||||
{primaryFields.map(({ key, label }) => {
|
||||
const v = meta[key];
|
||||
if (v === null || v === undefined || v === "") return null;
|
||||
const str = typeof v === "object" ? JSON.stringify(v) : String(v);
|
||||
return (
|
||||
<FieldPair key={key} label={label}>
|
||||
<TruncatedValue value={str} />
|
||||
</FieldPair>
|
||||
);
|
||||
})}
|
||||
|
||||
{extraEntries.map(([k, v]) => {
|
||||
const str = typeof v === "object" ? JSON.stringify(v) : String(v);
|
||||
return (
|
||||
<FieldPair key={k} label={k}>
|
||||
<TruncatedValue value={str} />
|
||||
</FieldPair>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
const FieldPair: React.FC<{ label: string; children: React.ReactNode }> = ({
|
||||
label,
|
||||
children,
|
||||
}) => (
|
||||
<div style={{ display: "flex", flexDirection: "column", gap: 1 }}>
|
||||
<span style={{ fontSize: 10, color: "#a1a1aa", textTransform: "uppercase", letterSpacing: "0.06em" }}>
|
||||
{label}
|
||||
</span>
|
||||
<span style={{ fontSize: 12 }}>{children}</span>
|
||||
</div>
|
||||
);
|
||||
|
||||
// ── gantt timeline ────────────────────────────────────────────────────────────
|
||||
|
||||
const GanttTimeline: React.FC<{
|
||||
run: WorkflowRun;
|
||||
events: WorkflowRunEvent[];
|
||||
}> = ({ run, events }) => {
|
||||
if (events.length === 0) {
|
||||
return (
|
||||
<div style={{ padding: "16px 0", color: "#a1a1aa", fontSize: 12, fontFamily: "monospace" }}>
|
||||
No events recorded
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
const runStart = new Date(run.created_at).getTime();
|
||||
const eventTimes = events.map((e) => new Date(e.created_at).getTime());
|
||||
const lastTime = Math.max(...eventTimes);
|
||||
const totalSpan = Math.max(lastTime - runStart, 1);
|
||||
const totalDur = fmtDuration(lastTime - runStart);
|
||||
|
||||
return (
|
||||
<div style={{ fontFamily: "monospace", fontSize: 12 }}>
|
||||
{/* ruler */}
|
||||
<div style={{ display: "grid", gridTemplateColumns: "160px 1fr", gap: "0 12px", marginBottom: 2 }}>
|
||||
<div />
|
||||
<div style={{ position: "relative", height: 16 }}>
|
||||
{[0, 100].map((pct) => (
|
||||
<span
|
||||
key={pct}
|
||||
style={{
|
||||
position: "absolute",
|
||||
left: `${pct}%`,
|
||||
transform: pct === 100 ? "translateX(-100%)" : undefined,
|
||||
fontSize: 10,
|
||||
color: "#a1a1aa",
|
||||
}}
|
||||
>
|
||||
{pct === 0 ? "0" : totalDur}
|
||||
</span>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* outer run bar */}
|
||||
<div style={{ display: "grid", gridTemplateColumns: "160px 1fr", gap: "0 12px", marginBottom: 4 }}>
|
||||
<div style={{ color: "#3f3f46", overflow: "hidden", textOverflow: "ellipsis", whiteSpace: "nowrap", paddingTop: 2 }}>
|
||||
{runTitle(run)}
|
||||
</div>
|
||||
<div
|
||||
style={{
|
||||
height: 24,
|
||||
background: "#f4f4f5",
|
||||
border: "1px solid #d4d4d8",
|
||||
borderRadius: 4,
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
paddingLeft: 8,
|
||||
}}
|
||||
>
|
||||
<span style={{ color: "#71717a", fontSize: 11 }}>{totalDur}</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* event rows */}
|
||||
<div style={{ display: "grid", gridTemplateColumns: "160px 1fr", gap: "0 12px", rowGap: 3 }}>
|
||||
{events.map((ev) => {
|
||||
const evTime = new Date(ev.created_at).getTime();
|
||||
const leftPct = ((evTime - runStart) / totalSpan) * 100;
|
||||
|
||||
const nextIdx = events.findIndex((e) => e.sequence_number > ev.sequence_number);
|
||||
const nextTime =
|
||||
nextIdx >= 0
|
||||
? new Date(events[nextIdx].created_at).getTime()
|
||||
: lastTime + Math.max(totalSpan * 0.12, 500);
|
||||
const widthPct = Math.max(8, ((nextTime - evTime) / totalSpan) * 100);
|
||||
const style = eventStyle(ev.event_type);
|
||||
const dur = fmtDuration(nextTime - evTime);
|
||||
|
||||
return (
|
||||
<React.Fragment key={ev.event_id}>
|
||||
<div
|
||||
style={{
|
||||
color: style.text,
|
||||
overflow: "hidden",
|
||||
textOverflow: "ellipsis",
|
||||
whiteSpace: "nowrap",
|
||||
paddingTop: 2,
|
||||
paddingLeft: 12,
|
||||
}}
|
||||
>
|
||||
{ev.step_name || ev.event_type}
|
||||
</div>
|
||||
<div style={{ position: "relative", height: 24 }}>
|
||||
<Tooltip
|
||||
title={
|
||||
<div style={{ fontFamily: "monospace", fontSize: 11, lineHeight: 1.6 }}>
|
||||
<div><span style={{ color: "#a1a1aa" }}>type: </span><span style={{ color: style.text }}>{ev.event_type}</span></div>
|
||||
<div><span style={{ color: "#a1a1aa" }}>step: </span>{ev.step_name}</div>
|
||||
<div><span style={{ color: "#a1a1aa" }}>seq: </span>{ev.sequence_number}</div>
|
||||
<div><span style={{ color: "#a1a1aa" }}>time: </span>{timeAgo(ev.created_at)}</div>
|
||||
{ev.data && Object.keys(ev.data).length > 0 && (
|
||||
<div><span style={{ color: "#a1a1aa" }}>data: </span>{JSON.stringify(ev.data)}</div>
|
||||
)}
|
||||
</div>
|
||||
}
|
||||
>
|
||||
<div
|
||||
style={{
|
||||
position: "absolute",
|
||||
left: `${Math.min(leftPct, 92)}%`,
|
||||
width: `${Math.min(widthPct, 100 - Math.min(leftPct, 92))}%`,
|
||||
height: "100%",
|
||||
background: style.bar,
|
||||
border: `1px solid ${style.border}`,
|
||||
borderRadius: 4,
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
paddingLeft: 8,
|
||||
cursor: "default",
|
||||
overflow: "hidden",
|
||||
gap: 6,
|
||||
}}
|
||||
>
|
||||
<span style={{ color: style.text, whiteSpace: "nowrap", fontSize: 11 }}>{ev.event_type}</span>
|
||||
{dur && <span style={{ color: "#a1a1aa", whiteSpace: "nowrap", fontSize: 11 }}>{dur}</span>}
|
||||
</div>
|
||||
</Tooltip>
|
||||
</div>
|
||||
</React.Fragment>
|
||||
);
|
||||
})}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
// ── message row ───────────────────────────────────────────────────────────────
|
||||
|
||||
const MessageRow: React.FC<{ msg: WorkflowRunMessage }> = ({ msg }) => {
|
||||
const roleColor: Record<string, string> = {
|
||||
user: "#2563eb",
|
||||
assistant: "#16a34a",
|
||||
system: "#7c3aed",
|
||||
tool_result: "#d97706",
|
||||
};
|
||||
const color = roleColor[msg.role] ?? "#52525b";
|
||||
|
||||
return (
|
||||
<div
|
||||
style={{
|
||||
display: "grid",
|
||||
gridTemplateColumns: "80px 1fr",
|
||||
gap: "0 16px",
|
||||
padding: "10px 0",
|
||||
borderBottom: "1px solid #f4f4f5",
|
||||
fontFamily: "monospace",
|
||||
fontSize: 12,
|
||||
alignItems: "start",
|
||||
}}
|
||||
>
|
||||
<span style={{ color, paddingTop: 1 }}>[{msg.role}]</span>
|
||||
<div>
|
||||
<span
|
||||
style={{
|
||||
color: "#27272a",
|
||||
lineHeight: 1.6,
|
||||
whiteSpace: "pre-wrap",
|
||||
wordBreak: "break-word",
|
||||
display: "block",
|
||||
}}
|
||||
>
|
||||
{msg.content}
|
||||
</span>
|
||||
<span style={{ color: "#a1a1aa", fontSize: 11, marginTop: 2, display: "block" }}>
|
||||
{timeAgo(msg.created_at)}
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
// ── main component ────────────────────────────────────────────────────────────
|
||||
|
||||
const WorkflowRuns: React.FC<WorkflowRunsProps> = ({ accessToken }) => {
|
||||
const [runs, setRuns] = useState<WorkflowRun[]>([]);
|
||||
const [loadingRuns, setLoadingRuns] = useState(false);
|
||||
const [selectedRun, setSelectedRun] = useState<WorkflowRun | null>(null);
|
||||
const [events, setEvents] = useState<WorkflowRunEvent[]>([]);
|
||||
const [messages, setMessages] = useState<WorkflowRunMessage[]>([]);
|
||||
const [loadingDetail, setLoadingDetail] = useState(false);
|
||||
const [drawerOpen, setDrawerOpen] = useState(false);
|
||||
|
||||
const fetchRuns = useCallback(async () => {
|
||||
if (!accessToken) return;
|
||||
setLoadingRuns(true);
|
||||
try {
|
||||
const res = await fetch(`${proxyBaseUrl ?? ""}/v1/workflows/runs?limit=100`, {
|
||||
headers: { Authorization: `Bearer ${accessToken}` },
|
||||
});
|
||||
if (!res.ok) throw new Error(`HTTP ${res.status}`);
|
||||
const data = await res.json();
|
||||
setRuns(data.runs ?? []);
|
||||
} catch (err) {
|
||||
console.error("workflow runs fetch failed:", err);
|
||||
} finally {
|
||||
setLoadingRuns(false);
|
||||
}
|
||||
}, [accessToken]);
|
||||
|
||||
const fetchRunDetail = useCallback(
|
||||
async (run: WorkflowRun) => {
|
||||
if (!accessToken) return;
|
||||
setSelectedRun(run);
|
||||
setDrawerOpen(true);
|
||||
setLoadingDetail(true);
|
||||
setEvents([]);
|
||||
setMessages([]);
|
||||
try {
|
||||
const base = proxyBaseUrl ?? "";
|
||||
const [evRes, msgRes] = await Promise.all([
|
||||
fetch(`${base}/v1/workflows/runs/${run.run_id}/events`, {
|
||||
headers: { Authorization: `Bearer ${accessToken}` },
|
||||
}),
|
||||
fetch(`${base}/v1/workflows/runs/${run.run_id}/messages`, {
|
||||
headers: { Authorization: `Bearer ${accessToken}` },
|
||||
}),
|
||||
]);
|
||||
const evData = evRes.ok ? await evRes.json() : { events: [] };
|
||||
const msgData = msgRes.ok ? await msgRes.json() : { messages: [] };
|
||||
setEvents(
|
||||
[...(evData.events ?? [])].sort(
|
||||
(a: WorkflowRunEvent, b: WorkflowRunEvent) => a.sequence_number - b.sequence_number
|
||||
)
|
||||
);
|
||||
setMessages(
|
||||
[...(msgData.messages ?? [])].sort(
|
||||
(a: WorkflowRunMessage, b: WorkflowRunMessage) => a.sequence_number - b.sequence_number
|
||||
)
|
||||
);
|
||||
} catch (err) {
|
||||
console.error("workflow run detail fetch failed:", err);
|
||||
} finally {
|
||||
setLoadingDetail(false);
|
||||
}
|
||||
},
|
||||
[accessToken]
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
fetchRuns();
|
||||
}, [fetchRuns]);
|
||||
|
||||
const columns = [
|
||||
{
|
||||
title: "Run",
|
||||
dataIndex: "run_id",
|
||||
key: "run",
|
||||
render: (_: string, run: WorkflowRun) => (
|
||||
<div style={{ display: "flex", alignItems: "center", gap: 8 }}>
|
||||
<StatusDot status={run.status} size={7} />
|
||||
<div>
|
||||
<div style={{ fontSize: 13, color: "#18181b", fontWeight: 500, lineHeight: 1.4 }}>
|
||||
{runTitle(run)}
|
||||
</div>
|
||||
<div style={{ fontFamily: "monospace", fontSize: 11, color: "#a1a1aa" }}>
|
||||
{shortId(run.run_id)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
),
|
||||
},
|
||||
{
|
||||
title: "Type",
|
||||
dataIndex: "workflow_type",
|
||||
key: "workflow_type",
|
||||
render: (v: string) => (
|
||||
<span style={{ fontFamily: "monospace", fontSize: 12, color: "#71717a" }}>{v}</span>
|
||||
),
|
||||
},
|
||||
{
|
||||
title: "Status",
|
||||
dataIndex: "status",
|
||||
key: "status",
|
||||
render: (status: RunStatus, run: WorkflowRun) => {
|
||||
const state = run.metadata?.state;
|
||||
return (
|
||||
<div style={{ display: "flex", alignItems: "center", gap: 6 }}>
|
||||
<StatusDot status={status} size={7} />
|
||||
<span style={{ fontSize: 12, color: "#52525b", textTransform: "capitalize" }}>
|
||||
{state ?? status}
|
||||
</span>
|
||||
</div>
|
||||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
title: "Created",
|
||||
dataIndex: "created_at",
|
||||
key: "created_at",
|
||||
render: (v: string) => (
|
||||
<span style={{ fontSize: 12, color: "#a1a1aa" }}>{timeAgo(v)}</span>
|
||||
),
|
||||
},
|
||||
];
|
||||
|
||||
return (
|
||||
<div
|
||||
style={{
|
||||
padding: "24px 32px",
|
||||
fontFamily: '-apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif',
|
||||
minHeight: "calc(100vh - 64px)",
|
||||
background: "#fff",
|
||||
}}
|
||||
>
|
||||
{/* page header */}
|
||||
<div
|
||||
style={{
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
justifyContent: "space-between",
|
||||
marginBottom: 20,
|
||||
}}
|
||||
>
|
||||
<div>
|
||||
<div style={{ fontSize: 18, fontWeight: 600, color: "#18181b" }}>Workflow Runs</div>
|
||||
<div style={{ fontSize: 13, color: "#71717a", marginTop: 2 }}>
|
||||
Durable state tracking for agents and automated workflows
|
||||
</div>
|
||||
</div>
|
||||
<Button
|
||||
icon={<ReloadOutlined />}
|
||||
onClick={fetchRuns}
|
||||
loading={loadingRuns}
|
||||
style={{ color: "#71717a", borderColor: "#e4e4e7" }}
|
||||
>
|
||||
Refresh
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{/* runs table — matches logs page density */}
|
||||
<div className="rounded-lg custom-border overflow-x-auto w-full">
|
||||
<Table
|
||||
dataSource={runs}
|
||||
columns={columns}
|
||||
rowKey="run_id"
|
||||
loading={loadingRuns}
|
||||
size="small"
|
||||
pagination={{ pageSize: 50, hideOnSinglePage: true, size: "small" }}
|
||||
onRow={(run) => ({
|
||||
onClick: () => fetchRunDetail(run),
|
||||
style: { cursor: "pointer" },
|
||||
})}
|
||||
locale={{
|
||||
emptyText: (
|
||||
<Empty
|
||||
description={<span style={{ color: "#a1a1aa", fontSize: 13 }}>No workflow runs yet</span>}
|
||||
image={Empty.PRESENTED_IMAGE_SIMPLE}
|
||||
/>
|
||||
),
|
||||
}}
|
||||
className="[&_.ant-table-cell]:py-0.5 [&_.ant-table-thead_.ant-table-cell]:py-1"
|
||||
style={{ border: "none" }}
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* detail drawer */}
|
||||
<Drawer
|
||||
open={drawerOpen}
|
||||
onClose={() => setDrawerOpen(false)}
|
||||
width={680}
|
||||
title={null}
|
||||
closable={false}
|
||||
bodyStyle={{ padding: 0 }}
|
||||
styles={{ body: { padding: 0 } }}
|
||||
>
|
||||
{!selectedRun ? null : loadingDetail ? (
|
||||
<div style={{ display: "flex", justifyContent: "center", padding: 80 }}>
|
||||
<Spin />
|
||||
</div>
|
||||
) : (
|
||||
<div style={{ padding: "24px 28px", fontFamily: '-apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif' }}>
|
||||
{/* drawer close + refresh */}
|
||||
<div
|
||||
style={{
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
justifyContent: "space-between",
|
||||
marginBottom: 16,
|
||||
}}
|
||||
>
|
||||
<button
|
||||
onClick={() => setDrawerOpen(false)}
|
||||
style={{
|
||||
background: "none",
|
||||
border: "none",
|
||||
cursor: "pointer",
|
||||
padding: "4px 0",
|
||||
fontSize: 12,
|
||||
color: "#a1a1aa",
|
||||
display: "flex",
|
||||
alignItems: "center",
|
||||
gap: 4,
|
||||
}}
|
||||
>
|
||||
← close
|
||||
</button>
|
||||
<Button
|
||||
size="small"
|
||||
icon={<ReloadOutlined />}
|
||||
onClick={() => fetchRunDetail(selectedRun)}
|
||||
loading={loadingDetail}
|
||||
style={{ color: "#71717a", borderColor: "#e4e4e7" }}
|
||||
>
|
||||
Refresh
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{/* metadata card — top */}
|
||||
<MetadataCard run={selectedRun} />
|
||||
|
||||
{/* collapsible sections */}
|
||||
<Collapse
|
||||
defaultActiveKey={["timeline"]}
|
||||
ghost={false}
|
||||
style={{ border: "1px solid #e4e4e7", borderRadius: 8, overflow: "hidden" }}
|
||||
items={[
|
||||
{
|
||||
key: "timeline",
|
||||
label: (
|
||||
<span style={{ fontSize: 12, fontWeight: 500, color: "#3f3f46" }}>
|
||||
Timeline
|
||||
<span style={{ marginLeft: 6, fontSize: 11, color: "#a1a1aa", fontWeight: 400 }}>
|
||||
{events.length} {events.length === 1 ? "event" : "events"}
|
||||
</span>
|
||||
</span>
|
||||
),
|
||||
children: (
|
||||
<div style={{ padding: "4px 4px 12px" }}>
|
||||
<GanttTimeline run={selectedRun} events={events} />
|
||||
</div>
|
||||
),
|
||||
},
|
||||
{
|
||||
key: "messages",
|
||||
label: (
|
||||
<span style={{ fontSize: 12, fontWeight: 500, color: "#3f3f46" }}>
|
||||
Messages
|
||||
<span style={{ marginLeft: 6, fontSize: 11, color: "#a1a1aa", fontWeight: 400 }}>
|
||||
{messages.length}
|
||||
</span>
|
||||
</span>
|
||||
),
|
||||
children: messages.length === 0 ? (
|
||||
<div style={{ padding: "12px 4px", color: "#a1a1aa", fontSize: 12, fontFamily: "monospace" }}>
|
||||
No messages
|
||||
</div>
|
||||
) : (
|
||||
<div style={{ paddingBottom: 4 }}>
|
||||
{messages.map((msg) => (
|
||||
<MessageRow key={msg.message_id} msg={msg} />
|
||||
))}
|
||||
</div>
|
||||
),
|
||||
},
|
||||
]}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</Drawer>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default WorkflowRuns;
|
||||
Loading…
Add table
Reference in a new issue