This commit is contained in:
xuhd committed 2026-05-15 15:26:16 +08:00
commit 1574d44cc3
80 files changed
+9366 -1299

No files matched your search

@@ -27,6 +27,7 @@ from app.gateway.routers import (
auth,
channels,
feedback,
llm_metrics,
mcp,
memory,
models,
@@ -36,6 +37,7 @@ from app.gateway.routers import (
scheduled_tasks,
skills,
suggestions,
system_settings,
thread_runs,
threads,
uploads,
@@ -572,6 +574,12 @@ This gateway provides custom endpoints for models, MCP configuration, skills, an
# Scheduled task API is mounted at /api/scheduled-tasks
app.include_router(scheduled_tasks.router)
# System settings API is mounted at /api/system-settings
app.include_router(system_settings.router)
# Admin LLM metrics audit + export
app.include_router(llm_metrics.router)
# Channels API is mounted at /api/channels
app.include_router(channels.router)
@@ -62,6 +62,15 @@ class LocalAuthProvider(AuthProvider):
"""Get user by ID."""
return await self._repo.get_user_by_id(user_id)
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
"""Find a user whose email starts with ``local_part@`` regardless of domain.
Used by the passwordless username-login flow to keep a single
username mapped to the same account across legacy / canonical
email-domain defaults.
"""
return await self._repo.find_user_by_email_local_part(local_part)
async def create_user(self, email: str, password: str | None = None, system_role: str = "user", needs_setup: bool = False) -> User:
"""Create a new local user.
@@ -61,6 +61,26 @@ class UserRepository(ABC):
"""
raise NotImplementedError
@abstractmethod
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
"""Find a user by the local part of their email (the bit before ``@``).
Used by the passwordless username-login flow so that a single
username keeps mapping to the same account even when the synthetic
email domain (e.g. ``cm.com`` vs the legacy ``local.deerflow``)
differs between registration and current login.
Returns the oldest matching user when multiple rows share the same
local part, so collisions resolve deterministically.
Args:
local_part: Substring to match before the ``@``. Case-insensitive.
Returns:
User if found, None otherwise.
"""
raise NotImplementedError
@abstractmethod
async def update_user(self, user: User) -> User:
"""Update an existing user.
@@ -88,6 +88,27 @@ class SQLiteUserRepository(UserRepository):
row = result.scalar_one_or_none()
return self._row_to_user(row) if row is not None else None
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
local = (local_part or "").strip().lower()
if not local:
return None
# Escape SQL LIKE wildcards so usernames like ``foo_bar`` or ``a%b`` do
# not match unrelated rows. ``\`` is the chosen escape char below.
escaped = local.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
pattern = f"{escaped}@%"
# Case-insensitive prefix match on the email field, oldest first so a
# subsequent registration cannot shadow an existing account.
stmt = (
select(UserRow)
.where(func.lower(UserRow.email).like(pattern, escape="\\"))
.order_by(UserRow.created_at.asc())
.limit(1)
)
async with self._sf() as session:
result = await session.execute(stmt)
row = result.scalar_one_or_none()
return self._row_to_user(row) if row is not None else None
async def update_user(self, user: User) -> User:
async with self._sf() as session:
row = await session.get(UserRow, str(user.id))
@@ -92,6 +92,10 @@ async def langgraph_runtime(app: FastAPI) -> AsyncGenerator[None, None]:
app.state.scheduled_task_store = make_scheduled_task_store(sf)
from deerflow.persistence.llm_metrics import make_llm_metrics_store
app.state.llm_metrics_store = make_llm_metrics_store(sf)
from deerflow.persistence.agents import make_agent_store
app.state.agent_store = make_agent_store(sf)
@@ -176,6 +180,11 @@ def get_scheduled_task_store(request: Request) -> ScheduledTaskStore:
return val
def get_llm_metrics_store(request: Request):
"""Return the LLM metrics store, or None when not configured."""
return getattr(request.app.state, "llm_metrics_store", None)
def get_agent_store(request: Request) -> AgentStore:
"""Return the custom-agent store."""
val = getattr(request.app.state, "agent_store", None)
@@ -164,25 +164,40 @@ class MessageResponse(BaseModel):
# ── Helpers ───────────────────────────────────────────────────────────────
def _normalize_username_local_part(username: str) -> str | None:
"""Return the canonical local-part for a username, or None if user typed an email.
Used both to build the synthetic email and to look up legacy accounts that
were registered with a different default domain. Returning ``None`` for
raw-email input tells callers to skip the domain-agnostic fallback —
if you typed ``alice@old.com``, you meant that specific row.
"""
raw = username.strip()
if not raw:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username cannot be empty")
if "@" in raw:
return None
normalized = "".join(ch if (ch.isalnum() or ch in "._-") else "-" for ch in raw.lower()).strip("._-")
if not normalized:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username is invalid")
return normalized
def _username_to_email(username: str) -> str:
"""Map a frontend username to a valid email used by user storage.
- If username already looks like an email (`contains @`), use it directly.
- Otherwise generate `<normalized-username>@<domain>`.
"""
raw = username.strip()
if not raw:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username cannot be empty")
local = _normalize_username_local_part(username)
if local is None:
return username.strip().lower()
if "@" in raw:
return raw.lower()
normalized = "".join(ch if (ch.isalnum() or ch in "._-") else "-" for ch in raw.lower()).strip("._-")
if not normalized:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="username is invalid")
domain = os.getenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", "local.deerflow").strip().lower() or "local.deerflow"
return f"{normalized}@{domain}"
# Default mail domain for passwordless / username-only registrations.
# Kept overridable via env so individual deployments can swap it out, but the
# built-in default targets the cm.com tenant.
domain = os.getenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", "cm.com").strip().lower() or "cm.com"
return f"{local}@{domain}"
def _set_session_cookie(response: Response, token: str, request: Request) -> None:
@@ -368,9 +383,21 @@ async def login_by_username(request: Request, response: Response, body: Username
"""
provider = get_local_provider()
email = _username_to_email(body.username)
local_part = _normalize_username_local_part(body.username)
created = False
user = await provider.get_user_by_email(email)
# Domain-agnostic fallback: if no row exists at the canonical
# ``<username>@<current-default-domain>``, look for an existing user with
# the same local part under any domain. This is the safety net that keeps
# usernames globally unique when the default email domain changes (e.g.
# ``local.deerflow`` → ``cm.com``) — without it, switching the default
# would let the same username log into two different accounts.
if user is None and local_part is not None:
legacy = await provider.find_user_by_email_local_part(local_part)
if legacy is not None:
user = legacy
logger.info("Matched username %r to legacy account %s", body.username, legacy.email)
if user is None:
try:
role = "admin" if body.username == "admin" else "user"
@@ -380,6 +407,8 @@ async def login_by_username(request: Request, response: Response, body: Username
except ValueError:
# Concurrent first login may create the same user in parallel.
user = await provider.get_user_by_email(email)
if user is None and local_part is not None:
user = await provider.find_user_by_email_local_part(local_part)
if user is None:
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Failed to create or load user")
elif body.username == "admin" and user.system_role != "admin":
@@ -0,0 +1,323 @@
"""Admin-only LLM call metrics: list + Excel export.
Lets administrators audit every provider invocation made by the runtime —
including failures — and download the full table as ``.xlsx``. The recording
is unconditional on the backend, so this view always reflects every call.
"""
from __future__ import annotations
import logging
from datetime import datetime
from io import BytesIO
from typing import Any
from urllib.parse import quote
from fastapi import APIRouter, HTTPException, Query, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from app.gateway.deps import get_llm_metrics_store, get_local_provider, get_optional_user_from_request
from deerflow.persistence.llm_metrics.base import LlmMetricsQuery, LlmMetricsStore
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/admin/llm-metrics", tags=["llm-metrics"])
class LlmMetricResponse(BaseModel):
id: str
created_at: str | None = None
user_id: str | None = None
user_name: str | None = Field(
default=None,
description="Local part (before @) of the user's email, resolved from user_id",
)
thread_id: str | None = None
run_id: str | None = None
agent_name: str | None = None
model_name: str | None = None
duration_ms: int = 0
input_tokens: int | None = None
output_tokens: int | None = None
total_tokens: int | None = None
tokens_per_sec: float | None = None
status: str = "success"
error_type: str | None = None
error_message: str | None = None
class LlmMetricsListResponse(BaseModel):
items: list[LlmMetricResponse]
total: int = Field(..., description="Total rows matching the filter, ignoring pagination")
limit: int
offset: int
async def _require_admin(request: Request) -> None:
user = await get_optional_user_from_request(request)
if user is None:
return
if getattr(user, "system_role", None) != "admin":
raise HTTPException(status_code=403, detail="LLM metrics admin view is admin-only")
def _require_store(request: Request) -> LlmMetricsStore:
store = get_llm_metrics_store(request)
if store is None:
raise HTTPException(status_code=503, detail="LLM metrics store is not configured")
return store
def _parse_dt(value: str | None) -> datetime | None:
if not value:
return None
try:
# Accept both ``Z`` suffix and ``+00:00``.
return datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError as exc:
raise HTTPException(status_code=422, detail=f"Invalid datetime: {value!r} ({exc})") from exc
def _email_local_part(email: str | None) -> str | None:
if not email:
return None
at = email.find("@")
return email[:at] if at > 0 else email
def _derive_tokens_per_sec(row: dict[str, Any]) -> None:
"""Fill in ``tokens_per_sec`` when the row was recorded before we knew how
to extract ``output_tokens`` from response_metadata.token_usage.
Uses ``output_tokens / elapsed`` when available, otherwise falls back to
``total_tokens / elapsed`` as an approximation. Newer rows already have
the field set during recording, so this only kicks in for legacy data.
"""
if not isinstance(row, dict):
return
if row.get("tokens_per_sec") not in (None, 0):
return
duration_ms = row.get("duration_ms") or 0
if duration_ms <= 0:
return
numerator = row.get("output_tokens") or row.get("total_tokens")
if not numerator:
return
row["tokens_per_sec"] = round((numerator / duration_ms) * 1000, 3)
async def _resolve_user_name_to_id(name: str | None) -> str | None:
"""Look up a UUID for the given username (email local-part)."""
if not name:
return None
try:
provider = get_local_provider()
except Exception:
return None
try:
user = await provider.find_user_by_email_local_part(name)
except Exception:
return None
return str(user.id) if user is not None else None
async def _attach_usernames(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Resolve each row's ``user_id`` UUID to a friendly ``user_name``.
Batches lookups so one query per distinct user_id, not per row. Tolerates
missing accounts (deleted users) by leaving ``user_name`` blank.
"""
if not rows:
return rows
user_ids = {row["user_id"] for row in rows if isinstance(row, dict) and row.get("user_id")}
if not user_ids:
return rows
try:
provider = get_local_provider()
except Exception:
provider = None
email_by_id: dict[str, str] = {}
if provider is not None:
for uid in user_ids:
try:
user = await provider.get_user(uid)
except Exception:
user = None
if user is not None and user.email:
email_by_id[uid] = user.email
for row in rows:
if isinstance(row, dict) and row.get("user_id"):
email = email_by_id.get(row["user_id"])
row["user_name"] = _email_local_part(email) or row["user_id"]
return rows
def _build_query(
*,
user_id: str | None,
model_name: str | None,
status: str | None,
since: str | None,
until: str | None,
limit: int,
offset: int,
) -> LlmMetricsQuery:
return LlmMetricsQuery(
user_id=user_id or None,
model_name=model_name or None,
status=status or None,
since=_parse_dt(since),
until=_parse_dt(until),
limit=limit,
offset=offset,
)
@router.get("", response_model=LlmMetricsListResponse)
async def list_llm_metrics(
request: Request,
user_id: str | None = Query(default=None),
user_name: str | None = Query(default=None, description="Username (email local part) to filter by"),
model_name: str | None = Query(default=None),
status: str | None = Query(default=None, pattern="^(success|error)$"),
since: str | None = Query(default=None, description="ISO 8601 datetime, inclusive lower bound"),
until: str | None = Query(default=None, description="ISO 8601 datetime, inclusive upper bound"),
limit: int = Query(default=50, ge=1, le=500),
offset: int = Query(default=0, ge=0),
) -> LlmMetricsListResponse:
await _require_admin(request)
store = _require_store(request)
if not user_id and user_name:
resolved = await _resolve_user_name_to_id(user_name)
if resolved is None:
# Username unknown → no rows; short-circuit instead of a global scan.
return LlmMetricsListResponse(items=[], total=0, limit=limit, offset=offset)
user_id = resolved
query = _build_query(
user_id=user_id,
model_name=model_name,
status=status,
since=since,
until=until,
limit=limit,
offset=offset,
)
rows = await store.list(query)
total = await store.count(query)
await _attach_usernames(rows)
for row in rows:
_derive_tokens_per_sec(row)
return LlmMetricsListResponse(
items=[LlmMetricResponse(**row) for row in rows],
total=total,
limit=limit,
offset=offset,
)
def _format_cell(value: Any) -> Any:
if value is None:
return ""
if isinstance(value, datetime):
return value.replace(tzinfo=None) if value.tzinfo else value
return value
@router.get("/export")
async def export_llm_metrics(
request: Request,
user_id: str | None = Query(default=None),
user_name: str | None = Query(default=None),
model_name: str | None = Query(default=None),
status: str | None = Query(default=None, pattern="^(success|error)$"),
since: str | None = Query(default=None),
until: str | None = Query(default=None),
):
"""Stream the full filtered metric set as an ``.xlsx`` file."""
await _require_admin(request)
store = _require_store(request)
if not user_id and user_name:
resolved = await _resolve_user_name_to_id(user_name)
if resolved is None:
# No matching user — return an empty xlsx so the user still gets a
# file (matching the list view's empty state).
user_id = "__no_such_user__"
else:
user_id = resolved
try:
from openpyxl import Workbook
except ImportError as exc:
raise HTTPException(status_code=500, detail=f"openpyxl is not installed: {exc}") from exc
query = _build_query(
user_id=user_id,
model_name=model_name,
status=status,
since=since,
until=until,
limit=1, # unused: iter_all paginates internally
offset=0,
)
workbook = Workbook(write_only=True)
sheet = workbook.create_sheet("llm_call_metrics")
columns = [
"id",
"created_at",
"status",
"user_name",
"user_id",
"thread_id",
"run_id",
"agent_name",
"model_name",
"duration_ms",
"input_tokens",
"output_tokens",
"total_tokens",
"tokens_per_sec",
"error_type",
"error_message",
]
sheet.append(columns)
# Pre-resolve all distinct user_ids in chunks so the export doesn't pay one
# round trip per row.
buffered: list[dict[str, Any]] = []
FLUSH_AT = 200
async def flush() -> int:
if not buffered:
return 0
await _attach_usernames(buffered)
for row in buffered:
_derive_tokens_per_sec(row)
sheet.append([_format_cell(row.get(col)) for col in columns])
n = len(buffered)
buffered.clear()
return n
count = 0
async for row in store.iter_all(query):
buffered.append(row)
if len(buffered) >= FLUSH_AT:
count += await flush()
count += await flush()
buffer = BytesIO()
workbook.save(buffer)
buffer.seek(0)
logger.info("Exported %d llm_call_metrics rows", count)
timestamp = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
filename = f"llm_call_metrics-{timestamp}.xlsx"
headers = {
"Content-Disposition": f"attachment; filename=\"{filename}\"; filename*=UTF-8''{quote(filename)}",
}
return StreamingResponse(
buffer,
media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
headers=headers,
)
@@ -3,15 +3,31 @@ import logging
from pathlib import Path
from typing import Literal
from fastapi import APIRouter, HTTPException
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel, Field
from app.gateway.deps import get_optional_user_from_request
from deerflow.config.extensions_config import ExtensionsConfig, get_extensions_config, reload_extensions_config
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api", tags=["mcp"])
async def _require_admin(request: Request) -> None:
"""Reject the request unless the caller is an authenticated admin user.
Allows the call through silently when auth is disabled (no user attached),
so existing single-user / no-auth deployments keep working.
"""
user = await get_optional_user_from_request(request)
if user is None:
# No authenticated user — auth might be disabled or middleware bypassed.
# Defer to the AuthMiddleware decision and let the call through.
return
if getattr(user, "system_role", None) != "admin":
raise HTTPException(status_code=403, detail="MCP configuration requires admin privileges")
class McpOAuthConfigResponse(BaseModel):
"""OAuth configuration for an MCP server."""
@@ -43,6 +59,10 @@ class McpServerConfigResponse(BaseModel):
headers: dict[str, str] = Field(default_factory=dict, description="HTTP headers to send (for sse or http type)")
oauth: McpOAuthConfigResponse | None = Field(default=None, description="OAuth configuration for MCP HTTP/SSE servers")
description: str = Field(default="", description="Human-readable description of what this MCP server provides")
disabled_tools: list[str] = Field(
default_factory=list,
description="Tool names (without server-name prefix) to hide from the agent",
)
class McpConfigResponse(BaseModel):
@@ -101,7 +121,7 @@ async def get_mcp_configuration() -> McpConfigResponse:
summary="Update MCP Configuration",
description="Update Model Context Protocol (MCP) server configurations and save to file.",
)
async def update_mcp_configuration(request: McpConfigUpdateRequest) -> McpConfigResponse:
async def update_mcp_configuration(request: McpConfigUpdateRequest, http_request: Request) -> McpConfigResponse:
"""Update the MCP configuration.
This will:
@@ -133,6 +153,22 @@ async def update_mcp_configuration(request: McpConfigUpdateRequest) -> McpConfig
}
```
"""
await _require_admin(http_request)
# Light validation: each server must have the fields required by its transport.
for server_name, server in request.mcp_servers.items():
if not server_name.strip():
raise HTTPException(status_code=400, detail="MCP server name cannot be empty")
transport = (server.type or "stdio").lower()
if transport == "stdio":
if not (server.command or "").strip():
raise HTTPException(status_code=400, detail=f"MCP server '{server_name}' uses stdio transport but command is empty")
elif transport in ("http", "sse"):
if not (server.url or "").strip():
raise HTTPException(status_code=400, detail=f"MCP server '{server_name}' uses {transport} transport but url is empty")
else:
raise HTTPException(status_code=400, detail=f"MCP server '{server_name}' has unsupported transport: {transport}")
try:
# Get the current config path (or determine where to save it)
config_path = ExtensionsConfig.resolve_config_path()
@@ -14,7 +14,7 @@ from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import FileResponse, Response, StreamingResponse
from pydantic import BaseModel, Field
from app.gateway.deps import get_current_user
from app.gateway.deps import get_current_user, get_local_provider
from deerflow.runtime.scheduler import get_scheduled_task_service
router = APIRouter(prefix="/api/scheduled-tasks", tags=["scheduled-tasks"])
@@ -61,6 +61,50 @@ async def _user_id(request: Request) -> str:
return user_id
async def _enrich_with_creator(tasks: list[dict[str, Any]] | dict[str, Any] | None) -> Any:
"""Attach ``creator_email`` / ``creator_user_id`` to one or many task dicts.
Looks up the auth user once per distinct ``user_id`` in the batch so we do
not pay N+1 round trips when rendering a long task list. Tolerates missing
user records (deleted accounts) by leaving the email blank.
"""
if tasks is None:
return tasks
items = tasks if isinstance(tasks, list) else [tasks]
user_ids = {str(t.get("user_id")) for t in items if isinstance(t, dict) and t.get("user_id")}
if not user_ids:
return tasks
try:
provider = get_local_provider()
except Exception:
# Auth provider not initialised (e.g. memory-mode test harness): keep
# the bare user_id and let the frontend render the raw id.
provider = None
email_by_id: dict[str, str] = {}
if provider is not None:
for uid in user_ids:
try:
user = await provider.get_user(uid)
except Exception:
user = None
if user is not None:
email_by_id[uid] = user.email
for item in items:
if not isinstance(item, dict):
continue
uid = item.get("user_id")
if uid is None:
continue
uid_str = str(uid)
item["creator_user_id"] = uid_str
item["creator_email"] = email_by_id.get(uid_str, "")
return tasks
def _safe_filename(value: str, default: str = "scheduled-task-result") -> str:
name = re.sub(r'[\\/:*?"<>|\r\n]+', "-", value).strip(" .")
return name or default
@@ -133,7 +177,9 @@ async def get_scheduler_thread(request: Request) -> dict:
async def list_scheduled_tasks(request: Request) -> list[dict]:
service = get_scheduled_task_service()
user_id = await _user_id(request)
return await service.list_tasks(user_id)
tasks = await service.list_tasks(user_id)
await _enrich_with_creator(tasks)
return tasks
@router.get("/{task_id}/runs")
@@ -158,7 +204,10 @@ async def get_scheduled_task_run(run_id: str, request: Request) -> dict:
async def delete_scheduled_task_run(run_id: str, request: Request) -> dict:
service = get_scheduled_task_service()
user_id = await _user_id(request)
deleted = await service.delete_run(user_id, run_id)
try:
deleted = await service.delete_run(user_id, run_id)
except PermissionError as exc:
raise HTTPException(status_code=403, detail=str(exc))
if not deleted:
raise HTTPException(status_code=404, detail=f"Run {run_id} not found")
return {"success": True}
@@ -298,7 +347,7 @@ async def update_delivery_profile(body: ScheduledTaskDeliveryProfileUpdateReques
async def create_scheduled_task(body: ScheduledTaskCreateRequest, request: Request) -> dict:
service = get_scheduled_task_service()
user_id = await _user_id(request)
return await service.create_task(
task = await service.create_task(
user_id=user_id,
name=body.name,
prompt=body.prompt,
@@ -307,6 +356,8 @@ async def create_scheduled_task(body: ScheduledTaskCreateRequest, request: Reque
execution_agent_name=body.execution_agent_name,
execution_context=body.execution_context,
)
await _enrich_with_creator(task)
return task
@router.patch("/{task_id}")
@@ -316,6 +367,7 @@ async def update_scheduled_task(task_id: str, body: ScheduledTaskUpdateRequest,
task = await service.update_task(user_id, task_id, body.model_dump(exclude_unset=True))
if task is None:
raise HTTPException(status_code=404, detail=f"Task {task_id} not found")
await _enrich_with_creator(task)
return task
@@ -366,7 +418,17 @@ async def trigger_scheduled_task(task_id: str, request: Request) -> dict:
async def list_published_scheduled_tasks(request: Request, limit: int = 50, offset: int = 0) -> list[dict]:
service = get_scheduled_task_service()
await _user_id(request)
return await service.list_published(limit=max(1, min(limit, 200)), offset=max(0, offset))
tasks = await service.list_published(limit=max(1, min(limit, 200)), offset=max(0, offset))
# published list uses owner_user_id; resolve it into a creator_email here too
# so the discovery view can show the author's account.
for task in tasks:
if isinstance(task, dict) and task.get("owner_user_id"):
task["user_id"] = task["owner_user_id"]
await _enrich_with_creator(tasks)
for task in tasks:
if isinstance(task, dict):
task.pop("user_id", None)
return tasks
@router.post("/{task_id}/publish")
@@ -816,6 +816,36 @@ async def get_skill(skill_name: str, config: AppConfig = Depends(get_config)) ->
raise HTTPException(status_code=500, detail=f"Failed to get skill: {str(e)}")
class SkillContentResponse(BaseModel):
"""SKILL.md 原文内容,适用于公共和自定义技能。"""
name: str
category: SkillCategory
content: str
@router.get(
"/skills/{skill_name}/content",
response_model=SkillContentResponse,
summary="Get Skill Markdown Content",
description="Retrieve the raw SKILL.md content of a specific skill (public or custom).",
)
async def get_skill_content(skill_name: str, config: AppConfig = Depends(get_config)) -> SkillContentResponse:
try:
skill_name = skill_name.replace("\r\n", "").replace("\n", "")
skills = get_or_new_skill_storage(app_config=config).load_skills(enabled_only=False)
skill = next((s for s in skills if s.name == skill_name), None)
if skill is None:
raise HTTPException(status_code=404, detail=f"Skill '{skill_name}' not found")
content = skill.skill_file.read_text(encoding="utf-8")
return SkillContentResponse(name=skill.name, category=skill.category, content=content)
except HTTPException:
raise
except Exception as e:
logger.error("Failed to get skill content %s: %s", skill_name, e, exc_info=True)
raise HTTPException(status_code=500, detail=f"Failed to get skill content: {str(e)}")
@router.put(
"/skills/{skill_name}",
response_model=SkillResponse,
@@ -0,0 +1,65 @@
"""Admin-configurable system-wide runtime settings."""
from __future__ import annotations
import logging
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel, Field
from app.gateway.deps import get_optional_user_from_request
from deerflow.config.system_settings import (
PromptPrefixSettings,
SystemSettings,
load_system_settings,
save_system_settings,
)
logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/system-settings", tags=["system-settings"])
class PromptPrefixResponse(BaseModel):
enabled: bool = Field(..., description="Whether the toggle is exposed to users")
switch_label: str = Field(..., description="Toggle label shown in the chat input")
prompt_prefix: str = Field(..., description="Prefix text prepended to the first user message")
class PromptPrefixUpdateRequest(BaseModel):
enabled: bool | None = None
switch_label: str | None = None
prompt_prefix: str | None = None
def _to_response(settings: PromptPrefixSettings) -> PromptPrefixResponse:
return PromptPrefixResponse(
enabled=settings.enabled,
switch_label=settings.switch_label,
prompt_prefix=settings.prompt_prefix,
)
async def _require_admin(request: Request) -> None:
user = await get_optional_user_from_request(request)
if user is None:
return
if getattr(user, "system_role", None) != "admin":
raise HTTPException(status_code=403, detail="System settings can only be modified by an admin")
@router.get("/prompt-prefix", response_model=PromptPrefixResponse)
async def get_prompt_prefix() -> PromptPrefixResponse:
"""Return the current prompt-prefix configuration (readable by any user)."""
return _to_response(load_system_settings().prompt_prefix)
@router.put("/prompt-prefix", response_model=PromptPrefixResponse)
async def update_prompt_prefix(body: PromptPrefixUpdateRequest, request: Request) -> PromptPrefixResponse:
"""Update the prompt-prefix configuration. Admin only."""
await _require_admin(request)
current = load_system_settings()
fields = body.model_dump(exclude_unset=True)
updated_prefix = current.prompt_prefix.model_copy(update=fields)
saved = save_system_settings(SystemSettings(prompt_prefix=updated_prefix))
return _to_response(saved.prompt_prefix)
@@ -19,6 +19,7 @@ from langchain_core.messages import HumanMessage
from app.gateway.deps import get_run_context, get_run_manager, get_stream_bridge
from app.gateway.utils import sanitize_log_param
from deerflow.config.system_settings import load_system_settings
from deerflow.runtime import (
END_SENTINEL,
HEARTBEAT_SENTINEL,
@@ -98,6 +99,45 @@ def normalize_input(raw_input: dict[str, Any] | None) -> dict[str, Any]:
_DEFAULT_ASSISTANT_ID = "lead_agent"
def _prepend_prompt_prefix(graph_input: dict[str, Any], prefix: str) -> dict[str, Any]:
"""Return ``graph_input`` with the prefix prepended to its first HumanMessage.
Leaves the input unchanged when there are no messages or the first message
is not a human turn. Idempotent in the sense that it only touches the first
human message, so re-running this with the same prefix would double-prepend
— callers must gate it on first-turn detection.
"""
if not prefix:
return graph_input
messages = graph_input.get("messages")
if not isinstance(messages, list) or not messages:
return graph_input
new_messages = list(messages)
for idx, msg in enumerate(new_messages):
if isinstance(msg, HumanMessage):
existing = msg.content
if isinstance(existing, str):
new_messages[idx] = HumanMessage(content=f"{prefix}\n\n{existing}".strip("\n"))
elif isinstance(existing, list):
# Structured content (text blocks + media). Prepend a new text block.
new_messages[idx] = HumanMessage(content=[{"type": "text", "text": prefix}, *existing])
else:
new_messages[idx] = HumanMessage(content=f"{prefix}\n\n{existing!s}")
break
return {**graph_input, "messages": new_messages}
def _prompt_prefix_opted_in(body: Any) -> bool:
"""Detect whether the caller asked for the prompt prefix to be applied."""
context = getattr(body, "context", None)
if isinstance(context, Mapping) and context.get("prompt_prefix_enabled"):
return True
metadata = getattr(body, "metadata", None)
if isinstance(metadata, Mapping) and metadata.get("prompt_prefix_enabled"):
return True
return False
# Whitelist of run-context keys that the langgraph-compat layer forwards from
# ``body.context`` into the run config. ``config["context"]`` exists in
# LangGraph >=0.6, but these values must be written to both ``configurable``
@@ -340,6 +380,7 @@ async def start_run(
# Upsert thread metadata so the thread appears in /threads/search,
# even for threads that were never explicitly created via POST /threads
# (e.g. stateless runs).
thread_existed = False
try:
existing = await run_ctx.thread_store.get(thread_id)
if existing is None:
@@ -349,12 +390,24 @@ async def start_run(
metadata=body.metadata,
)
else:
thread_existed = True
await run_ctx.thread_store.update_status(thread_id, "running")
except Exception:
logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id))
agent_factory = resolve_agent_factory(body.assistant_id)
graph_input = normalize_input(body.input)
# Server-side safety net for the admin-configurable prompt prefix: when the
# caller opted in via context.prompt_prefix_enabled AND this is the very
# first turn of the thread, prepend the configured prefix to the human
# message. We deliberately skip already-running threads so the prefix is
# never inserted twice — the frontend toggle is a per-message hint, but
# this check is what makes the feature safe against client bugs / replays.
if not thread_existed and _prompt_prefix_opted_in(body):
prefix_settings = load_system_settings().prompt_prefix
if prefix_settings.enabled and prefix_settings.prompt_prefix.strip():
graph_input = _prepend_prompt_prefix(graph_input, prefix_settings.prompt_prefix.strip())
config = build_run_config(thread_id, body.config, body.metadata, assistant_id=body.assistant_id)
# Merge DeerFlow-specific context overrides into both ``configurable`` and ``context``.
@@ -0,0 +1,335 @@
"""Middleware that records per-LLM-call metrics for the admin audit page.
Captures **every** provider invocation regardless of the frontend speed-meter
toggle: the goal is a complete admin-side audit log of LLM usage, including
failures. The user-facing speed indicator (controlled by
``LocalSettings.display.streamSpeedMeterEnabled``) is a separate, purely
cosmetic feature — turning it off does not stop this recording.
Each call ends with one row in ``llm_call_metrics``:
- Successes record duration, token counts, and ``tokens_per_sec``.
- Errors record the duration spent before the failure, the exception class as
``error_type`` and a one-line ``error_message``.
Writes go through a best-effort store: a DB outage logs a warning and the
agent runtime keeps streaming. We intentionally place this middleware
*inside* :class:`LLMErrorHandlingMiddleware` (i.e. nested deeper in the
``wrap_model_call`` chain) so we observe the **real** provider call timings —
not the retry-wrapped successful response.
"""
from __future__ import annotations
import asyncio
import logging
import time
import uuid
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from typing import Any, override
from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from langchain.agents.middleware.types import (
ModelCallResult,
ModelRequest,
ModelResponse,
)
from langchain_core.messages import AIMessage
from langgraph.errors import GraphBubbleUp
from deerflow.persistence.engine import get_session_factory
from deerflow.persistence.llm_metrics import LlmMetricRecord, make_llm_metrics_store
from deerflow.persistence.llm_metrics.base import LlmMetricsStore
from deerflow.runtime.user_context import get_effective_user_id
logger = logging.getLogger(__name__)
def _now() -> datetime:
return datetime.now(UTC)
def _safe_get_thread_id(request: ModelRequest) -> str | None:
state = getattr(request, "state", None)
if isinstance(state, dict):
thread_id = state.get("thread_id")
if isinstance(thread_id, str) and thread_id:
return thread_id
runtime = getattr(request, "runtime", None)
if runtime is None:
return None
config = getattr(runtime, "config", None)
configurable = (config or {}).get("configurable", {}) if isinstance(config, dict) else {}
thread_id = configurable.get("thread_id") if isinstance(configurable, dict) else None
if isinstance(thread_id, str) and thread_id:
return thread_id
return None
def _safe_get_run_id(request: ModelRequest) -> str | None:
runtime = getattr(request, "runtime", None)
if runtime is None:
return None
context = getattr(runtime, "context", None)
if isinstance(context, dict):
for key in ("run_id", "langgraph_run_id"):
value = context.get(key)
if isinstance(value, str) and value:
return value
config = getattr(runtime, "config", None)
if isinstance(config, dict):
metadata = config.get("metadata") if isinstance(config.get("metadata"), dict) else {}
if isinstance(metadata, dict):
for key in ("run_id", "langgraph_run_id"):
value = metadata.get(key)
if isinstance(value, str) and value:
return value
return None
def _safe_get_model_name(request: ModelRequest) -> str | None:
model = getattr(request, "model", None)
for attr in ("model_name", "model", "name"):
value = getattr(model, attr, None)
if isinstance(value, str) and value:
return value
return None
def _safe_get_agent_name(request: ModelRequest) -> str | None:
runtime = getattr(request, "runtime", None)
if runtime is None:
return None
context = getattr(runtime, "context", None)
if isinstance(context, dict):
for key in ("agent_name", "agent_id"):
value = context.get(key)
if isinstance(value, str) and value:
return value
return None
def _coerce_int(value: Any) -> int | None:
if isinstance(value, bool):
return None
if isinstance(value, int):
return value
if isinstance(value, float):
return int(value)
if isinstance(value, str):
try:
return int(value)
except (TypeError, ValueError):
return None
return None
def _extract_usage(response: Any) -> dict[str, int]:
"""Pull token counts out of an AIMessage, falling back to provider-native shapes.
The handler in ``wrap_model_call`` actually hands us a
:class:`langchain.agents.middleware.types.ModelResponse` dataclass whose
``.result`` is a ``list[BaseMessage]``. Earlier versions of this helper
only matched ``AIMessage`` and ``list``, which meant we silently dropped
every successful call and stored None for input/output/total. To stay
robust we now also handle ``ModelResponse`` (and ``ExtendedModelResponse``,
which wraps a ModelResponse), and we walk the messages list for the first
``AIMessage``.
LangChain normalises usage into ``message.usage_metadata`` with the keys
``input_tokens`` / ``output_tokens`` / ``total_tokens``. Several
OpenAI-compatible providers (DeepSeek, vLLM, some Azure deployments) leave
that dict partial — typically only ``total_tokens`` is populated — and the
real breakdown sits under ``response_metadata.token_usage`` with the
provider-native key names ``prompt_tokens`` / ``completion_tokens``. We
merge both so the admin audit view always has the most complete numbers
available.
"""
message: AIMessage | None = None
def _find_ai(candidate: Any) -> AIMessage | None:
if isinstance(candidate, AIMessage):
return candidate
if isinstance(candidate, list):
for item in candidate:
found = _find_ai(item)
if found is not None:
return found
return None
message = _find_ai(response)
if message is None:
# ModelResponse / ExtendedModelResponse — defined as dataclasses with a
# ``result`` field holding ``list[BaseMessage]``. Recurse via attribute
# access so we don't tightly couple to the dataclass class itself.
for attr in ("result", "model_response"):
inner = getattr(response, attr, None)
if inner is None:
continue
message = _find_ai(inner)
if message is None and attr == "model_response":
inner_result = getattr(inner, "result", None)
message = _find_ai(inner_result)
if message is not None:
break
if message is None:
return {}
out: dict[str, int] = {}
usage = getattr(message, "usage_metadata", None) or {}
if isinstance(usage, dict):
for key in ("input_tokens", "output_tokens", "total_tokens"):
coerced = _coerce_int(usage.get(key))
if coerced is not None:
out[key] = coerced
response_metadata = getattr(message, "response_metadata", None) or {}
token_usage = response_metadata.get("token_usage") if isinstance(response_metadata, dict) else None
if isinstance(token_usage, dict):
# Map both the LangChain-canonical names and OpenAI-native names; later
# writes win so a more-precise provider-native value can override a
# zero placeholder from usage_metadata.
candidates = {
"input_tokens": ("input_tokens", "prompt_tokens"),
"output_tokens": ("output_tokens", "completion_tokens"),
"total_tokens": ("total_tokens",),
}
for canonical, keys in candidates.items():
for key in keys:
coerced = _coerce_int(token_usage.get(key))
if coerced is not None and coerced > 0:
out[canonical] = coerced
break
# Derive a missing or implausibly-zero field if the other two are present.
# Some providers report ``input_tokens=0`` / ``output_tokens=0`` next to a
# real ``total_tokens`` — treat those zeros as missing so we don't silently
# store nonsense breakdowns.
def _missing_or_zero(key: str) -> bool:
return out.get(key) in (None, 0)
if "total_tokens" not in out and out.get("input_tokens") and out.get("output_tokens"):
out["total_tokens"] = out["input_tokens"] + out["output_tokens"]
if _missing_or_zero("output_tokens") and out.get("total_tokens") and out.get("input_tokens"):
derived = out["total_tokens"] - out["input_tokens"]
if derived > 0:
out["output_tokens"] = derived
if _missing_or_zero("input_tokens") and out.get("total_tokens") and out.get("output_tokens"):
derived = out["total_tokens"] - out["output_tokens"]
if derived > 0:
out["input_tokens"] = derived
return out
def _classify_exception(exc: BaseException) -> tuple[str, str]:
"""Return (error_type, error_message). Message is trimmed to 500 chars."""
name = exc.__class__.__name__
raw = (str(exc) or name).strip()
return name, (raw[:500] if len(raw) > 500 else raw)
class LlmMetricsMiddleware(AgentMiddleware[AgentState]):
"""Wrap each model call to time it and persist a metric row."""
def __init__(self, *, store: LlmMetricsStore | None = None, **kwargs: Any) -> None:
super().__init__(**kwargs)
self._store_override = store
def _store(self) -> LlmMetricsStore | None:
if self._store_override is not None:
return self._store_override
try:
sf = get_session_factory()
except Exception:
return None
return make_llm_metrics_store(sf)
def _build_base_record(self, request: ModelRequest) -> LlmMetricRecord:
return LlmMetricRecord(
id=uuid.uuid4().hex,
created_at=_now(),
user_id=get_effective_user_id(),
thread_id=_safe_get_thread_id(request),
run_id=_safe_get_run_id(request),
agent_name=_safe_get_agent_name(request),
model_name=_safe_get_model_name(request),
)
def _finalize_success(self, record: LlmMetricRecord, response: ModelResponse, started: float) -> None:
elapsed = max(0.0, time.monotonic() - started)
record.duration_ms = int(elapsed * 1000)
usage = _extract_usage(response)
record.input_tokens = usage.get("input_tokens")
record.output_tokens = usage.get("output_tokens")
record.total_tokens = usage.get("total_tokens") or (
(record.input_tokens or 0) + (record.output_tokens or 0) or None
)
if record.output_tokens is not None and elapsed > 0:
record.tokens_per_sec = round(record.output_tokens / elapsed, 3)
record.status = "success"
def _finalize_error(self, record: LlmMetricRecord, exc: BaseException, started: float) -> None:
record.duration_ms = int(max(0.0, time.monotonic() - started) * 1000)
record.status = "error"
record.error_type, record.error_message = _classify_exception(exc)
def _flush(self, record: LlmMetricRecord) -> None:
store = self._store()
if store is None:
return
try:
try:
loop = asyncio.get_running_loop()
except RuntimeError:
loop = None
if loop is not None and loop.is_running():
# Fire-and-forget — never let metric IO block the request path.
loop.create_task(store.record(record))
else:
asyncio.run(store.record(record))
except Exception:
logger.warning("LlmMetricsMiddleware: failed to flush metric %s", record.id, exc_info=True)
@override
def wrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], ModelResponse],
) -> ModelCallResult:
record = self._build_base_record(request)
started = time.monotonic()
try:
response = handler(request)
except GraphBubbleUp:
# LangGraph control flow — not a real LLM failure.
raise
except Exception as exc:
self._finalize_error(record, exc, started)
self._flush(record)
raise
self._finalize_success(record, response, started)
self._flush(record)
return response
@override
async def awrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
) -> ModelCallResult:
record = self._build_base_record(request)
started = time.monotonic()
try:
response = await handler(request)
except GraphBubbleUp:
raise
except Exception as exc:
self._finalize_error(record, exc, started)
self._flush(record)
raise
self._finalize_success(record, response, started)
self._flush(record)
return response
@@ -96,6 +96,14 @@ def _build_runtime_middlewares(
middlewares.append(LLMErrorHandlingMiddleware(app_config=app_config))
# LLM call metrics — always on, even when the user-facing speed indicator
# is disabled. We append *after* the retry middleware so the metrics
# middleware sits deeper in the wrap chain and observes individual provider
# attempts (including the failed ones).
from deerflow.agents.middlewares.llm_metrics_middleware import LlmMetricsMiddleware
middlewares.append(LlmMetricsMiddleware())
# Guardrail middleware (if configured)
guardrails_config = app_config.guardrails
if guardrails_config.enabled and guardrails_config.provider:
@@ -45,7 +45,12 @@ class McpServerConfig(BaseModel):
headers: dict[str, str] = Field(default_factory=dict, description="HTTP headers to send (for sse or http type)")
oauth: McpOAuthConfig | None = Field(default=None, description="OAuth configuration (for sse or http type)")
description: str = Field(default="", description="Human-readable description of what this MCP server provides")
model_config = ConfigDict(extra="allow")
disabled_tools: list[str] = Field(
default_factory=list,
description="List of tool names (without server-name prefix) to hide from the agent",
alias="disabledTools",
)
model_config = ConfigDict(extra="allow", populate_by_name=True)
class SkillStateConfig(BaseModel):
@@ -0,0 +1,77 @@
"""System-wide runtime settings editable from the admin UI.
Backed by a single JSON file under the DeerFlow runtime home (default
``.deer-flow/system_settings.json``). Unlike :mod:`extensions_config` this
file is **not** part of the project source tree — it is per-deployment
state, and it stays writable at runtime.
The only feature stored here at the moment is the **prompt prefix toggle**
shown in the chat input: an administrator can configure a label and a
prefix string, and any user can flip the toggle on/off; when enabled, the
prefix is prepended to the first human message of a new thread.
"""
from __future__ import annotations
import json
import logging
import threading
from pathlib import Path
from pydantic import BaseModel, ConfigDict, Field
from deerflow.config.runtime_paths import runtime_home
logger = logging.getLogger(__name__)
_SETTINGS_FILE = "system_settings.json"
_lock = threading.Lock()
class PromptPrefixSettings(BaseModel):
"""Admin-configurable first-message prompt prefix."""
enabled: bool = Field(default=False, description="Whether the prefix toggle is offered to users at all")
switch_label: str = Field(default="加入立场前缀", description="Toggle label shown in the chat input")
prompt_prefix: str = Field(default="", description="Text prepended to the first user message when enabled")
model_config = ConfigDict(extra="ignore")
class SystemSettings(BaseModel):
"""Top-level system settings document."""
prompt_prefix: PromptPrefixSettings = Field(default_factory=PromptPrefixSettings)
model_config = ConfigDict(extra="ignore")
def _settings_path() -> Path:
home = runtime_home()
home.mkdir(parents=True, exist_ok=True)
return home / _SETTINGS_FILE
def load_system_settings() -> SystemSettings:
"""Read the on-disk system settings, returning defaults when missing."""
path = _settings_path()
if not path.is_file():
return SystemSettings()
try:
raw = json.loads(path.read_text(encoding="utf-8"))
if not isinstance(raw, dict):
logger.warning("System settings file %s is not a JSON object; ignoring", path)
return SystemSettings()
return SystemSettings.model_validate(raw)
except Exception as exc:
logger.warning("Failed to read system settings from %s: %s", path, exc)
return SystemSettings()
def save_system_settings(settings: SystemSettings) -> SystemSettings:
"""Persist the settings document atomically, returning the saved value."""
path = _settings_path()
payload = settings.model_dump(mode="json")
with _lock:
tmp = path.with_suffix(path.suffix + ".tmp")
tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
tmp.replace(path)
return settings
@@ -121,6 +121,25 @@ async def get_mcp_tools() -> list[BaseTool]:
# Get all tools from all servers
tools = await client.get_tools()
# Filter out tools listed in each server's `disabled_tools`. Tool names from
# MultiServerMCPClient are prefixed with the server name (because we pass
# `tool_name_prefix=True`), so we compare against `<server>_<tool>` here.
disabled_full_names: set[str] = set()
for server_name, server_cfg in extensions_config.get_enabled_mcp_servers().items():
for tool_name in server_cfg.disabled_tools or []:
name = (tool_name or "").strip()
if not name:
continue
disabled_full_names.add(f"{server_name}_{name}")
disabled_full_names.add(name)
if disabled_full_names:
kept = [t for t in tools if t.name not in disabled_full_names]
removed = len(tools) - len(kept)
if removed:
logger.info(f"Filtered out {removed} disabled MCP tool(s)")
tools = kept
logger.info(f"Successfully loaded {len(tools)} tool(s) from MCP servers")
# Patch tools to support sync invocation, as deerflow client streams synchronously
@@ -0,0 +1,25 @@
"""Persistence layer for per-LLM-call metrics.
Each completed LLM invocation (success or failure) writes a single row so
administrators can audit usage and speed. Designed to be cheap to write —
the middleware fires once per provider call and never blocks streaming.
"""
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsStore
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
from deerflow.persistence.llm_metrics.sql import SqlLlmMetricsStore
__all__ = [
"LlmCallMetricRow",
"LlmMetricRecord",
"LlmMetricsStore",
"SqlLlmMetricsStore",
"make_llm_metrics_store",
]
def make_llm_metrics_store(session_factory) -> LlmMetricsStore | None:
"""Return a SQL-backed metrics store, or ``None`` when no DB is configured."""
if session_factory is None:
return None
return SqlLlmMetricsStore(session_factory)
@@ -0,0 +1,66 @@
"""Abstract repository contract for LLM call metrics."""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from typing import Any
@dataclass
class LlmMetricRecord:
"""In-memory shape of a metric row before persistence."""
id: str
created_at: datetime
user_id: str | None = None
thread_id: str | None = None
run_id: str | None = None
agent_name: str | None = None
model_name: str | None = None
duration_ms: int = 0
input_tokens: int | None = None
output_tokens: int | None = None
total_tokens: int | None = None
tokens_per_sec: float | None = None
status: str = "success"
error_type: str | None = None
error_message: str | None = None
@dataclass
class LlmMetricsQuery:
"""Filter set for the admin list / export endpoints."""
user_id: str | None = None
model_name: str | None = None
status: str | None = None
since: datetime | None = None
until: datetime | None = None
limit: int = 100
offset: int = 0
class LlmMetricsStore(ABC):
"""Persistence contract for the metrics middleware + admin API."""
@abstractmethod
async def record(self, metric: LlmMetricRecord) -> None:
"""Persist one metric row. Must never raise on best-effort writes."""
raise NotImplementedError
@abstractmethod
async def list(self, query: LlmMetricsQuery) -> list[dict[str, Any]]:
"""Return the matching rows (newest first), paginated by ``query``."""
raise NotImplementedError
@abstractmethod
async def count(self, query: LlmMetricsQuery) -> int:
"""Return the unpaginated row count for ``query``."""
raise NotImplementedError
@abstractmethod
async def iter_all(self, query: LlmMetricsQuery):
"""Yield rows matching ``query`` without pagination — used by export."""
raise NotImplementedError
@@ -0,0 +1,62 @@
"""ORM model for per-LLM-call metrics."""
from __future__ import annotations
from datetime import UTC, datetime
from sqlalchemy import DateTime, Float, Integer, String, Text
from sqlalchemy.orm import Mapped, mapped_column
from deerflow.persistence.base import Base
class LlmCallMetricRow(Base):
"""One row per completed (or failed) LLM provider invocation."""
__tablename__ = "llm_call_metrics"
id: Mapped[str] = mapped_column(String(64), primary_key=True)
# Wall-clock time the LLM call ended. Indexed because admin views always
# sort/filter by recency.
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=lambda: datetime.now(UTC), index=True)
user_id: Mapped[str | None] = mapped_column(String(64), index=True)
thread_id: Mapped[str | None] = mapped_column(String(64), index=True)
run_id: Mapped[str | None] = mapped_column(String(64), index=True)
agent_name: Mapped[str | None] = mapped_column(String(128), index=True)
model_name: Mapped[str | None] = mapped_column(String(128), index=True)
# Wall-clock duration of the provider invocation in milliseconds. Always set.
duration_ms: Mapped[int] = mapped_column(Integer, default=0)
input_tokens: Mapped[int | None] = mapped_column(Integer)
output_tokens: Mapped[int | None] = mapped_column(Integer)
total_tokens: Mapped[int | None] = mapped_column(Integer)
# Output tokens / seconds; null when there is no completion (errors).
tokens_per_sec: Mapped[float | None] = mapped_column(Float)
# "success" or "error". Errors store the classification and original message
# so admins can scan for transient vs auth/quota issues.
status: Mapped[str] = mapped_column(String(16), default="success", index=True)
error_type: Mapped[str | None] = mapped_column(String(64))
error_message: Mapped[str | None] = mapped_column(Text)
def to_dict(self) -> dict:
return {
"id": self.id,
"created_at": self.created_at.isoformat() if isinstance(self.created_at, datetime) else self.created_at,
"user_id": self.user_id,
"thread_id": self.thread_id,
"run_id": self.run_id,
"agent_name": self.agent_name,
"model_name": self.model_name,
"duration_ms": self.duration_ms,
"input_tokens": self.input_tokens,
"output_tokens": self.output_tokens,
"total_tokens": self.total_tokens,
"tokens_per_sec": self.tokens_per_sec,
"status": self.status,
"error_type": self.error_type,
"error_message": self.error_message,
}
@@ -0,0 +1,109 @@
"""SQLAlchemy-backed LLM metrics repository."""
from __future__ import annotations
import logging
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from typing import Any
from sqlalchemy import desc, func, select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery, LlmMetricsStore
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
logger = logging.getLogger(__name__)
def _row_to_dict(row: LlmCallMetricRow) -> dict[str, Any]:
return row.to_dict()
def _normalize_created_at(value: datetime) -> datetime:
return value if value.tzinfo else value.replace(tzinfo=UTC)
class SqlLlmMetricsStore(LlmMetricsStore):
"""SQLite/Postgres-backed implementation; writes are intentionally fire-and-forget."""
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
self._sf = session_factory
async def record(self, metric: LlmMetricRecord) -> None:
row = LlmCallMetricRow(
id=metric.id,
created_at=_normalize_created_at(metric.created_at),
user_id=metric.user_id,
thread_id=metric.thread_id,
run_id=metric.run_id,
agent_name=metric.agent_name,
model_name=metric.model_name,
duration_ms=metric.duration_ms,
input_tokens=metric.input_tokens,
output_tokens=metric.output_tokens,
total_tokens=metric.total_tokens,
tokens_per_sec=metric.tokens_per_sec,
status=metric.status,
error_type=metric.error_type,
error_message=metric.error_message,
)
try:
async with self._sf() as session:
session.add(row)
await session.commit()
except Exception:
# Metrics are best-effort. Never let a persistence failure surface
# to the agent runtime — log and move on.
logger.warning("Failed to persist llm_call_metrics row %s", metric.id, exc_info=True)
def _apply_filters(self, stmt, query: LlmMetricsQuery):
if query.user_id:
stmt = stmt.where(LlmCallMetricRow.user_id == query.user_id)
if query.model_name:
stmt = stmt.where(LlmCallMetricRow.model_name == query.model_name)
if query.status:
stmt = stmt.where(LlmCallMetricRow.status == query.status)
if query.since is not None:
stmt = stmt.where(LlmCallMetricRow.created_at >= _normalize_created_at(query.since))
if query.until is not None:
stmt = stmt.where(LlmCallMetricRow.created_at <= _normalize_created_at(query.until))
return stmt
async def list(self, query: LlmMetricsQuery) -> list[dict[str, Any]]:
stmt = select(LlmCallMetricRow)
stmt = self._apply_filters(stmt, query)
stmt = stmt.order_by(desc(LlmCallMetricRow.created_at)).limit(max(1, min(query.limit, 1000))).offset(max(0, query.offset))
async with self._sf() as session:
result = await session.execute(stmt)
return [_row_to_dict(row) for row in result.scalars()]
async def count(self, query: LlmMetricsQuery) -> int:
stmt = select(func.count()).select_from(LlmCallMetricRow)
stmt = self._apply_filters(stmt, query)
async with self._sf() as session:
return (await session.scalar(stmt)) or 0
async def iter_all(self, query: LlmMetricsQuery) -> AsyncIterator[dict[str, Any]]:
# Page through the result so the Excel export stays bounded in memory
# for large date windows.
page_size = 500
offset = 0
while True:
page_query = LlmMetricsQuery(
user_id=query.user_id,
model_name=query.model_name,
status=query.status,
since=query.since,
until=query.until,
limit=page_size,
offset=offset,
)
page = await self.list(page_query)
if not page:
return
for item in page:
yield item
if len(page) < page_size:
return
offset += page_size
@@ -0,0 +1,59 @@
"""Create llm_call_metrics table.
Revision ID: 20260514_02
Revises: 20260514_01
Create Date: 2026-05-14
"""
from __future__ import annotations
import sqlalchemy as sa
from alembic import op
from sqlalchemy import inspect
revision = "20260514_02"
down_revision = "20260514_01"
branch_labels = None
depends_on = None
def _has_table(table_name: str) -> bool:
inspector = inspect(op.get_bind())
return table_name in inspector.get_table_names()
def upgrade() -> None:
if _has_table("llm_call_metrics"):
return
op.create_table(
"llm_call_metrics",
sa.Column("id", sa.String(64), primary_key=True),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
sa.Column("user_id", sa.String(64), nullable=True),
sa.Column("thread_id", sa.String(64), nullable=True),
sa.Column("run_id", sa.String(64), nullable=True),
sa.Column("agent_name", sa.String(128), nullable=True),
sa.Column("model_name", sa.String(128), nullable=True),
sa.Column("duration_ms", sa.Integer(), nullable=False, server_default="0"),
sa.Column("input_tokens", sa.Integer(), nullable=True),
sa.Column("output_tokens", sa.Integer(), nullable=True),
sa.Column("total_tokens", sa.Integer(), nullable=True),
sa.Column("tokens_per_sec", sa.Float(), nullable=True),
sa.Column("status", sa.String(16), nullable=False, server_default="success"),
sa.Column("error_type", sa.String(64), nullable=True),
sa.Column("error_message", sa.Text(), nullable=True),
)
op.create_index("ix_llm_call_metrics_created_at", "llm_call_metrics", ["created_at"])
op.create_index("ix_llm_call_metrics_user_id", "llm_call_metrics", ["user_id"])
op.create_index("ix_llm_call_metrics_model_name", "llm_call_metrics", ["model_name"])
op.create_index("ix_llm_call_metrics_status", "llm_call_metrics", ["status"])
op.create_index("ix_llm_call_metrics_thread_id", "llm_call_metrics", ["thread_id"])
op.create_index("ix_llm_call_metrics_run_id", "llm_call_metrics", ["run_id"])
op.create_index("ix_llm_call_metrics_agent_name", "llm_call_metrics", ["agent_name"])
def downgrade() -> None:
if _has_table("llm_call_metrics"):
op.drop_table("llm_call_metrics")
@@ -17,8 +17,10 @@ there is no matching entity directory.
from deerflow.persistence.agents.model import AgentRow
from deerflow.persistence.feedback.model import FeedbackRow
from deerflow.persistence.llm_metrics.model import LlmCallMetricRow
from deerflow.persistence.models.run_event import RunEventRow
from deerflow.persistence.notifications.model import NotificationRow
from deerflow.persistence.recommended_questions.model import RecommendedQuestionRow
from deerflow.persistence.run.model import RunRow
from deerflow.persistence.scheduled_tasks.model import (
ScheduledTaskDeliveryProfileRow,
@@ -34,7 +36,9 @@ from deerflow.persistence.user.model import UserRow
__all__ = [
"FeedbackRow",
"AgentRow",
"LlmCallMetricRow",
"NotificationRow",
"RecommendedQuestionRow",
"RunEventRow",
"RunRow",
"ScheduledTaskDeliveryProfileRow",
@@ -247,7 +247,11 @@ class ScheduledTaskService:
accessible = await self._accessible_task(user_id, task_id)
if accessible is None:
return []
return await self.store.list_runs(task_id, accessible["user_id"], limit=limit)
runs = await self.store.list_runs(task_id, accessible["user_id"], limit=limit)
can_edit = accessible.get("role") == "owner"
for run in runs:
run["can_edit"] = can_edit
return runs
async def get_run(self, user_id: str, run_id: str) -> dict[str, Any] | None:
run = await self.store.get_run(run_id, user_id)
@@ -316,7 +320,7 @@ class ScheduledTaskService:
except ValueError as exc:
raise PermissionError("Refusing to edit file outside the run archive") from exc
archive_path.write_text(new_content, encoding="utf-8")
archive_path.write_text(new_content, encoding="utf-8", newline="")
new_size = archive_path.stat().st_size
# Reflect the new size in the run row so subsequent list/preview calls
@@ -342,6 +346,10 @@ class ScheduledTaskService:
shutil.rmtree(archive_dir)
except Exception:
logger.warning("Failed to delete scheduled task run archive %s", archive_dir, exc_info=True)
return True
candidate = await self.store.get_run_any(run_id)
if candidate is not None and await self.store.get_subscription(candidate["task_id"], user_id) is not None:
raise PermissionError("Only the run owner may delete scheduled task runs")
return deleted
async def get_delivery_profile(self, user_id: str) -> dict[str, Any]:
@@ -0,0 +1,321 @@
"""Tests for the LLM call metrics middleware + repository.
Covers:
- Successful calls emit a row with duration, token counts and tokens_per_sec.
- Failed calls record the exception class as ``error_type`` plus a trimmed
``error_message``.
- The middleware never blocks on a backend exception (best-effort writes).
- The SQL store applies filters and pagination for the admin list endpoint.
"""
from __future__ import annotations
import asyncio
import uuid
from datetime import UTC, datetime, timedelta
import pytest
import pytest_asyncio
from langchain_core.messages import AIMessage
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from deerflow.agents.middlewares.llm_metrics_middleware import (
LlmMetricsMiddleware,
_extract_usage,
)
from deerflow.persistence.base import Base
from deerflow.persistence.llm_metrics import SqlLlmMetricsStore
from deerflow.persistence.llm_metrics.base import LlmMetricRecord, LlmMetricsQuery
class _RecordingStore:
"""Minimal in-memory store used to assert what the middleware persists."""
def __init__(self) -> None:
self.records: list[LlmMetricRecord] = []
async def record(self, metric: LlmMetricRecord) -> None:
self.records.append(metric)
async def list(self, query: LlmMetricsQuery):
return [r.__dict__ for r in self.records]
async def count(self, query: LlmMetricsQuery) -> int:
return len(self.records)
async def iter_all(self, query: LlmMetricsQuery):
for record in self.records:
yield record.__dict__
@pytest_asyncio.fixture
async def sql_store(tmp_path):
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'metrics.db'}")
async with engine.begin() as conn:
# Importing the package registers the LlmCallMetricRow with Base.metadata.
import deerflow.persistence.llm_metrics.model # noqa: F401
await conn.run_sync(Base.metadata.create_all)
sf = async_sessionmaker(engine, expire_on_commit=False)
store = SqlLlmMetricsStore(sf)
try:
yield store
finally:
await engine.dispose()
def _make_request():
return type(
"FakeRequest",
(),
{
"model": type("FakeModel", (), {"model_name": "gpt-test"}),
"runtime": type(
"FakeRuntime",
(),
{
"config": {"configurable": {"thread_id": "t1"}},
"context": {"agent_name": "lead_agent", "run_id": "r1"},
},
),
"state": {"thread_id": "t1"},
},
)()
def test_middleware_records_success_with_speed():
import time as _time
store = _RecordingStore()
mw = LlmMetricsMiddleware(store=store)
response = AIMessage(
content="hi",
usage_metadata={"input_tokens": 5, "output_tokens": 20, "total_tokens": 25},
)
def slow_handler(_req):
# Force a measurable elapsed window so tokens_per_sec is populated.
# 50ms is comfortably above Windows monotonic-clock granularity.
_time.sleep(0.05)
return response
result = mw.wrap_model_call(_make_request(), slow_handler)
assert result is response
# awrap_model_call uses a fire-and-forget task; here the sync path runs
# asyncio.run synchronously, so the record landed already.
assert len(store.records) == 1
record = store.records[0]
assert record.status == "success"
assert record.model_name == "gpt-test"
assert record.thread_id == "t1"
assert record.run_id == "r1"
assert record.input_tokens == 5
assert record.output_tokens == 20
assert record.total_tokens == 25
# tokens_per_sec may be very high in tests because the handler is
# ~instant; we only assert that the field is populated.
assert record.tokens_per_sec is not None
assert record.duration_ms >= 0
def test_middleware_records_error_with_classification():
store = _RecordingStore()
mw = LlmMetricsMiddleware(store=store)
class BoomError(RuntimeError):
pass
def explode(_req):
raise BoomError("provider exploded")
with pytest.raises(BoomError):
mw.wrap_model_call(_make_request(), explode)
assert len(store.records) == 1
record = store.records[0]
assert record.status == "error"
assert record.error_type == "BoomError"
assert record.error_message == "provider exploded"
assert record.tokens_per_sec is None
def test_middleware_swallows_store_errors():
"""A persistence outage must not break the agent runtime."""
class FaultyStore:
async def record(self, metric):
raise RuntimeError("disk gone")
async def list(self, q):
return []
async def count(self, q):
return 0
async def iter_all(self, q):
if False:
yield {}
mw = LlmMetricsMiddleware(store=FaultyStore())
# The wrap returns normally even though _flush would raise; the call below
# must not raise.
response = AIMessage(
content="ok",
usage_metadata={"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
)
out = mw.wrap_model_call(_make_request(), lambda req: response)
assert out is response
@pytest.mark.asyncio
async def test_sql_store_round_trip(sql_store):
base = datetime.now(UTC)
for i in range(3):
await sql_store.record(
LlmMetricRecord(
id=uuid.uuid4().hex,
created_at=base - timedelta(minutes=i),
user_id="alice",
model_name="gpt-test" if i % 2 == 0 else "gpt-other",
duration_ms=100 * (i + 1),
output_tokens=10 * (i + 1),
tokens_per_sec=10.0 * (i + 1),
status="success",
)
)
rows = await sql_store.list(LlmMetricsQuery(limit=10))
assert len(rows) == 3
# Newest first.
assert rows[0]["duration_ms"] == 100
filtered = await sql_store.list(LlmMetricsQuery(model_name="gpt-test", limit=10))
assert all(r["model_name"] == "gpt-test" for r in filtered)
assert len(filtered) == 2
assert await sql_store.count(LlmMetricsQuery(model_name="gpt-test")) == 2
@pytest.mark.asyncio
async def test_sql_store_filters_errors_only(sql_store):
base = datetime.now(UTC)
await sql_store.record(
LlmMetricRecord(id=uuid.uuid4().hex, created_at=base, status="success", duration_ms=10, output_tokens=1)
)
await sql_store.record(
LlmMetricRecord(
id=uuid.uuid4().hex,
created_at=base,
status="error",
duration_ms=5,
error_type="APITimeoutError",
error_message="timed out",
)
)
only_errors = await sql_store.list(LlmMetricsQuery(status="error", limit=10))
assert len(only_errors) == 1
assert only_errors[0]["error_type"] == "APITimeoutError"
@pytest.mark.asyncio
async def test_sql_store_iter_all_paginates(sql_store):
base = datetime.now(UTC)
for i in range(1200):
await sql_store.record(
LlmMetricRecord(
id=uuid.uuid4().hex,
created_at=base - timedelta(seconds=i),
duration_ms=1,
status="success",
)
)
seen = 0
async for _ in sql_store.iter_all(LlmMetricsQuery()):
seen += 1
assert seen == 1200
def test_extract_usage_prefers_usage_metadata():
response = AIMessage(
content="x",
usage_metadata={"input_tokens": 3, "output_tokens": 7, "total_tokens": 10},
)
out = _extract_usage(response)
assert out == {"input_tokens": 3, "output_tokens": 7, "total_tokens": 10}
def test_extract_usage_falls_back_to_response_metadata_token_usage():
"""DeepSeek / vLLM-style: only total in usage_metadata, breakdown elsewhere."""
response = AIMessage(
content="x",
usage_metadata={"input_tokens": 0, "output_tokens": 0, "total_tokens": 11093},
response_metadata={
"token_usage": {
"prompt_tokens": 9000,
"completion_tokens": 2093,
"total_tokens": 11093,
},
},
)
out = _extract_usage(response)
assert out["input_tokens"] == 9000
assert out["output_tokens"] == 2093
assert out["total_tokens"] == 11093
def test_extract_usage_unwraps_model_response():
"""The handler hands us a ModelResponse dataclass, not a bare AIMessage."""
from langchain.agents.middleware.types import ModelResponse
inner = AIMessage(
content="hi",
usage_metadata={"input_tokens": 7, "output_tokens": 9, "total_tokens": 16},
)
response = ModelResponse(result=[inner])
out = _extract_usage(response)
assert out == {"input_tokens": 7, "output_tokens": 9, "total_tokens": 16}
def test_extract_usage_unwraps_extended_model_response():
"""ExtendedModelResponse wraps a ModelResponse in a ``model_response`` field."""
from langchain.agents.middleware.types import ExtendedModelResponse, ModelResponse
inner = AIMessage(
content="hi",
usage_metadata={"input_tokens": 1, "output_tokens": 2, "total_tokens": 3},
)
extended = ExtendedModelResponse(model_response=ModelResponse(result=[inner]))
out = _extract_usage(extended)
assert out["output_tokens"] == 2
def test_extract_usage_derives_missing_field_when_two_present():
response = AIMessage(
content="x",
usage_metadata={"input_tokens": 9000, "output_tokens": 0, "total_tokens": 11093},
)
out = _extract_usage(response)
# output should be derived as total - input = 2093.
assert out["output_tokens"] == 2093
@pytest.mark.asyncio
async def test_async_middleware_records_success():
store = _RecordingStore()
mw = LlmMetricsMiddleware(store=store)
async def handler(_req):
return AIMessage(
content="hi",
usage_metadata={"input_tokens": 1, "output_tokens": 4, "total_tokens": 5},
)
response = await mw.awrap_model_call(_make_request(), handler)
assert isinstance(response, AIMessage)
# awrap_model_call schedules the write as an asyncio.task; give the loop a
# chance to drain.
await asyncio.sleep(0)
assert len(store.records) == 1
assert store.records[0].status == "success"
assert store.records[0].output_tokens == 4
@@ -0,0 +1,108 @@
"""Tests for the admin-configurable prompt-prefix toggle."""
from __future__ import annotations
import json
from pathlib import Path
from types import SimpleNamespace
import pytest
from langchain_core.messages import AIMessage, HumanMessage
from app.gateway.services import _prepend_prompt_prefix, _prompt_prefix_opted_in
from deerflow.config import system_settings as ss
@pytest.fixture
def tmp_runtime_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path:
"""Redirect ``runtime_home()`` to a per-test tmp dir."""
def fake_home() -> Path:
return tmp_path
monkeypatch.setattr(ss, "runtime_home", fake_home)
return tmp_path
def test_load_returns_defaults_when_missing(tmp_runtime_home: Path) -> None:
settings = ss.load_system_settings()
assert settings.prompt_prefix.enabled is False
assert settings.prompt_prefix.prompt_prefix == ""
def test_save_then_load_roundtrips(tmp_runtime_home: Path) -> None:
settings = ss.SystemSettings(
prompt_prefix=ss.PromptPrefixSettings(
enabled=True,
switch_label="加入立场前缀",
prompt_prefix="你需要站在中方角度",
)
)
ss.save_system_settings(settings)
file_path = tmp_runtime_home / "system_settings.json"
assert file_path.is_file()
saved = json.loads(file_path.read_text(encoding="utf-8"))
assert saved["prompt_prefix"]["prompt_prefix"] == "你需要站在中方角度"
reloaded = ss.load_system_settings()
assert reloaded.prompt_prefix.enabled is True
assert reloaded.prompt_prefix.switch_label == "加入立场前缀"
assert reloaded.prompt_prefix.prompt_prefix == "你需要站在中方角度"
def test_prepend_prefix_on_string_content() -> None:
graph_input = {"messages": [HumanMessage(content="评价一下中美贸易")]}
out = _prepend_prompt_prefix(graph_input, "你需要站在中方角度")
msg = out["messages"][0]
assert isinstance(msg, HumanMessage)
assert msg.content == "你需要站在中方角度\n\n评价一下中美贸易"
def test_prepend_prefix_on_block_content() -> None:
graph_input = {
"messages": [
HumanMessage(content=[{"type": "text", "text": "评价一下中美贸易"}])
]
}
out = _prepend_prompt_prefix(graph_input, "你需要站在中方角度")
blocks = out["messages"][0].content
assert isinstance(blocks, list)
assert blocks[0] == {"type": "text", "text": "你需要站在中方角度"}
assert blocks[1] == {"type": "text", "text": "评价一下中美贸易"}
def test_prepend_prefix_only_touches_first_human() -> None:
graph_input = {
"messages": [
AIMessage(content="hi"),
HumanMessage(content="q1"),
HumanMessage(content="q2"),
]
}
out = _prepend_prompt_prefix(graph_input, "PREFIX")
# first human is index 1; index 2 must remain untouched
assert out["messages"][1].content == "PREFIX\n\nq1"
assert out["messages"][2].content == "q2"
def test_prepend_prefix_no_messages_is_noop() -> None:
assert _prepend_prompt_prefix({"messages": []}, "X") == {"messages": []}
assert _prepend_prompt_prefix({}, "X") == {}
def test_opted_in_via_context() -> None:
body = SimpleNamespace(context={"prompt_prefix_enabled": True}, metadata=None)
assert _prompt_prefix_opted_in(body) is True
def test_opted_in_via_metadata() -> None:
body = SimpleNamespace(context=None, metadata={"prompt_prefix_enabled": True})
assert _prompt_prefix_opted_in(body) is True
def test_not_opted_in_by_default() -> None:
body = SimpleNamespace(context=None, metadata=None)
assert _prompt_prefix_opted_in(body) is False
body2 = SimpleNamespace(context={"mode": "ultra"}, metadata={})
assert _prompt_prefix_opted_in(body2) is False
@@ -147,13 +147,23 @@ async def test_subscriber_can_read_runs_and_get_run(store):
)
assert run is not None
alice_runs = await service.list_runs("alice", "t1", limit=10)
assert alice_runs[0]["can_edit"] is True
# Bob (subscriber) sees runs on the task.
bob_runs = await service.list_runs("bob", "t1", limit=10)
assert [r["id"] for r in bob_runs] == ["run-1"]
assert bob_runs[0]["can_edit"] is False
# Bob can fetch the run by id.
fetched = await service.get_run("bob", "run-1")
assert fetched is not None and fetched["id"] == "run-1"
assert fetched["can_edit"] is False
# Bob can read subscribed records, but cannot delete Alice's record.
with pytest.raises(PermissionError):
await service.delete_run("bob", "run-1")
assert await store.get_run("run-1", "alice") is not None
# Carol (no subscription) cannot.
assert await service.list_runs("carol", "t1", limit=10) == []
@@ -0,0 +1,148 @@
"""Tests for the username-based login domain handling.
Verifies that:
- The default mail domain is ``cm.com`` (with env-var override still honoured).
- ``_normalize_username_local_part`` distinguishes raw-email vs username input.
- The login flow falls back to a legacy-domain lookup so the same username
keeps mapping to the same account after the default domain changes.
"""
from __future__ import annotations
import os
from datetime import UTC, datetime
from uuid import UUID, uuid4
import pytest
from app.gateway.auth.models import User
from app.gateway.auth.repositories.base import UserRepository
from app.gateway.routers.auth import (
_normalize_username_local_part,
_username_to_email,
)
class FakeUserRepository(UserRepository):
"""In-memory repo just enough to exercise the local-part lookup."""
def __init__(self) -> None:
self.rows: dict[str, User] = {}
self.create_calls: list[User] = []
async def create_user(self, user: User) -> User:
for existing in self.rows.values():
if existing.email == user.email:
raise ValueError(f"Email already registered: {user.email}")
u = User(
id=user.id or uuid4(),
email=user.email,
password_hash=user.password_hash,
system_role=user.system_role,
created_at=user.created_at or datetime.now(UTC),
oauth_provider=user.oauth_provider,
oauth_id=user.oauth_id,
needs_setup=user.needs_setup,
token_version=user.token_version,
)
self.rows[str(u.id)] = u
self.create_calls.append(u)
return u
async def get_user_by_id(self, user_id: str) -> User | None:
return self.rows.get(user_id)
async def get_user_by_email(self, email: str) -> User | None:
for user in self.rows.values():
if user.email == email:
return user
return None
async def find_user_by_email_local_part(self, local_part: str) -> User | None:
prefix = f"{(local_part or '').strip().lower()}@"
candidates = [u for u in self.rows.values() if u.email.lower().startswith(prefix)]
candidates.sort(key=lambda u: u.created_at)
return candidates[0] if candidates else None
async def update_user(self, user: User) -> User:
self.rows[str(user.id)] = user
return user
async def count_users(self) -> int:
return len(self.rows)
async def count_admin_users(self) -> int:
return sum(1 for u in self.rows.values() if u.system_role == "admin")
async def get_user_by_oauth(self, provider: str, oauth_id: str) -> User | None:
return None
def _seed(repo: FakeUserRepository, email: str, *, role: str = "user") -> User:
user = User(
id=uuid4(),
email=email,
password_hash=None,
system_role=role, # type: ignore[arg-type]
created_at=datetime.now(UTC),
needs_setup=False,
)
repo.rows[str(user.id)] = user
return user
def test_normalize_returns_none_for_explicit_email() -> None:
assert _normalize_username_local_part("alice@old.com") is None
def test_normalize_lowercases_and_keeps_safe_chars() -> None:
assert _normalize_username_local_part("Alice_Smith.1") == "alice_smith.1"
def test_normalize_strips_unsafe_chars() -> None:
assert _normalize_username_local_part("hi there!") == "hi-there"
def test_username_to_email_uses_cm_com_by_default(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", raising=False)
assert _username_to_email("alice") == "alice@cm.com"
def test_username_to_email_honours_env_override(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("DEER_FLOW_USERNAME_LOGIN_DOMAIN", "example.org")
assert _username_to_email("alice") == "alice@example.org"
def test_username_to_email_passes_explicit_email_through() -> None:
assert _username_to_email("bob@legacy.io") == "bob@legacy.io"
@pytest.mark.asyncio
async def test_local_part_lookup_finds_legacy_domain_user() -> None:
"""Same username under legacy domain should match the new lookup."""
repo = FakeUserRepository()
_seed(repo, "alice@local.deerflow")
legacy = await repo.find_user_by_email_local_part("alice")
assert legacy is not None
assert legacy.email == "alice@local.deerflow"
@pytest.mark.asyncio
async def test_local_part_lookup_returns_oldest_when_collision() -> None:
"""When the same local part exists under multiple domains, prefer the
oldest row so re-logins land on the original account."""
repo = FakeUserRepository()
older = _seed(repo, "alice@local.deerflow")
older.created_at = datetime(2024, 1, 1, tzinfo=UTC)
newer = _seed(repo, "alice@cm.com")
newer.created_at = datetime(2026, 1, 1, tzinfo=UTC)
match = await repo.find_user_by_email_local_part("alice")
assert match is not None
assert match.email == "alice@local.deerflow"
@pytest.mark.asyncio
async def test_local_part_lookup_misses_unknown_username() -> None:
repo = FakeUserRepository()
_seed(repo, "alice@cm.com")
assert await repo.find_user_by_email_local_part("bob") is None