Merge branch 'main' into rayhpeng/hexagonal-feedback-slice

This commit is contained in:
rayhpeng 2026-07-31 10:02:24 +08:00 committed by GitHub
commit f27fa2188b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
55 changed files with 2712 additions and 203 deletions

View File

@ -245,6 +245,14 @@ This section accumulates work toward the **2.1.0** milestone
### Fixed
- **mcp:** Isolate Settings > Tools enable/disable updates to one MCP server, so
an unrelated disallowed stdio command no longer blocks every switch; allow
disabling a disallowed target while still rejecting its re-enable, preserve
the raw extensions config, honor the MCP-spec `transport` alias when enabling
SSE/HTTP servers, surface backend validation details in the UI, and atomically
replace the shared config for MCP, skill, and embedded-client updates so
interrupted writes cannot leave it truncated.
([#4574])
- **runtime:** Thread metadata now switches to `running` only after the run passes
the startup barrier, so pending-cancelled runs no longer briefly project
`running`; clients may observe the prior thread status during worker startup.

View File

@ -401,6 +401,9 @@ See the [Sandbox Configuration Guide](backend/docs/CONFIGURATION.md#sandbox) to
DeerFlow supports configurable MCP servers and skills to extend its capabilities.
For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`, `refresh_token`).
For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`.
Settings > Tools updates one MCP server at a time: an invalid stdio command on one server no longer blocks toggling another, while enabling that invalid server remains protected by the command allowlist and surfaces the backend validation message in the UI.
Targeted updates accept both DeerFlow's `type` field and the MCP-spec `transport` field for SSE/HTTP servers.
Runtime MCP and skill updates replace `extensions_config.json` atomically, so an interrupted write cannot leave the shared configuration truncated or partially written.
MCP routing hints can also prefer a specific MCP tool for matching requests without forbidding other tools. When `tool_search` defers MCP schemas, matching routing metadata can auto-promote up to `tool_search.auto_promote_top_k` deferred schemas before the model call.
See the [MCP Server Guide](backend/docs/MCP_SERVER.md) for detailed instructions.
@ -926,6 +929,8 @@ Image bytes loaded for a vision-model call are transient: DeerFlow removes the h
After each run, DeerFlow records a workspace change summary for the run-owned `workspace` and `outputs` directories. The Web UI shows a compact "files changed" badge on the assistant turn; opening it reveals created, modified, and deleted files with text diffs when safe to display. Uploads are excluded because they are user inputs, not agent-generated changes. Large, binary, or sensitive-looking files are shown as metadata only.
Files presented through `present_files` remain part of the thread's artifact state, and the Web UI restores the artifact panel and selected document after a page refresh. The currently selected formal artifact is refreshed once when the run finishes so edits become visible without a manual reload.
With `AioSandboxProvider`, shell execution runs inside isolated containers. With `LocalSandboxProvider`, file tools still map to per-thread directories on the host, but host `bash` is disabled by default because it is not a secure isolation boundary. Re-enable host bash only for fully trusted local workflows. Host bash commands have a wall-clock timeout, and long-lived processes should be started in the background with output redirected to a workspace log.
`AioSandboxProvider` normally detects thread-data mounts from its backend: local

View File

@ -486,7 +486,7 @@ Localhost persistence deliberately reads the direct request `Host` and ignores `
| **Models** (`/api/models`) | `GET /` - list models; `GET /{name}` - model details |
| **Features** (`/api/features`) | `GET /` - report config-gated feature availability (`agents_api.enabled`, `browser_control.enabled`) for frontend UI gating |
| **Console** (`/api/console`) | Read-only cross-thread observability for the current user (the data layer for an operations dashboard or external monitoring): `GET /stats` - headline counters (runs/threads/agents/tokens/cost); `GET /runs` - paginated run history joined with thread titles (per-run cost); `GET /usage` - zero-filled daily token series + per-model breakdown with spend. Queries `runs`/`threads_meta` directly as a reporting layer (no new `RunStore` methods); requires a SQL database backend — returns 503 on `database.backend: memory`. Real-cost estimation reads optional `models[*].pricing` (`currency`, `input_per_million`, `output_per_million`, `input_cache_hit_per_million`; `ModelConfig` is `extra="allow"`, so no schema change) and prices each run from its `token_usage_by_model` input/output split. Pricing is **cache-aware**: `RunJournal` accumulates prompt-cache hits from `usage_metadata.input_token_details.cache_read` into a sparse `cache_read_tokens` bucket key (also threaded through `SubagentTokenCollector``record_external_llm_usage_records`), and cache-hit input tokens are billed at `input_cache_hit_per_million` (omitted → billed at the miss price, a conservative upper bound). All priced models must use one currency; mixed currencies disable cost reporting and leave cost/currency fields null instead of producing invalid aggregates. Legacy rows fall back to run-level totals at `model_name`; unpriced models yield `cost: null` and cost fields are null when no pricing is configured |
| **MCP** (`/api/mcp`) | `GET /config` - get config; `PUT /config` - update config (saves to extensions_config.json) |
| **MCP** (`/api/mcp`) | `GET /config` - get config; `PUT /config` - replace the full config with whole-payload stdio validation; `PATCH /config` - toggle one server while preserving the raw extensions config and validating only an enabled target; both writes reload config and reset the process-local MCP cache |
| **Skills** (`/api/skills`) | `GET /` - list skills; `GET /{name}` - details; `PUT /{name}` - update enabled; `POST /install` - install from .skill archive (accepts standard optional frontmatter like `version`, `author`, `compatibility`); `POST /reload` - admin-only process-local prompt-cache invalidation after trusted external filesystem changes |
| **Integrations** (`/api/integrations`) | `GET /lark/status` - inspect managed Lark/Feishu CLI integration state, including `sandbox_runtime_mode` / `sandbox_runtime_ready` (whether `lark-cli` will actually be present in the sandbox at chat time); `POST /lark/install` - admin-only install of the official `lark-*` managed skill pack; `POST /lark/config/start` and `/lark/config/complete` - internal first-time Lark connection setup; `POST /lark/auth/start` and `/lark/auth/complete` - browser device-flow user authorization without terminal access, with optional `domains` / exact `scope` for incremental permission grants |
| **Memory** (`/api/memory`) | `GET /` - memory data; `POST /reload` - force reload; `GET /config` - config; `GET /status` - config + data |
@ -752,7 +752,7 @@ E2B output sync records remote file versions and actual host file metadata in a
add a parallel routing middleware for PR1-style preference hints.
- **Stdio file outputs**: Persistent stdio sessions are scoped by `user_id:thread_id`. For stdio transports only, DeerFlow pins the subprocess default `cwd` to the thread workspace and `TMPDIR`/`TMP`/`TEMP` to `workspace/.mcp/tmp/`, unless the operator explicitly configured `cwd` or temp env values. SSE/HTTP transports skip this filesystem prep entirely.
- **Stdio path translation**: MCP-returned local file references are not copied. If a `ResourceLink` or conservative free-text path resolves to an existing file inside the thread's mounted user-data tree, it is translated deterministically to `/mnt/user-data/...`; paths outside that tree remain unchanged.
- **Runtime updates**: Gateway API saves to extensions_config.json; the Gateway-embedded runtime detects changes via the resolved-path + content-signature check above, so multi-worker / stale-mtime deployments still pick up an added/removed MCP server without a restart (the `PUT /api/mcp/config` reset only clears the cache in its own worker)
- **Runtime updates**: Gateway API saves to extensions_config.json; the Gateway-embedded runtime detects changes via the resolved-path + content-signature check above, so multi-worker / stale-mtime deployments still pick up an added/removed MCP server without a restart (`PUT /api/mcp/config` keeps whole-payload validation, while `PATCH /api/mcp/config` changes only one server's `enabled` field, normalizes the same `type`/MCP-spec `transport` alias as the runtime config model, and validates the target only when enabling it; either endpoint's reset clears the cache only in its own worker). MCP, skill, and embedded-client writers share `atomic_write_extensions_config()`, which writes and fsyncs a same-directory temporary file before `os.replace()` and preserves an existing file's mode and symlink target; failed serialization or replacement leaves the prior config intact and cleans up the temporary file.
### Skills System (`packages/harness/deerflow/skills/`)
@ -1219,7 +1219,7 @@ Config is env-driven like the others — `MonocleTracingConfig`, built in `get_t
- `skills` - Map of skill name → state (enabled)
- `middlewares` - Zero-argument `AgentMiddleware` class paths for lead and subagent runtime extension. `config.yaml -> extensions` can override these fields after validation; overrides are replace-per-field, not list concatenation.
Gateway API endpoints and `DeerFlowClient` methods can modify MCP servers and skill state at runtime; `middlewares` remains an operator-controlled config-file extension point.
Gateway API endpoints and `DeerFlowClient` methods can modify MCP servers and skill state at runtime; their `extensions_config.json` writes use the shared atomic replacement helper, while `middlewares` remains an operator-controlled config-file extension point.
### Embedded Client (`packages/harness/deerflow/client.py`)

View File

@ -827,7 +827,8 @@ async def oauth_callback(
# ── Issue DeerFlow session ───────────────────────────────────────
token = create_access_token(str(user.id), token_version=user.token_version)
redirect_target = state_payload.next_path or "/workspace"
# Revalidate as defense-in-depth if future state writers populate this target.
redirect_target = validate_next_param(state_payload.next_path) or "/workspace"
frontend_base = oidc_config.frontend_base_url or ""
callback_redirect = f"{frontend_base}/auth/callback?next={urllib.parse.quote(redirect_target)}"
@ -855,7 +856,8 @@ def validate_next_param(next_param: str | None) -> str | None:
"""Validate and sanitize the ``next`` redirect parameter.
Only allows relative paths starting with ``/``. Rejects protocol-relative
URLs (``//``), absolute URLs, and URLs with embedded protocols.
URLs (``//``), absolute URLs, URLs with embedded protocols, and backslashes
that URL parsers may reinterpret as forward slashes.
"""
if not next_param:
return None
@ -863,6 +865,8 @@ def validate_next_param(next_param: str | None) -> str | None:
return None
if next_param.startswith("//") or next_param.startswith("http://") or next_param.startswith("https://"):
return None
if "\\" in next_param:
return None
if ":" in next_param:
return None
return next_param

View File

@ -7,10 +7,19 @@ from pathlib import Path
from typing import Any, Literal
from fastapi import APIRouter, HTTPException, Request, status
from pydantic import BaseModel, ConfigDict, Field
from pydantic import BaseModel, ConfigDict, Field, model_validator
from app.gateway.deps import require_admin_user
from deerflow.config.extensions_config import ExtensionsConfig, McpRoutingConfig, McpToolOverride, extensions_config_write_lock, get_extensions_config, reload_extensions_config
from deerflow.config.extensions_config import (
ExtensionsConfig,
McpRoutingConfig,
McpToolOverride,
atomic_write_extensions_config,
extensions_config_write_lock,
get_extensions_config,
normalize_mcp_transport_alias,
reload_extensions_config,
)
from deerflow.mcp.cache import reset_mcp_tools_cache
logger = logging.getLogger(__name__)
@ -60,6 +69,12 @@ class McpServerConfigResponse(BaseModel):
tool_call_timeout: float | None = Field(default=None, description="Timeout in seconds for individual stdio MCP tool calls")
model_config = ConfigDict(extra="allow")
@model_validator(mode="before")
@classmethod
def _accept_transport_alias(cls, data: Any) -> Any:
"""Keep API parsing aligned with the runtime MCP config model."""
return normalize_mcp_transport_alias(data)
class McpConfigResponse(BaseModel):
"""Response model for MCP configuration."""
@ -79,6 +94,17 @@ class McpConfigUpdateRequest(BaseModel):
)
class McpServerStateUpdateRequest(BaseModel):
"""Request model for enabling or disabling one MCP server."""
server_name: str = Field(
...,
min_length=1,
description="Name of the MCP server to update",
)
enabled: bool = Field(..., description="Whether the MCP server is enabled")
class McpCacheResetResponse(BaseModel):
"""Response model for resetting the MCP tools cache."""
@ -383,9 +409,7 @@ def _apply_mcp_config_update(body: McpConfigUpdateRequest) -> dict:
config_data["mcpServers"] = {name: server.model_dump() for name, server in merged_servers.items()}
config_data["skills"] = {name: {"enabled": skill.enabled} for name, skill in current_config.skills.items()}
# Write the configuration to file
with open(config_path, "w", encoding="utf-8") as f:
json.dump(config_data, f, indent=2)
atomic_write_extensions_config(config_path, config_data)
logger.info(f"MCP configuration updated and saved to: {config_path}")
@ -396,6 +420,43 @@ def _apply_mcp_config_update(body: McpConfigUpdateRequest) -> dict:
return reloaded_config.mcp_servers
def _apply_mcp_server_state_update(body: McpServerStateUpdateRequest) -> dict:
"""Update one server state while preserving the raw extensions config."""
with extensions_config_write_lock:
config_path = ExtensionsConfig.resolve_config_path()
if config_path is None or not config_path.exists():
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"MCP server '{body.server_name}' not found",
)
with open(config_path, encoding="utf-8") as f:
raw_data = json.load(f)
raw_servers = raw_data.get("mcpServers", {})
raw_server = raw_servers.get(body.server_name) if isinstance(raw_servers, dict) else None
if not isinstance(raw_server, dict):
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"MCP server '{body.server_name}' not found",
)
if body.enabled:
target_server = McpServerConfigResponse(**raw_server)
_validate_mcp_update_request(
McpConfigUpdateRequest(
mcp_servers={body.server_name: target_server},
)
)
raw_server["enabled"] = body.enabled
atomic_write_extensions_config(config_path, raw_data)
logger.info("MCP server %s enabled state updated to %s", body.server_name, body.enabled)
reloaded_config = reload_extensions_config()
return reloaded_config.mcp_servers
@router.post(
"/mcp/cache/reset",
response_model=McpCacheResetResponse,
@ -475,3 +536,25 @@ async def update_mcp_configuration(request: Request, body: McpConfigUpdateReques
except Exception as e:
logger.error(f"Failed to update MCP configuration: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=f"Failed to update MCP configuration: {str(e)}")
@router.patch(
"/mcp/config",
response_model=McpConfigResponse,
summary="Update MCP Server State",
description="Enable or disable one MCP server without replacing the full extensions configuration.",
)
async def update_mcp_server_state(request: Request, body: McpServerStateUpdateRequest) -> McpConfigResponse:
"""Enable or disable one MCP server and reload the MCP tool cache."""
try:
await require_admin_user(request, detail=_ADMIN_REQUIRED_DETAIL)
reloaded_servers = await asyncio.to_thread(_apply_mcp_server_state_update, body)
servers = {name: _mask_server_config(McpServerConfigResponse(**server.model_dump())) for name, server in reloaded_servers.items()}
reset_mcp_tools_cache()
return McpConfigResponse(mcp_servers=servers)
except HTTPException:
raise
except Exception as e:
logger.error("Failed to update MCP server %s state: %s", body.server_name, e, exc_info=True)
raise HTTPException(status_code=500, detail=f"Failed to update MCP server state: {str(e)}")

View File

@ -1,5 +1,4 @@
import asyncio
import json
import logging
import tempfile
from pathlib import Path
@ -12,7 +11,14 @@ from app.gateway.deps import get_config, require_admin_user
from app.gateway.path_utils import resolve_thread_virtual_path
from deerflow.agents.lead_agent.prompt import clear_skills_system_prompt_cache, refresh_skills_system_prompt_cache_async, refresh_user_skills_system_prompt_cache_async
from deerflow.config.app_config import AppConfig
from deerflow.config.extensions_config import ExtensionsConfig, SkillStateConfig, extensions_config_write_lock, get_extensions_config, reload_extensions_config
from deerflow.config.extensions_config import (
ExtensionsConfig,
SkillStateConfig,
atomic_write_extensions_config,
extensions_config_write_lock,
get_extensions_config,
reload_extensions_config,
)
from deerflow.runtime.user_context import get_effective_user_id
from deerflow.skills import Skill
from deerflow.skills.installer import SkillAlreadyExistsError, SkillSecurityScanError
@ -445,8 +451,7 @@ def _write_extensions_skill_state(skill_name: str, enabled: bool) -> None:
config_data = extensions_config.to_file_dict()
with open(config_path, "w", encoding="utf-8") as f:
json.dump(config_data, f, indent=2)
atomic_write_extensions_config(config_path, config_data)
logger.info(f"Skills configuration updated and saved to: {config_path}")
reload_extensions_config()

View File

@ -393,6 +393,36 @@ deployment needs additional trusted launchers.
}
```
#### Update One MCP Server State
Enable or disable one configured MCP server without replacing the full
extensions configuration.
```http
PATCH /api/mcp/config
Content-Type: application/json
```
Requires an authenticated admin session. Enabling a `stdio` server validates
that server's `command` against the same allowlist used by the full `PUT`
endpoint. Disabling a server does not require its command to be allowlisted, and
invalid commands on other servers do not block the update. The endpoint
preserves secrets, environment-variable placeholders, skills, custom server
fields, and other top-level extensions config. SSE/HTTP targets may use either
DeerFlow's `type` field or the MCP-spec `transport` field.
**Request Body:**
```json
{
"server_name": "semantic-scholar",
"enabled": false
}
```
The response is the full masked MCP configuration, matching `GET` and `PUT`.
An unknown `server_name` returns `404`; attempting to enable a server with a
disallowed `stdio` command returns `400`.
#### Reset MCP Tools Cache
Clear cached MCP tools and persistent MCP sessions process-wide. This affects

View File

@ -18,12 +18,10 @@ Usage:
import asyncio
import concurrent.futures
import copy
import json
import logging
import mimetypes
import os
import shutil
import tempfile
import uuid
from collections.abc import Generator, Mapping, Sequence
from dataclasses import dataclass, field
@ -41,7 +39,13 @@ from deerflow.agents.thread_state import get_thread_state_schema, normalize_midd
from deerflow.authz.principal import build_principal_from_context
from deerflow.config.agents_config import AGENT_NAME_PATTERN
from deerflow.config.app_config import get_app_config, is_trace_correlation_enabled, reload_app_config
from deerflow.config.extensions_config import ExtensionsConfig, SkillStateConfig, get_extensions_config, reload_extensions_config
from deerflow.config.extensions_config import (
ExtensionsConfig,
SkillStateConfig,
atomic_write_extensions_config,
get_extensions_config,
reload_extensions_config,
)
from deerflow.config.paths import get_paths
from deerflow.models import create_chat_model
from deerflow.runtime import CheckpointStateAccessor
@ -229,20 +233,7 @@ class DeerFlowClient:
@staticmethod
def _atomic_write_json(path: Path, data: dict) -> None:
"""Write JSON to *path* atomically (temp file + replace)."""
fd = tempfile.NamedTemporaryFile(
mode="w",
dir=path.parent,
suffix=".tmp",
delete=False,
)
try:
json.dump(data, fd, indent=2)
fd.close()
Path(fd.name).replace(path)
except BaseException:
fd.close()
Path(fd.name).unlink(missing_ok=True)
raise
atomic_write_extensions_config(path, data)
def _get_runnable_config(self, thread_id: str, **overrides) -> RunnableConfig:
"""Build a RunnableConfig for agent invocation."""

View File

@ -446,6 +446,16 @@ class AppConfig(BaseModel):
if previous_checkpointer_config != config.checkpointer:
# These runtime singletons derive their backend from checkpointer config.
# Keep imports local to avoid cycles: both providers import get_app_config.
#
# The unified ``database`` section is intentionally NOT handled here.
# ``database`` is a restart-required field (reload_boundary.STARTUP_ONLY_FIELDS):
# ``init_engine_from_config()`` builds the ORM engine once at startup and
# never rebuilds it on a config.yaml edit. Resetting only the sync
# checkpointer/store singletons on a live ``database``/``postgres_schema``
# change would half-migrate the deployment -- new checkpoint/store tables
# would land in the new schema while ORM rows keep landing in the old one,
# with no error surfaced. Requiring the documented restart keeps the
# deployment self-consistent.
from deerflow.runtime.checkpointer import reset_checkpointer
from deerflow.runtime.store import reset_store

View File

@ -2,7 +2,9 @@
from typing import Literal
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, field_validator
from deerflow.config.postgres_schema import POSTGRES_SCHEMA_PATTERN, validate_postgres_schema
CheckpointerType = Literal["memory", "sqlite", "postgres"]
@ -24,6 +26,15 @@ class CheckpointerConfig(BaseModel):
"For sqlite, use a file path like '.deer-flow/checkpoints.db' or ':memory:' for in-memory. "
"For postgres, use a DSN like 'postgresql://user:pass@localhost:5432/db'.",
)
postgres_schema: str = Field(
default="",
description=(f"PostgreSQL schema for legacy checkpointer/store tables (postgres only). Empty string keeps the server default search_path (usually 'public'). Only plain identifiers are allowed: {POSTGRES_SCHEMA_PATTERN}."),
)
@field_validator("postgres_schema")
@classmethod
def _validate_postgres_schema(cls, value: str) -> str:
return validate_postgres_schema(value)
# Global configuration instance — None means no checkpointer is configured.

View File

@ -35,7 +35,9 @@ import logging
import os
from typing import Any, Literal
from pydantic import BaseModel, Field, model_validator
from pydantic import BaseModel, Field, field_validator, model_validator
from deerflow.config.postgres_schema import POSTGRES_SCHEMA_PATTERN, validate_postgres_schema
logger = logging.getLogger(__name__)
@ -151,6 +153,22 @@ class DatabaseConfig(BaseModel):
gt=0,
description="Timeout in seconds for app ORM PostgreSQL commands. Set to null to disable the command timeout.",
)
postgres_schema: str = Field(
default="",
description=(
"PostgreSQL schema for both app ORM tables and LangGraph "
"checkpointer/store tables (postgres only). Empty string keeps "
"the server default search_path (usually 'public'). When set, "
"the schema is created automatically at startup and applied via "
"connection-level search_path. Only plain identifiers are "
f"allowed: {POSTGRES_SCHEMA_PATTERN}."
),
)
@field_validator("postgres_schema")
@classmethod
def _validate_postgres_schema(cls, value: str) -> str:
return validate_postgres_schema(value)
# -- Legacy key migration (not user-configured) --

View File

@ -3,6 +3,8 @@
import json
import logging
import os
import stat
import tempfile
import threading
from pathlib import Path
from typing import Any, Literal
@ -14,6 +16,15 @@ from deerflow.config.runtime_paths import existing_project_file
logger = logging.getLogger(__name__)
def normalize_mcp_transport_alias(data: Any) -> Any:
"""Promote MCP-spec ``transport`` to ``type`` when ``type`` is absent."""
if isinstance(data, dict):
transport = data.get("transport")
if transport and not data.get("type"):
return {**data, "type": transport}
return data
class McpRoutingConfig(BaseModel):
"""Soft routing hints for MCP tool preference."""
@ -105,11 +116,7 @@ class McpServerConfig(BaseModel):
``stdio`` (the default). This validator normalizes the two so either
spelling works, with ``type`` taking precedence when both are provided.
"""
if isinstance(data, dict):
transport = data.get("transport")
if transport and not data.get("type"):
data = {**data, "type": transport}
return data
return normalize_mcp_transport_alias(data)
def resolve_effective_mcp_routing(server_config: McpServerConfig | None, original_tool_name: str) -> dict[str, Any]:
@ -323,6 +330,70 @@ class ExtensionsConfig(BaseModel):
_extensions_config: ExtensionsConfig | None = None
def _fsync_directory_best_effort(directory: Path) -> None:
"""Persist a directory entry update where the platform supports it."""
if os.name == "nt":
return
try:
directory_fd = os.open(directory, os.O_RDONLY)
except OSError:
return
try:
os.fsync(directory_fd)
except OSError:
logger.debug("Could not fsync extensions config directory: %s", directory, exc_info=True)
finally:
try:
os.close(directory_fd)
except OSError:
logger.debug("Could not close extensions config directory: %s", directory, exc_info=True)
def atomic_write_extensions_config(path: Path, data: dict[str, Any]) -> None:
"""Write extensions config without exposing a truncated or partial file."""
path = Path(path)
target_path = path.resolve(strict=False) if path.is_symlink() else path
target_path.parent.mkdir(parents=True, exist_ok=True)
existing_mode: int | None = None
try:
existing_mode = stat.S_IMODE(target_path.stat().st_mode)
except FileNotFoundError:
pass
temporary_path: Path | None = None
try:
with tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
dir=target_path.parent,
prefix=f".{target_path.name}.",
suffix=".tmp",
delete=False,
) as temporary_file:
temporary_path = Path(temporary_file.name)
json.dump(data, temporary_file, indent=2)
if existing_mode is not None:
temporary_path.chmod(existing_mode)
temporary_file.flush()
os.fsync(temporary_file.fileno())
os.replace(temporary_path, target_path)
_fsync_directory_best_effort(target_path.parent)
finally:
if temporary_path is not None:
try:
temporary_path.unlink(missing_ok=True)
except OSError:
logger.warning(
"Could not remove temporary extensions config file: %s",
temporary_path,
exc_info=True,
)
def get_extensions_config() -> ExtensionsConfig:
"""Get the extensions config instance.

View File

@ -0,0 +1,27 @@
"""Shared validation for PostgreSQL schema names."""
from __future__ import annotations
import re
# Lowercase-only on purpose. The schema is created *quoted*
# (``CREATE SCHEMA IF NOT EXISTS "<schema>"``, case-preserved) but pinned via an
# *unquoted* ``search_path`` token, which PostgreSQL folds to lowercase. Allowing
# uppercase here would let the two diverge so tables silently land in ``public``.
# No anchors: validation uses ``re.fullmatch`` so the whole value must match.
# A ``$``-anchored ``re.match`` would accept a trailing newline (``"deerflow\n"``):
# Python's ``$`` matches just before a single trailing ``\n``, which then
# creates a *quoted* schema literally named ``deerflow\n`` while the *unquoted*
# ``search_path`` folds to ``deerflow`` and misses it, silently landing tables
# in ``public``.
POSTGRES_SCHEMA_PATTERN = r"[a-z_][a-z0-9_]{0,62}"
_POSTGRES_SCHEMA_RE = re.compile(POSTGRES_SCHEMA_PATTERN)
def validate_postgres_schema(value: str) -> str:
"""Validate the v1 plain-identifier PostgreSQL schema contract."""
if value == "":
return value
if not _POSTGRES_SCHEMA_RE.fullmatch(value):
raise ValueError(f"postgres_schema must be a plain lowercase PostgreSQL identifier matching {POSTGRES_SCHEMA_PATTERN}; got {value!r}. Mixed-case and quoted identifiers are not supported.")
return value

View File

@ -117,10 +117,10 @@ def _local_uri_to_virtual_path(
try:
real = src.resolve()
if not real.is_file():
return None
except OSError:
return None
if not real.is_file():
return None
try:
user_data_root = get_paths().sandbox_user_data_dir(thread_id, user_id=user_id).resolve()

View File

@ -234,16 +234,25 @@ def _alembic_safe_url(engine: AsyncEngine) -> str:
return _escape_url_for_alembic(rendered)
def _get_alembic_config(engine: AsyncEngine) -> AlembicConfig:
def _get_alembic_config(engine: AsyncEngine, *, postgres_schema: str = "") -> AlembicConfig:
"""Build an in-process alembic config pointing at our migrations dir.
Avoids reading ``alembic.ini`` from disk so the production runtime doesn't
depend on a working-directory-relative file lookup. The ``script_location``
is anchored at the package path on disk.
When *postgres_schema* is set it is forwarded as the ``deerflow_pg_schema``
main option so ``env.py`` can pin its alembic-spawned engine's
``search_path`` to the same schema the app engine uses. Without it,
alembic's own engine -- built from the bare URL -- would create
``alembic_version`` and all migration DDL in the default (``public``)
schema while the app tables land in the custom schema.
"""
cfg = AlembicConfig()
cfg.set_main_option("script_location", str(_MIGRATIONS_DIR))
cfg.set_main_option("sqlalchemy.url", _alembic_safe_url(engine))
if postgres_schema:
cfg.set_main_option("deerflow_pg_schema", postgres_schema)
return cfg
@ -470,7 +479,7 @@ def _bootstrap_lock(engine: AsyncEngine, *, backend: str):
# ---------------------------------------------------------------------------
async def bootstrap_schema(engine: AsyncEngine, *, backend: str) -> None:
async def bootstrap_schema(engine: AsyncEngine, *, backend: str, postgres_schema: str = "") -> None:
"""Bring the DB schema to head.
Postgres calls are serialised across processes with an advisory lock.
@ -480,9 +489,14 @@ async def bootstrap_schema(engine: AsyncEngine, *, backend: str) -> None:
Branch dispatch is documented at module top. ``alembic.command.stamp`` and
``alembic.command.upgrade`` are synchronous and would block the event
loop; both are wrapped in ``asyncio.to_thread``.
*postgres_schema*, when set, is forwarded to the alembic config so the
alembic-spawned engine pins its ``search_path`` to that schema. The target
schema must already exist (``init_engine`` issues ``CREATE SCHEMA`` before
calling this). Ignored for non-postgres backends.
"""
head = _get_head_revision()
cfg = _get_alembic_config(engine)
cfg = _get_alembic_config(engine, postgres_schema=postgres_schema if backend == "postgres" else "")
async with _bootstrap_lock(engine, backend=backend):
async with engine.connect() as conn:

View File

@ -33,17 +33,18 @@ def _postgres_engine_kwargs(
pool_size: int,
pool_recycle: int = POSTGRES_POOL_RECYCLE_SECONDS,
command_timeout: float | None = POSTGRES_COMMAND_TIMEOUT_SECONDS,
connect_args: dict[str, object] | None = None,
) -> dict[str, object]:
"""Build the shared SQLAlchemy engine options for PostgreSQL."""
connect_args = {}
merged_connect_args = dict(connect_args or {})
if command_timeout is not None:
connect_args["command_timeout"] = command_timeout
merged_connect_args["command_timeout"] = command_timeout
return {
"echo": echo,
"pool_size": pool_size,
"pool_pre_ping": True,
"pool_recycle": pool_recycle,
"connect_args": connect_args,
"connect_args": merged_connect_args,
"json_serializer": _json_serializer,
}
@ -90,6 +91,7 @@ async def init_engine(
pool_recycle: int = POSTGRES_POOL_RECYCLE_SECONDS,
command_timeout: float | None = POSTGRES_COMMAND_TIMEOUT_SECONDS,
sqlite_dir: str = "",
postgres_schema: str = "",
) -> None:
"""Create the async engine and session factory, then auto-create tables.
@ -101,6 +103,10 @@ async def init_engine(
pool_recycle: Seconds before Postgres connections are recycled.
command_timeout: Timeout in seconds for app ORM Postgres commands, or None to disable.
sqlite_dir: Directory to create for SQLite (ensured to exist).
postgres_schema: Target PostgreSQL schema. When set, the engine
pins the connection ``search_path`` to it via asyncpg
``server_settings`` and the schema is created (if missing)
before tables are auto-created. Ignored for non-postgres.
"""
global _engine, _session_factory
@ -161,6 +167,9 @@ async def init_engine(
finally:
cursor.close()
elif backend == "postgres":
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
pg_connect_args = build_asyncpg_connect_args(postgres_schema)
_engine = create_async_engine(
url,
**_postgres_engine_kwargs(
@ -168,6 +177,7 @@ async def init_engine(
pool_size=pool_size,
pool_recycle=pool_recycle,
command_timeout=command_timeout,
connect_args=pg_connect_args,
),
)
else:
@ -186,13 +196,28 @@ async def init_engine(
# See deerflow.persistence.bootstrap for the full state machine.
from deerflow.persistence.bootstrap import bootstrap_schema
async def _ensure_postgres_schema() -> None:
# CREATE SCHEMA is DDL and is unaffected by search_path, so it is
# safe even though the connection's search_path already points at
# the (not-yet-existing) target schema. It must run before
# ``bootstrap_schema`` so the subsequent ``create_all`` / alembic
# DDL lands in the target schema instead of failing on a missing one.
if backend == "postgres" and postgres_schema:
from sqlalchemy.schema import CreateSchema
async with _engine.begin() as conn:
await conn.execute(CreateSchema(postgres_schema, if_not_exists=True))
try:
await bootstrap_schema(_engine, backend=backend)
await _ensure_postgres_schema()
await bootstrap_schema(_engine, backend=backend, postgres_schema=postgres_schema)
except Exception as exc:
if backend == "postgres" and "does not exist" in str(exc):
# Database not yet created -- attempt to auto-create it, then retry.
await _auto_create_postgres_db(url)
# Rebuild engine against the now-existing database
# Rebuild engine against the now-existing database. The rebuilt
# engine MUST keep the same connect_args so the retried bootstrap
# lands in the target schema, not the default one.
await _engine.dispose()
_engine = create_async_engine(
url,
@ -201,10 +226,12 @@ async def init_engine(
pool_size=pool_size,
pool_recycle=pool_recycle,
command_timeout=command_timeout,
connect_args=pg_connect_args,
),
)
_session_factory = async_sessionmaker(_engine, expire_on_commit=False)
await bootstrap_schema(_engine, backend=backend)
await _ensure_postgres_schema()
await bootstrap_schema(_engine, backend=backend, postgres_schema=postgres_schema)
else:
raise
@ -224,6 +251,7 @@ async def init_engine_from_config(config) -> None:
pool_recycle=config.pool_recycle,
command_timeout=config.command_timeout,
sqlite_dir=config.sqlite_dir if config.backend == "sqlite" else "",
postgres_schema=config.postgres_schema if config.backend == "postgres" else "",
)

View File

@ -72,7 +72,26 @@ def do_run_migrations(connection):
async def run_migrations_online() -> None:
connectable = create_async_engine(config.get_main_option("sqlalchemy.url"))
url = config.get_main_option("sqlalchemy.url")
# When a custom Postgres schema is configured, pin the alembic-spawned
# engine's search_path to it. This engine is built from the bare URL and
# does NOT inherit the asyncpg ``server_settings`` the app engine sets via
# ``connect_args``, so without this both ``alembic_version`` and every
# migration's DDL would land in the default (``public``) schema while the
# ORM tables land in the custom schema. ``init_engine`` has already created
# the schema (``CREATE SCHEMA IF NOT EXISTS``) before bootstrap runs.
pg_schema = config.get_main_option("deerflow_pg_schema")
connect_args: dict = {}
# Accept both the canonical ``postgresql`` scheme and libpq's ``postgres``
# short scheme (with or without a SQLAlchemy ``+driver`` suffix) so a
# ``postgres://`` DSN still gets its search_path pinned instead of silently
# writing ``alembic_version`` + migration DDL to the default schema.
if pg_schema and url and url.split("+", 1)[0].split(":", 1)[0] in {"postgresql", "postgres"}:
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
connect_args = build_asyncpg_connect_args(pg_schema)
connectable = create_async_engine(url, connect_args=connect_args)
# Cross-process bootstrap safety for SQLite: every connection alembic
# opens needs a wide ``busy_timeout`` so that when another process holds

View File

@ -0,0 +1,257 @@
"""PostgreSQL schema helpers (Issue #3380).
Centralizes the driver-specific ways of pinning a connection's
``search_path`` to a target schema. The two PostgreSQL drivers DeerFlow
uses expect different mechanisms:
- **asyncpg** (app ORM engine): only honours ``server_settings`` passed
via SQLAlchemy ``connect_args``. It does not understand libpq's
``options=-c ...`` syntax.
- **psycopg** (LangGraph checkpointer/store): uses the libpq
``options=-c search_path=...`` connection parameter, either as a pool
kwarg or encoded into the DSN query string.
Schema names are validated upstream by
:class:`deerflow.config.database_config.DatabaseConfig` to be plain
identifiers. SQL-emitting helpers re-validate at the boundary as
defense-in-depth; connection-argument helpers only assemble driver payloads.
"""
from __future__ import annotations
import re
from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit
def build_asyncpg_connect_args(schema: str) -> dict:
"""Return SQLAlchemy ``connect_args`` that pin asyncpg's search_path.
Empty *schema* yields ``{}`` so the engine keeps the server default.
"""
if not schema:
return {}
return {"server_settings": {"search_path": schema}}
def build_psycopg_options(schema: str) -> str | None:
"""Return the libpq ``options`` value for psycopg pool kwargs.
Empty *schema* yields ``None`` so callers can skip setting the kwarg.
"""
if not schema:
return None
return f"-c search_path={schema}"
def _split_libpq_options(options: str) -> list[str]:
"""Tokenize a libpq ``options`` string.
libpq splits on unescaped whitespace; a backslash escapes the next
character (so ``\\ `` is a literal space and ``\\\\`` a literal backslash).
This is NOT POSIX shell quoting -- single/double quotes are literal here.
"""
tokens: list[str] = []
current: list[str] = []
in_token = False
escaped = False
for char in options:
if escaped:
current.append(char)
escaped = False
in_token = True
continue
if char == "\\":
escaped = True
in_token = True
continue
if char.isspace():
if in_token:
tokens.append("".join(current))
current = []
in_token = False
continue
current.append(char)
in_token = True
if in_token:
tokens.append("".join(current))
return tokens
def _join_libpq_options(tokens: list[str]) -> str:
"""Join tokens into a libpq ``options`` string.
Whitespace and backslashes inside a token are backslash-escaped so libpq
keeps each token intact. ``shlex.join`` cannot be used: it emits POSIX
shell quoting (single quotes), which libpq treats as literal characters.
All whitespace bytes are escaped, not just spaces: ``_split_libpq_options``
preserves a backslash-escaped TAB/CR/LF as part of one token, so re-joining
with a bare whitespace byte would let libpq re-tokenize on it and corrupt a
caller's pre-existing ``options`` value.
"""
escaped = [re.sub(r"([\\\s])", r"\\\1", token) for token in tokens]
return " ".join(escaped)
def _merge_search_path_option(existing_options: str, schema: str) -> str:
"""Return libpq options with search_path replaced while preserving others."""
new_option = build_psycopg_options(schema)
if not new_option:
return existing_options
if not existing_options:
return new_option
tokens = _split_libpq_options(existing_options)
merged: list[str] = []
index = 0
while index < len(tokens):
token = tokens[index]
if token == "-c" and index + 1 < len(tokens):
setting = tokens[index + 1]
if setting.split("=", 1)[0] == "search_path":
index += 2
continue
merged.extend([token, setting])
index += 2
continue
if token.startswith("-csearch_path="):
index += 1
continue
merged.append(token)
index += 1
merged.extend(_split_libpq_options(new_option))
return _join_libpq_options(merged)
def create_schema_sql(schema: str) -> str | None:
"""Return a safe CREATE SCHEMA statement for a validated plain identifier.
Defense-in-depth: the identifier is re-validated here rather than trusting
the distant pydantic validator. ``create_schema_sql`` is publicly exported
and psycopg accepts multiple ``;``-separated statements, so a future caller
that bypasses ``DatabaseConfig``/``CheckpointerConfig`` (e.g. a test helper)
must not be able to inject SQL through this f-string boundary.
"""
if not schema:
return None
from deerflow.config.postgres_schema import validate_postgres_schema
validate_postgres_schema(schema)
return f'CREATE SCHEMA IF NOT EXISTS "{schema}"'
def normalize_libpq_dsn(dsn: str) -> str:
"""Return *dsn* with any SQLAlchemy ``+driver`` suffix dropped.
``DatabaseConfig.postgres_url`` may carry a SQLAlchemy driver suffix such
as ``postgresql+asyncpg://``. psycopg's libpq only understands the bare
``postgres``/``postgresql`` scheme, so a raw ``+asyncpg`` DSN handed to
``psycopg.connect`` raises an opaque parse error. Keyword/DSN strings
without a URL scheme (``host=... dbname=...``) are returned unchanged.
Raises ``ValueError`` for URL schemes that are not a PostgreSQL variant.
"""
parts = urlsplit(dsn)
if not parts.scheme:
return dsn
scheme_base = parts.scheme.split("+", 1)[0]
if scheme_base not in {"postgres", "postgresql"}:
raise ValueError(f"Unsupported PostgreSQL DSN scheme for schema injection: {parts.scheme!r}")
if scheme_base == parts.scheme:
return dsn
return urlunsplit((scheme_base, parts.netloc, parts.path, parts.query, parts.fragment))
def dsn_with_search_path(dsn: str, schema: str) -> str:
"""Return *dsn* with an ``options=-c search_path=<schema>`` query param.
Used for psycopg ``from_conn_string`` call sites that take a DSN
string rather than pool kwargs. The value contains a space and ``=``;
both are percent-encoded so libpq parses the URL correctly.
libpq only recognizes ``%XX`` percent-encoding in URI query values; it
does NOT treat ``+`` as a space (that is an HTML-form convention). So
the space MUST be encoded as ``%20`` rather than ``+``, otherwise libpq
sees a single broken token ``-c+search_path=...`` and the search_path is
never applied. Existing query parameters are preserved. Empty *schema*
returns *dsn* unchanged.
"""
if not schema:
return dsn
parts = urlsplit(dsn)
if not parts.scheme:
from psycopg.conninfo import conninfo_to_dict, make_conninfo
params = conninfo_to_dict(dsn)
params["options"] = _merge_search_path_option(params.get("options", ""), schema)
return make_conninfo(**params)
# DatabaseConfig.postgres_url may carry a SQLAlchemy driver suffix such as
# ``postgresql+asyncpg://``. psycopg's libpq only understands the bare
# ``postgres``/``postgresql`` scheme, so accept the compound form but emit
# a psycopg-consumable DSN by dropping the ``+driver`` part.
scheme_base = parts.scheme.split("+", 1)[0]
if scheme_base not in {"postgres", "postgresql"}:
raise ValueError(f"Unsupported PostgreSQL DSN scheme for schema injection: {parts.scheme!r}")
options_values: list[str] = []
query_pairs = []
for key, value in parse_qsl(parts.query, keep_blank_values=True):
if key == "options":
options_values.append(value)
else:
query_pairs.append((key, value))
options = _merge_search_path_option(" ".join(options_values), schema)
query_pairs.append(("options", options))
# quote_via=quote encodes space as %20 (libpq-safe), not + (form-style).
query = urlencode(query_pairs, quote_via=quote)
return urlunsplit((scheme_base, parts.netloc, parts.path, query, parts.fragment))
def ensure_postgres_schema(conn_string: str, schema: str, *, install_hint: str) -> None:
"""Create *schema* over a fresh sync psycopg connection.
No-op when *schema* is empty. A missing ``psycopg`` dependency is mapped to
*install_hint* so callers surface the same actionable message they use for
the rest of the backend. The DSN is normalized so a SQLAlchemy ``+driver``
suffix does not reach libpq.
"""
statement = create_schema_sql(schema)
if statement is None:
return
try:
import psycopg
except ImportError as exc:
raise ImportError(install_hint) from exc
# psycopg 3's ``Connection.__exit__`` only commits/rolls back -- it does NOT
# close the connection (a documented psycopg2->3 change). Use try/finally so
# the libpq connection is released deterministically, mirroring the async
# counterpart, instead of leaking it until GC.
conn = psycopg.connect(normalize_libpq_dsn(conn_string), autocommit=True)
try:
conn.execute(statement)
finally:
conn.close()
async def ensure_postgres_schema_async(conn_string: str, schema: str, *, install_hint: str) -> None:
"""Async counterpart of :func:`ensure_postgres_schema`."""
statement = create_schema_sql(schema)
if statement is None:
return
try:
import psycopg
except ImportError as exc:
raise ImportError(install_hint) from exc
conn = await psycopg.AsyncConnection.connect(normalize_libpq_dsn(conn_string), autocommit=True)
try:
await conn.execute(statement)
finally:
await conn.close()

View File

@ -25,6 +25,7 @@ from collections.abc import AsyncIterator
from langgraph.types import Checkpointer
from deerflow.config.app_config import AppConfig, get_app_config
from deerflow.persistence.postgres_schema import create_schema_sql, dsn_with_search_path, normalize_libpq_dsn
from deerflow.runtime.checkpointer.provider import (
POSTGRES_CONN_REQUIRED,
POSTGRES_INSTALL,
@ -47,26 +48,43 @@ def _prepare_database_sqlite_checkpointer_path(db_config) -> str:
return conn_str
def _build_postgres_pool(conn_string: str):
def _build_postgres_pool(conn_string: str, schema: str = ""):
"""Build an AsyncConnectionPool with TCP keepalive and connection checking."""
from psycopg.rows import dict_row
from psycopg_pool import AsyncConnectionPool
kwargs = {
"autocommit": True,
"prepare_threshold": 0,
"row_factory": dict_row,
"keepalives": 1,
"keepalives_idle": 60,
"keepalives_interval": 10,
"keepalives_count": 6,
}
# Inject search_path into the DSN (merging with any libpq options already in
# the conn string) rather than via kwargs["options"], which psycopg applies
# *on top of* the conninfo and would silently drop a DSN-supplied option
# such as statement_timeout. This also strips a SQLAlchemy ``+driver``
# suffix so libpq can parse the DSN. Matches the sync/DSN paths.
dsn = dsn_with_search_path(normalize_libpq_dsn(conn_string), schema)
return AsyncConnectionPool(
conn_string,
kwargs={
"autocommit": True,
"prepare_threshold": 0,
"row_factory": dict_row,
"keepalives": 1,
"keepalives_idle": 60,
"keepalives_interval": 10,
"keepalives_count": 6,
},
dsn,
kwargs=kwargs,
check=AsyncConnectionPool.check_connection,
)
async def _ensure_postgres_schema_with_pool(pool, schema: str) -> None:
"""Create the configured schema before LangGraph creates its tables."""
statement = create_schema_sql(schema)
if statement is None:
return
async with pool.connection() as conn:
await conn.execute(statement)
def _ensure_postgres_imports():
"""Import and return (AsyncPostgresSaver, AsyncConnectionPool), raising ImportError on failure."""
try:
@ -113,8 +131,9 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
raise ValueError(POSTGRES_CONN_REQUIRED)
AsyncPostgresSaver, _ = _ensure_postgres_imports()
pool = _build_postgres_pool(config.connection_string)
pool = _build_postgres_pool(config.connection_string, config.postgres_schema)
async with pool:
await _ensure_postgres_schema_with_pool(pool, config.postgres_schema)
saver = AsyncPostgresSaver(conn=pool)
await saver.setup()
yield saver
@ -154,8 +173,9 @@ async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpoi
raise ValueError("database.postgres_url is required for the postgres backend")
AsyncPostgresSaver, _ = _ensure_postgres_imports()
pool = _build_postgres_pool(db_config.postgres_url)
pool = _build_postgres_pool(db_config.postgres_url, db_config.postgres_schema)
async with pool:
await _ensure_postgres_schema_with_pool(pool, db_config.postgres_schema)
saver = AsyncPostgresSaver(conn=pool)
await saver.setup()
yield saver

View File

@ -28,6 +28,7 @@ from langgraph.types import Checkpointer
from deerflow.config.app_config import AppConfig, get_app_config
from deerflow.config.checkpointer_config import CheckpointerConfig, ensure_config_loaded, get_checkpointer_config
from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
logger = logging.getLogger(__name__)
@ -43,6 +44,11 @@ POSTGRES_INSTALL = (
POSTGRES_CONN_REQUIRED = "checkpointer.connection_string is required for the postgres backend"
def _ensure_postgres_schema(conn_string: str, schema: str) -> None:
"""Create the configured schema before LangGraph creates its tables."""
ensure_postgres_schema(conn_string, schema, install_hint=POSTGRES_INSTALL)
# ---------------------------------------------------------------------------
# Config resolution
# ---------------------------------------------------------------------------
@ -68,7 +74,7 @@ def _resolve_checkpointer_config(app_config: AppConfig) -> CheckpointerConfig:
if database.backend == "postgres":
if not database.postgres_url:
raise ValueError("database.postgres_url is required for the postgres backend")
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url)
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url, postgres_schema=database.postgres_schema)
raise ValueError(f"Unknown database backend: {database.backend!r}")
@ -131,7 +137,9 @@ def _sync_checkpointer_cm(config: CheckpointerConfig) -> Iterator[Checkpointer]:
if not config.connection_string:
raise ValueError(POSTGRES_CONN_REQUIRED)
with PostgresSaver.from_conn_string(config.connection_string) as saver:
_ensure_postgres_schema(config.connection_string, config.postgres_schema)
conn_string = dsn_with_search_path(config.connection_string, config.postgres_schema)
with PostgresSaver.from_conn_string(conn_string) as saver:
saver.setup()
logger.info("Checkpointer: using PostgresSaver")
yield saver

View File

@ -25,6 +25,7 @@ from collections.abc import AsyncIterator
from langgraph.store.base import BaseStore
from deerflow.config.app_config import AppConfig, get_app_config
from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema_async
from deerflow.runtime.store.provider import (
POSTGRES_CONN_REQUIRED,
POSTGRES_STORE_INSTALL,
@ -36,6 +37,12 @@ from deerflow.runtime.store.provider import (
logger = logging.getLogger(__name__)
async def _ensure_postgres_schema(conn_string: str, schema: str) -> None:
"""Create the configured schema before LangGraph creates its store tables."""
await ensure_postgres_schema_async(conn_string, schema, install_hint=POSTGRES_STORE_INSTALL)
# ---------------------------------------------------------------------------
# Internal backend factory
# ---------------------------------------------------------------------------
@ -79,7 +86,9 @@ async def _async_store(config) -> AsyncIterator[BaseStore]:
if not config.connection_string:
raise ValueError(POSTGRES_CONN_REQUIRED)
async with AsyncPostgresStore.from_conn_string(config.connection_string) as store:
await _ensure_postgres_schema(config.connection_string, config.postgres_schema)
conn_string = dsn_with_search_path(config.connection_string, config.postgres_schema)
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
await store.setup()
logger.info("Store: using AsyncPostgresStore")
yield store

View File

@ -30,6 +30,7 @@ from langgraph.store.base import BaseStore
from deerflow.config.app_config import AppConfig, get_app_config
from deerflow.config.checkpointer_config import CheckpointerConfig, ensure_config_loaded, get_checkpointer_config
from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
logger = logging.getLogger(__name__)
@ -45,12 +46,19 @@ POSTGRES_STORE_INSTALL = (
POSTGRES_CONN_REQUIRED = "checkpointer.connection_string is required for the postgres backend"
def _ensure_postgres_schema(conn_string: str, schema: str) -> None:
"""Create the configured schema before LangGraph creates its store tables."""
ensure_postgres_schema(conn_string, schema, install_hint=POSTGRES_STORE_INSTALL)
def _resolve_store_config(app_config: AppConfig) -> CheckpointerConfig:
"""Resolve the Store backend from legacy or unified application config.
The legacy ``checkpointer`` section remains authoritative when present so
Store and Checkpointer continue to use the same backend. Otherwise the
unified ``database`` section drives the Store as documented.
unified ``database`` section drives the Store as documented. The unified
``postgres_schema`` is forwarded so Store tables land in the configured
schema alongside the checkpointer and app tables.
"""
if app_config.checkpointer is not None:
return app_config.checkpointer
@ -63,7 +71,7 @@ def _resolve_store_config(app_config: AppConfig) -> CheckpointerConfig:
if database.backend == "postgres":
if not database.postgres_url:
raise ValueError("database.postgres_url is required for the postgres backend")
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url)
return CheckpointerConfig(type="postgres", connection_string=database.postgres_url, postgres_schema=database.postgres_schema)
raise ValueError(f"Unknown database backend: {database.backend!r}")
@ -126,7 +134,9 @@ def _sync_store_cm(config) -> Iterator[BaseStore]:
if not config.connection_string:
raise ValueError(POSTGRES_CONN_REQUIRED)
with PostgresStore.from_conn_string(config.connection_string) as store:
_ensure_postgres_schema(config.connection_string, config.postgres_schema)
conn_string = dsn_with_search_path(config.connection_string, config.postgres_schema)
with PostgresStore.from_conn_string(conn_string) as store:
store.setup()
logger.info("Store: using PostgresStore")
yield store

View File

@ -1,9 +1,9 @@
"""Regression anchor: updating MCP config must not block the event loop.
``update_mcp_configuration`` resolves the extensions config path, probes its
existence, reads the raw JSON, writes the merged config, and reloads it all
blocking filesystem IO. The handler offloads the whole read-modify-write via
``asyncio.to_thread``; if it regresses back onto the event loop, the strict
The PUT and PATCH handlers resolve the extensions config path, probe its
existence, read raw JSON, atomically write it, and reload it all blocking
filesystem IO. They offload the whole read-modify-write via
``asyncio.to_thread``; if either regresses back onto the event loop, the strict
Blockbuster gate raises ``BlockingError`` and this test fails.
The admin check is patched to a no-op so the anchor exercises the handler's own
@ -22,7 +22,13 @@ from types import SimpleNamespace
import pytest
from app.gateway.routers import mcp as mcp_router
from app.gateway.routers.mcp import McpConfigUpdateRequest, McpServerConfigResponse, update_mcp_configuration
from app.gateway.routers.mcp import (
McpConfigUpdateRequest,
McpServerConfigResponse,
McpServerStateUpdateRequest,
update_mcp_configuration,
update_mcp_server_state,
)
pytestmark = pytest.mark.asyncio
@ -52,15 +58,40 @@ async def test_update_mcp_configuration_does_not_block_event_loop(tmp_path: Path
assert await asyncio.to_thread(config_path.exists)
async def test_concurrent_mcp_updates_are_serialized(tmp_path: Path, monkeypatch) -> None:
async def test_update_mcp_server_state_does_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
config_path = tmp_path / "extensions_config.json"
await asyncio.to_thread(
config_path.write_text,
'{"mcpServers":{"remote":{"enabled":false,"transport":"http","url":"https://example.test/mcp"}},"skills":{}}',
encoding="utf-8",
)
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(config_path))
async def _noop_admin(_request, **_kwargs) -> None:
return None
monkeypatch.setattr(mcp_router, "require_admin_user", _noop_admin)
monkeypatch.setattr(mcp_router, "reset_mcp_tools_cache", lambda: None)
response = await update_mcp_server_state(
request=None,
body=McpServerStateUpdateRequest(server_name="remote", enabled=True),
)
assert response.mcp_servers["remote"].enabled is True
assert response.mcp_servers["remote"].type == "http"
async def test_concurrent_mcp_put_and_patch_updates_are_serialized(tmp_path: Path, monkeypatch) -> None:
"""The write lock keeps the offloaded read-modify-write atomic within the process.
Offloading the RMW to a worker thread dropped the implicit serialization the
single-threaded event loop provided. ``extensions_config_write_lock`` restores
it and, being shared with the skills router (the other writer of this file),
also serializes against skill toggles: even with several concurrent
``PUT /api/mcp/config`` calls, only one RMW is inside the critical section at a
time. (Without the lock the tracked max concurrency would exceed 1.)
mix of ``PUT /api/mcp/config`` and ``PATCH /api/mcp/config`` calls, only one
RMW is inside the critical section at a time. (Without the lock the tracked
max concurrency would exceed 1.)
The tracker is injected *inside* the real worker (via ``reload_extensions_config``,
the last step under the lock) rather than replacing ``_apply_mcp_config_update``,
@ -68,7 +99,11 @@ async def test_concurrent_mcp_updates_are_serialized(tmp_path: Path, monkeypatch
the very thing under test.
"""
config_path = tmp_path / "extensions_config.json"
await asyncio.to_thread(config_path.write_text, '{"mcpServers": {}, "skills": {}}', encoding="utf-8")
await asyncio.to_thread(
config_path.write_text,
'{"mcpServers":{"s":{"enabled":true,"type":"http","url":"https://example.test/mcp"}},"skills":{}}',
encoding="utf-8",
)
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(config_path))
async def _noop_admin(_request, **_kwargs) -> None:
@ -96,6 +131,12 @@ async def test_concurrent_mcp_updates_are_serialized(tmp_path: Path, monkeypatch
mcp_servers={"s": McpServerConfigResponse(type="http", url="https://example.test/mcp")},
)
await asyncio.gather(*[update_mcp_configuration(request=None, body=body) for _ in range(5)])
await asyncio.gather(
*[update_mcp_configuration(request=None, body=body) for _ in range(4)],
update_mcp_server_state(
request=None,
body=McpServerStateUpdateRequest(server_name="s", enabled=False),
),
)
assert counters["max"] == 1, f"config updates were not serialized (max concurrency {counters['max']})"

View File

@ -20,8 +20,8 @@ performs the same RMW on the same file.
go back to separate module-local locks.
Only the config-infra boundaries (storage / ``get_extensions_config`` / reload /
path resolution) are stubbed; the real ``open(config_path, "w")`` write to a tmp
file is exercised.
path resolution) are stubbed; the real same-directory temporary write and atomic
replacement are exercised.
"""
from __future__ import annotations
@ -141,7 +141,7 @@ async def test_update_skill_serializes_concurrent_writes(tmp_path: Path, monkeyp
update_skill("demo-skill", SkillUpdateRequest(enabled=True), _admin_request(), SimpleNamespace()),
)
# The shared asyncio.Lock must serialize the offloaded read-modify-write.
# The shared threading.Lock must serialize the offloaded read-modify-write.
assert counters["max"] == 1

View File

@ -646,6 +646,56 @@ def test_get_app_config_keeps_persistence_runtime_singletons_when_checkpointer_u
_reset_config_singletons()
def test_get_app_config_does_not_reset_persistence_singletons_when_database_changes(tmp_path, monkeypatch):
# ``database`` is a restart-required field: the ORM engine is built once at
# startup and never rebuilt on a config.yaml edit. Resetting only the sync
# checkpointer/store singletons on a live ``postgres_schema`` change would
# half-migrate the deployment (new checkpoint/store tables in the new schema,
# ORM rows still in the old one). So a ``database`` change must NOT trigger a
# partial reset -- the operator must restart.
config_path = tmp_path / "config.yaml"
extensions_path = tmp_path / "extensions_config.json"
_write_extensions_config(extensions_path)
_write_config_with_sections(
config_path,
{"database": {"backend": "postgres", "postgres_url": "postgresql://localhost/db", "postgres_schema": "schema_a"}},
)
monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(config_path))
monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(extensions_path))
_reset_config_singletons()
reset_calls = {"checkpointer": 0, "store": 0}
def _reset_checkpointer() -> None:
reset_calls["checkpointer"] += 1
def _reset_store() -> None:
reset_calls["store"] += 1
monkeypatch.setattr("deerflow.runtime.checkpointer.reset_checkpointer", _reset_checkpointer)
monkeypatch.setattr("deerflow.runtime.store.reset_store", _reset_store)
try:
get_app_config()
reset_calls["checkpointer"] = 0
reset_calls["store"] = 0
_write_config_with_sections(
config_path,
{"database": {"backend": "postgres", "postgres_url": "postgresql://localhost/db", "postgres_schema": "schema_b"}},
)
next_mtime = config_path.stat().st_mtime + 5
os.utime(config_path, (next_mtime, next_mtime))
get_app_config()
assert get_checkpointer_config() is None
assert reset_calls == {"checkpointer": 0, "store": 0}
finally:
_reset_config_singletons()
def test_get_app_config_does_not_mutate_singletons_when_reload_validation_fails(tmp_path, monkeypatch):
config_path = tmp_path / "config.yaml"
extensions_path = tmp_path / "extensions_config.json"

View File

@ -1130,8 +1130,13 @@ def test_authenticate_skips_rehash_for_v2_hash():
mock_repo.update_user.assert_not_called()
def test_validate_next_param_rejects_colon_paths():
def test_validate_next_param_rejects_unsafe_paths():
from app.gateway.routers.auth import validate_next_param
assert validate_next_param("/workspace") == "/workspace"
assert validate_next_param("/workspace/chats/new?tab=recent#top") == "/workspace/chats/new?tab=recent#top"
assert validate_next_param("/:evil") is None
assert validate_next_param("/\\evil.example") is None
assert validate_next_param("/foo\\bar") is None
assert validate_next_param("//evil.example") is None
assert validate_next_param("https://evil.example") is None

View File

@ -144,6 +144,18 @@ class TestCheckpointerConfig:
assert config is not None
assert config.type == "postgres"
assert config.connection_string == "postgresql://localhost/db"
assert config.postgres_schema == ""
def test_postgres_schema_accepts_valid_identifier(self):
config = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema="deerflow")
assert config.postgres_schema == "deerflow"
@pytest.mark.parametrize("schema", ["1abc", "a b", "a;b", "a-b", "a" * 64, 'a"b', "MySchema", "Orders", "Public"])
def test_postgres_schema_rejects_invalid_identifier(self, schema):
from pydantic import ValidationError
with pytest.raises(ValidationError):
CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema=schema)
def test_default_connection_string_is_none(self):
config = CheckpointerConfig(type="memory")
@ -393,6 +405,76 @@ class TestGetCheckpointer:
mock_saver_cls.from_conn_string.assert_called_once_with("postgresql://localhost/db")
mock_saver_instance.setup.assert_called_once()
def test_postgres_schema_creates_schema_and_sets_search_path(self):
"""Sync Postgres checkpointer should create schema before setup."""
load_checkpointer_config_from_dict({"type": "postgres", "connection_string": "postgresql://localhost/db", "postgres_schema": "deerflow"})
mock_saver_instance = MagicMock()
mock_cm = MagicMock()
mock_cm.__enter__ = MagicMock(return_value=mock_saver_instance)
mock_cm.__exit__ = MagicMock(return_value=False)
mock_saver_cls = MagicMock()
mock_saver_cls.from_conn_string = MagicMock(return_value=mock_cm)
mock_pg_module = MagicMock()
mock_pg_module.PostgresSaver = mock_saver_cls
mock_conn = MagicMock()
mock_psycopg = MagicMock()
mock_psycopg.connect.return_value = mock_conn
with (
patch.dict(sys.modules, {"langgraph.checkpoint.postgres": mock_pg_module}),
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
):
reset_checkpointer()
cp = get_checkpointer()
assert cp is mock_saver_instance
mock_psycopg.connect.assert_called_once_with("postgresql://localhost/db", autocommit=True)
mock_conn.execute.assert_called_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
# psycopg 3 __exit__ does not close(); the sync path must close explicitly.
mock_conn.close.assert_called_once_with()
called_dsn = mock_saver_cls.from_conn_string.call_args.args[0]
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
mock_saver_instance.setup.assert_called_once()
def test_store_postgres_schema_creates_schema_and_sets_search_path(self):
"""Sync Postgres store should use the legacy checkpointer schema."""
load_checkpointer_config_from_dict({"type": "postgres", "connection_string": "postgresql://localhost/db", "postgres_schema": "deerflow"})
mock_store_instance = MagicMock()
mock_cm = MagicMock()
mock_cm.__enter__ = MagicMock(return_value=mock_store_instance)
mock_cm.__exit__ = MagicMock(return_value=False)
mock_store_cls = MagicMock()
mock_store_cls.from_conn_string = MagicMock(return_value=mock_cm)
mock_pg_module = MagicMock()
mock_pg_module.PostgresStore = mock_store_cls
mock_conn = MagicMock()
mock_psycopg = MagicMock()
mock_psycopg.connect.return_value = mock_conn
with (
patch.dict(sys.modules, {"langgraph.store.postgres": mock_pg_module}),
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
):
reset_store()
store = get_store()
assert store is mock_store_instance
mock_psycopg.connect.assert_called_once_with("postgresql://localhost/db", autocommit=True)
mock_conn.execute.assert_called_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
# psycopg 3 __exit__ does not close(); the sync path must close explicitly.
mock_conn.close.assert_called_once_with()
called_dsn = mock_store_cls.from_conn_string.call_args.args[0]
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
mock_store_instance.setup.assert_called_once()
class TestSyncSingletonThreadSafety:
def test_store_reset_clears_singleton(self):
@ -576,7 +658,7 @@ class TestAsyncCheckpointer:
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
mock_config = MagicMock()
mock_config.checkpointer = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db")
mock_config.checkpointer = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema="deerflow")
mock_saver = AsyncMock()
@ -585,6 +667,11 @@ class TestAsyncCheckpointer:
mock_pool_instance = AsyncMock()
mock_pool_instance.__aenter__.return_value = mock_pool_instance
mock_pool_instance.__aexit__.return_value = False
mock_conn = AsyncMock()
mock_conn_cm = AsyncMock()
mock_conn_cm.__aenter__.return_value = mock_conn
mock_conn_cm.__aexit__.return_value = False
mock_pool_instance.connection = MagicMock(return_value=mock_conn_cm)
mock_pool_cls = MagicMock(return_value=mock_pool_instance)
mock_pool_cls.check_connection = AsyncMock()
@ -610,8 +697,12 @@ class TestAsyncCheckpointer:
# Verify the pool was constructed with check Connection
mock_pool_cls.assert_called_once()
call_kwargs = mock_pool_cls.call_args
assert call_kwargs[0][0] == "postgresql://localhost/db"
# search_path is injected into the DSN (merged with any existing libpq
# options), not via kwargs["options"] which would clobber DSN options.
assert "options=-c%20search_path%3Ddeerflow" in call_kwargs[0][0]
assert call_kwargs[1]["check"] is mock_pool_cls.check_connection
assert "options" not in call_kwargs[1]["kwargs"]
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
# Verify saver was constructed with the pool (not via from_conn_string)
mock_saver_cls.assert_called_once_with(conn=mock_pool_instance)
@ -623,7 +714,7 @@ class TestAsyncCheckpointer:
from deerflow.config.database_config import DatabaseConfig
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
db_config = DatabaseConfig(backend="postgres", postgres_url="postgresql://localhost/db")
db_config = DatabaseConfig(backend="postgres", postgres_url="postgresql://localhost/db", postgres_schema="deerflow")
mock_config = MagicMock()
mock_config.checkpointer = None
mock_config.database = db_config
@ -635,6 +726,11 @@ class TestAsyncCheckpointer:
mock_pool_instance = AsyncMock()
mock_pool_instance.__aenter__.return_value = mock_pool_instance
mock_pool_instance.__aexit__.return_value = False
mock_conn = AsyncMock()
mock_conn_cm = AsyncMock()
mock_conn_cm.__aenter__.return_value = mock_conn
mock_conn_cm.__aexit__.return_value = False
mock_pool_instance.connection = MagicMock(return_value=mock_conn_cm)
mock_pool_cls = MagicMock(return_value=mock_pool_instance)
mock_pool_cls.check_connection = AsyncMock()
@ -657,8 +753,10 @@ class TestAsyncCheckpointer:
mock_pool_cls.assert_called_once()
call_kwargs = mock_pool_cls.call_args
assert call_kwargs[0][0] == "postgresql://localhost/db"
assert "options=-c%20search_path%3Ddeerflow" in call_kwargs[0][0]
assert call_kwargs[1]["check"] is mock_pool_cls.check_connection
assert "options" not in call_kwargs[1]["kwargs"]
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
mock_saver_cls.assert_called_once_with(conn=mock_pool_instance)
mock_saver.setup.assert_awaited_once()
@ -705,6 +803,88 @@ class TestAsyncCheckpointer:
mock_saver.setup.assert_awaited_once()
class TestAsyncStore:
@pytest.mark.anyio
async def test_postgres_schema_creates_schema_and_sets_search_path(self):
"""Async Postgres store should use the legacy checkpointer schema."""
from deerflow.runtime.store.async_provider import make_store
mock_config = MagicMock()
mock_config.checkpointer = CheckpointerConfig(type="postgres", connection_string="postgresql://localhost/db", postgres_schema="deerflow")
mock_store = AsyncMock()
mock_cm = AsyncMock()
mock_cm.__aenter__.return_value = mock_store
mock_cm.__aexit__.return_value = False
mock_store_cls = MagicMock()
mock_store_cls.from_conn_string.return_value = mock_cm
mock_pg_module = MagicMock()
mock_pg_module.AsyncPostgresStore = mock_store_cls
mock_conn = AsyncMock()
mock_async_connection = MagicMock()
mock_async_connection.connect = AsyncMock(return_value=mock_conn)
mock_psycopg = MagicMock()
mock_psycopg.AsyncConnection = mock_async_connection
with (
patch.dict(sys.modules, {"langgraph.store.postgres.aio": mock_pg_module}),
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
):
async with make_store(mock_config) as store:
assert store is mock_store
mock_async_connection.connect.assert_awaited_once_with("postgresql://localhost/db", autocommit=True)
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
mock_conn.close.assert_awaited_once()
called_dsn = mock_store_cls.from_conn_string.call_args.args[0]
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
mock_store.setup.assert_awaited_once()
@pytest.mark.anyio
async def test_database_postgres_schema_creates_schema_and_sets_search_path(self):
"""Unified database postgres store should use database.postgres_schema."""
from deerflow.config.database_config import DatabaseConfig
from deerflow.runtime.store.async_provider import make_store
mock_config = MagicMock()
mock_config.checkpointer = None
mock_config.database = DatabaseConfig(backend="postgres", postgres_url="postgresql://localhost/db", postgres_schema="deerflow")
mock_store = AsyncMock()
mock_cm = AsyncMock()
mock_cm.__aenter__.return_value = mock_store
mock_cm.__aexit__.return_value = False
mock_store_cls = MagicMock()
mock_store_cls.from_conn_string.return_value = mock_cm
mock_pg_module = MagicMock()
mock_pg_module.AsyncPostgresStore = mock_store_cls
mock_conn = AsyncMock()
mock_async_connection = MagicMock()
mock_async_connection.connect = AsyncMock(return_value=mock_conn)
mock_psycopg = MagicMock()
mock_psycopg.AsyncConnection = mock_async_connection
with (
patch.dict(sys.modules, {"langgraph.store.postgres.aio": mock_pg_module}),
patch.dict(sys.modules, {"psycopg": mock_psycopg}),
):
async with make_store(mock_config) as store:
assert store is mock_store
mock_async_connection.connect.assert_awaited_once_with("postgresql://localhost/db", autocommit=True)
mock_conn.execute.assert_awaited_once_with('CREATE SCHEMA IF NOT EXISTS "deerflow"')
mock_conn.close.assert_awaited_once()
called_dsn = mock_store_cls.from_conn_string.call_args.args[0]
assert "options=-c%20search_path%3Ddeerflow" in called_dsn
mock_store.setup.assert_awaited_once()
class TestCheckpointerDatabaseConfig:
"""The sync checkpointer must follow the unified ``database`` section when no
legacy ``checkpointer`` section is configured matching the async

View File

@ -0,0 +1,144 @@
"""Regression tests for crash-safe extensions config writes."""
from __future__ import annotations
import json
import os
import stat
from pathlib import Path
import pytest
from deerflow.config import extensions_config as extensions_config_module
from deerflow.config.extensions_config import atomic_write_extensions_config
def _temporary_files_for(path: Path) -> list[Path]:
return list(path.parent.glob(f".{path.name}.*.tmp"))
def test_atomic_write_replaces_config_without_leaving_temp_files(tmp_path: Path) -> None:
config_path = tmp_path / "extensions_config.json"
config_path.write_text('{"old": true}', encoding="utf-8")
atomic_write_extensions_config(
config_path,
{
"mcpServers": {"github": {"enabled": False}},
"skills": {"research": {"enabled": True}},
},
)
assert json.loads(config_path.read_text(encoding="utf-8")) == {
"mcpServers": {"github": {"enabled": False}},
"skills": {"research": {"enabled": True}},
}
assert _temporary_files_for(config_path) == []
def test_atomic_write_preserves_original_when_json_dump_fails_mid_write(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
config_path = tmp_path / "extensions_config.json"
original = '{"mcpServers": {"github": {"enabled": true}}, "skills": {}}'
config_path.write_text(original, encoding="utf-8")
def fail_after_partial_write(_data, file_handle, **_kwargs) -> None:
file_handle.write('{"mcpServers":')
file_handle.flush()
raise OSError("disk full")
monkeypatch.setattr(extensions_config_module.json, "dump", fail_after_partial_write)
with pytest.raises(OSError, match="disk full"):
atomic_write_extensions_config(
config_path,
{"mcpServers": {"github": {"enabled": False}}, "skills": {}},
)
assert config_path.read_text(encoding="utf-8") == original
assert _temporary_files_for(config_path) == []
def test_atomic_write_preserves_original_when_replace_fails(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
config_path = tmp_path / "extensions_config.json"
original = '{"mcpServers": {}, "skills": {}}'
config_path.write_text(original, encoding="utf-8")
def fail_replace(_source, _destination) -> None:
raise OSError("replace failed")
monkeypatch.setattr(extensions_config_module.os, "replace", fail_replace)
with pytest.raises(OSError, match="replace failed"):
atomic_write_extensions_config(
config_path,
{"mcpServers": {"github": {"enabled": True}}, "skills": {}},
)
assert config_path.read_text(encoding="utf-8") == original
assert _temporary_files_for(config_path) == []
def test_atomic_write_preserves_original_when_file_fsync_fails(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
config_path = tmp_path / "extensions_config.json"
original = '{"mcpServers": {}, "skills": {}}'
config_path.write_text(original, encoding="utf-8")
def fail_fsync(_file_descriptor) -> None:
raise OSError("fsync failed")
monkeypatch.setattr(extensions_config_module.os, "fsync", fail_fsync)
with pytest.raises(OSError, match="fsync failed"):
atomic_write_extensions_config(
config_path,
{"mcpServers": {"github": {"enabled": True}}, "skills": {}},
)
assert config_path.read_text(encoding="utf-8") == original
assert _temporary_files_for(config_path) == []
@pytest.mark.skipif(os.name == "nt", reason="POSIX mode bits unavailable")
def test_atomic_write_preserves_existing_file_mode(tmp_path: Path) -> None:
config_path = tmp_path / "extensions_config.json"
config_path.write_text('{"mcpServers": {}, "skills": {}}', encoding="utf-8")
config_path.chmod(0o640)
atomic_write_extensions_config(
config_path,
{"mcpServers": {}, "skills": {"research": {"enabled": False}}},
)
assert stat.S_IMODE(config_path.stat().st_mode) == 0o640
def test_atomic_write_updates_symlink_target_without_replacing_symlink(
tmp_path: Path,
) -> None:
target_path = tmp_path / "actual-extensions-config.json"
target_path.write_text('{"mcpServers": {}, "skills": {}}', encoding="utf-8")
config_path = tmp_path / "extensions_config.json"
try:
config_path.symlink_to(target_path)
except OSError as error:
pytest.skip(f"Symlinks are unavailable: {error}")
atomic_write_extensions_config(
config_path,
{"mcpServers": {"github": {"enabled": False}}, "skills": {}},
)
assert config_path.is_symlink()
assert json.loads(target_path.read_text(encoding="utf-8")) == {
"mcpServers": {"github": {"enabled": False}},
"skills": {},
}

View File

@ -3,6 +3,8 @@
Verifies that GET /api/mcp/config masks sensitive fields (env values,
header values, OAuth secrets) and that PUT /api/mcp/config correctly
preserves existing secrets when the frontend round-trips masked values.
PATCH /api/mcp/config coverage pins targeted state changes, raw-config
preservation, transport aliases, authorization, and command validation.
"""
from __future__ import annotations
@ -21,13 +23,15 @@ from app.gateway.routers.mcp import (
McpConfigUpdateRequest,
McpOAuthConfigResponse,
McpServerConfigResponse,
McpServerStateUpdateRequest,
_mask_server_config,
_merge_preserving_secrets,
_validate_mcp_update_request,
reset_mcp_tools_cache_endpoint,
update_mcp_configuration,
update_mcp_server_state,
)
from deerflow.config.extensions_config import ExtensionsConfig
from deerflow.config.extensions_config import ExtensionsConfig, McpServerConfig
# ---------------------------------------------------------------------------
# _mask_server_config
@ -589,6 +593,202 @@ async def test_update_mcp_configuration_preserves_server_extra_fields(monkeypatc
assert response.mcp_servers["playwright"].model_extra["api_key"] == "***"
@pytest.mark.asyncio
@pytest.mark.parametrize("enabled", [False, True])
async def test_update_mcp_server_state_updates_valid_target_despite_unrelated_disallowed_command(
monkeypatch,
tmp_path,
enabled: bool,
):
config_path = tmp_path / "extensions_config.json"
original = {
"mcpServers": {
"semantic-scholar": {
"enabled": True,
"type": "stdio",
"command": "s2-mcp-server",
"env": {"S2_API_KEY": "$S2_API_KEY"},
"customFlag": "keep-me",
},
"github": {
"enabled": not enabled,
"type": "stdio",
"command": "npx",
"args": ["-y", "@modelcontextprotocol/server-github"],
},
},
"skills": {"research": {"enabled": False}},
"middlewares": ["example.middleware:Middleware"],
"customTopLevel": {"preserve": True},
}
config_path.write_text(json.dumps(original), encoding="utf-8")
reset_calls = 0
def fake_reload_extensions_config():
return ExtensionsConfig.model_validate(json.loads(config_path.read_text(encoding="utf-8")))
def fake_reset_mcp_tools_cache():
nonlocal reset_calls
reset_calls += 1
monkeypatch.setattr(mcp_router.ExtensionsConfig, "resolve_config_path", lambda: config_path)
monkeypatch.setattr(mcp_router, "reload_extensions_config", fake_reload_extensions_config)
monkeypatch.setattr(mcp_router, "reset_mcp_tools_cache", fake_reset_mcp_tools_cache)
monkeypatch.delenv(_MCP_STDIO_COMMAND_ALLOWLIST_ENV, raising=False)
response = await update_mcp_server_state(
_request_with_role("admin"),
McpServerStateUpdateRequest(server_name="github", enabled=enabled),
)
persisted = json.loads(config_path.read_text(encoding="utf-8"))
assert persisted["mcpServers"]["github"]["enabled"] is enabled
assert persisted["mcpServers"]["semantic-scholar"] == original["mcpServers"]["semantic-scholar"]
assert persisted["skills"] == original["skills"]
assert persisted["middlewares"] == original["middlewares"]
assert persisted["customTopLevel"] == original["customTopLevel"]
assert response.mcp_servers["github"].enabled is enabled
assert response.mcp_servers["semantic-scholar"].env == {"S2_API_KEY": "***"}
assert reset_calls == 1
@pytest.mark.asyncio
async def test_update_mcp_server_state_allows_disabling_but_rejects_enabling_disallowed_command(monkeypatch, tmp_path):
config_path = tmp_path / "extensions_config.json"
config_path.write_text(
json.dumps(
{
"mcpServers": {
"semantic-scholar": {
"enabled": True,
"type": "stdio",
"command": "s2-mcp-server",
}
},
"skills": {},
}
),
encoding="utf-8",
)
reset_calls = 0
def fake_reload_extensions_config():
return ExtensionsConfig.model_validate(json.loads(config_path.read_text(encoding="utf-8")))
def fake_reset_mcp_tools_cache():
nonlocal reset_calls
reset_calls += 1
monkeypatch.setattr(mcp_router.ExtensionsConfig, "resolve_config_path", lambda: config_path)
monkeypatch.setattr(mcp_router, "reload_extensions_config", fake_reload_extensions_config)
monkeypatch.setattr(mcp_router, "reset_mcp_tools_cache", fake_reset_mcp_tools_cache)
monkeypatch.delenv(_MCP_STDIO_COMMAND_ALLOWLIST_ENV, raising=False)
response = await update_mcp_server_state(
_request_with_role("admin"),
McpServerStateUpdateRequest(server_name="semantic-scholar", enabled=False),
)
assert response.mcp_servers["semantic-scholar"].enabled is False
with pytest.raises(HTTPException) as exc_info:
await update_mcp_server_state(
_request_with_role("admin"),
McpServerStateUpdateRequest(server_name="semantic-scholar", enabled=True),
)
assert exc_info.value.status_code == 400
assert "s2-mcp-server" in exc_info.value.detail
persisted = json.loads(config_path.read_text(encoding="utf-8"))
assert persisted["mcpServers"]["semantic-scholar"]["enabled"] is False
assert reset_calls == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["sse", "http"])
async def test_update_mcp_server_state_enables_raw_transport_alias(
monkeypatch,
tmp_path,
transport: str,
):
config_path = tmp_path / "extensions_config.json"
original_server = {
"enabled": False,
"transport": transport,
"url": "https://mcp.example.com/mcp",
"customFlag": "keep-me",
}
config_path.write_text(
json.dumps(
{
"mcpServers": {"remote": original_server},
"skills": {},
}
),
encoding="utf-8",
)
reset_calls = 0
def fake_reload_extensions_config():
return ExtensionsConfig.model_validate(json.loads(config_path.read_text(encoding="utf-8")))
def fake_reset_mcp_tools_cache():
nonlocal reset_calls
reset_calls += 1
monkeypatch.setattr(mcp_router.ExtensionsConfig, "resolve_config_path", lambda: config_path)
monkeypatch.setattr(mcp_router, "reload_extensions_config", fake_reload_extensions_config)
monkeypatch.setattr(mcp_router, "reset_mcp_tools_cache", fake_reset_mcp_tools_cache)
monkeypatch.delenv(_MCP_STDIO_COMMAND_ALLOWLIST_ENV, raising=False)
response = await update_mcp_server_state(
_request_with_role("admin"),
McpServerStateUpdateRequest(server_name="remote", enabled=True),
)
persisted_server = json.loads(config_path.read_text(encoding="utf-8"))["mcpServers"]["remote"]
assert persisted_server == {**original_server, "enabled": True}
assert "type" not in persisted_server
assert response.mcp_servers["remote"].enabled is True
assert response.mcp_servers["remote"].type == transport
assert reset_calls == 1
@pytest.mark.asyncio
async def test_update_mcp_server_state_returns_404_without_writing_or_resetting_cache(monkeypatch, tmp_path):
config_path = tmp_path / "extensions_config.json"
original_text = '{"mcpServers": {}, "skills": {}}'
config_path.write_text(original_text, encoding="utf-8")
reset_calls = 0
def fake_reset_mcp_tools_cache():
nonlocal reset_calls
reset_calls += 1
monkeypatch.setattr(mcp_router.ExtensionsConfig, "resolve_config_path", lambda: config_path)
monkeypatch.setattr(mcp_router, "reset_mcp_tools_cache", fake_reset_mcp_tools_cache)
with pytest.raises(HTTPException) as exc_info:
await update_mcp_server_state(
_request_with_role("admin"),
McpServerStateUpdateRequest(server_name="missing", enabled=True),
)
assert exc_info.value.status_code == 404
assert config_path.read_text(encoding="utf-8") == original_text
assert reset_calls == 0
@pytest.mark.asyncio
async def test_update_mcp_server_state_requires_admin():
with pytest.raises(HTTPException) as exc_info:
await update_mcp_server_state(
_request_with_role("user"),
McpServerStateUpdateRequest(server_name="github", enabled=False),
)
assert exc_info.value.status_code == 403
def test_validate_mcp_update_allows_default_npx_stdio_command(monkeypatch):
monkeypatch.delenv(_MCP_STDIO_COMMAND_ALLOWLIST_ENV, raising=False)
request = McpConfigUpdateRequest(
@ -689,3 +889,48 @@ def test_validate_mcp_update_ignores_remote_transports(monkeypatch):
)
_validate_mcp_update_request(request)
@pytest.mark.parametrize(
("raw_server", "expected_type"),
[
({"transport": "sse", "url": "https://mcp.example.com/sse"}, "sse"),
({"transport": "http", "url": "https://mcp.example.com/mcp"}, "http"),
({"transport": "stdio", "command": "npx"}, "stdio"),
({"type": "http", "transport": "sse", "url": "https://mcp.example.com/mcp"}, "http"),
({}, "stdio"),
],
)
def test_api_and_runtime_mcp_models_normalize_transport_consistently(
raw_server: dict[str, object],
expected_type: str,
):
api_server = McpServerConfigResponse.model_validate(raw_server)
runtime_server = McpServerConfig.model_validate(raw_server)
assert api_server.type == expected_type
assert runtime_server.type == expected_type
assert api_server.type == runtime_server.type
if "transport" in raw_server:
assert api_server.model_extra["transport"] == raw_server["transport"]
assert runtime_server.model_extra["transport"] == raw_server["transport"]
def test_validate_mcp_update_enforces_stdio_transport_alias(monkeypatch):
monkeypatch.delenv(_MCP_STDIO_COMMAND_ALLOWLIST_ENV, raising=False)
request = McpConfigUpdateRequest.model_validate(
{
"mcp_servers": {
"disallowed": {
"transport": "stdio",
"command": "custom-mcp-server",
}
}
}
)
with pytest.raises(HTTPException) as exc_info:
_validate_mcp_update_request(request)
assert exc_info.value.status_code == 400
assert "custom-mcp-server" in exc_info.value.detail

View File

@ -203,6 +203,20 @@ class TestRewriteLocalPathsInText:
assert result == text
def test_oversized_path_like_text_is_left_untouched(self, paths: Paths):
workspace = paths.sandbox_work_dir("t1", user_id="u1")
text = f"手术室/重症监护室OR/ICU整体解决方案{'说明' * 200}"
with _patch_paths(paths):
result = mcp_tools._rewrite_local_paths_in_text(
text,
thread_id="t1",
user_id="u1",
source_base_dir=workspace,
)
assert result == text
def test_playwright_markdown_path_is_rewritten_twice_without_copy(self, paths: Paths):
workspace = paths.sandbox_work_dir("t1", user_id="u1")
_workspace_file(paths, ".playwright-mcp/page.png", content=b"png")

View File

@ -69,3 +69,34 @@ def test_escape_url_for_alembic_doubles_only_percent_signs() -> None:
# Idempotency is intentionally NOT a property -- doubling is one-way;
# callers must escape exactly once on the way into set_main_option.
assert _escape_url_for_alembic("a%%b") == "a%%%%b"
def test_alembic_config_forwards_postgres_schema_option() -> None:
# The custom schema must reach env.py so the alembic-spawned engine can
# pin its search_path; otherwise alembic_version + migration DDL land in
# ``public`` while the ORM tables land in the custom schema.
engine = _fake_engine("postgresql://a:b@h/d")
cfg = _get_alembic_config(engine, postgres_schema="deerflow")
assert cfg.get_main_option("deerflow_pg_schema") == "deerflow"
def test_alembic_config_omits_schema_option_when_unset() -> None:
engine = _fake_engine("postgresql://a:b@h/d")
cfg = _get_alembic_config(engine)
assert cfg.get_main_option("deerflow_pg_schema") is None
def test_env_module_pins_search_path_from_schema_option() -> None:
"""env.py must read ``deerflow_pg_schema`` and pin the alembic-spawned
engine's search_path. That engine is built from the bare URL and does not
inherit the app engine's asyncpg ``server_settings``, so without this both
``alembic_version`` and migration DDL would land in ``public`` while the
ORM tables land in the custom schema. Source-parity check mirroring
``test_env_module_wires_busy_timeout_for_sqlite``.
"""
from pathlib import Path
env_path = Path(__file__).resolve().parents[1] / "packages/harness/deerflow/persistence/migrations/env.py"
src = env_path.read_text(encoding="utf-8")
assert 'get_main_option("deerflow_pg_schema")' in src, "env.py must read the deerflow_pg_schema option set by _get_alembic_config"
assert "build_asyncpg_connect_args" in src, "env.py must pin the alembic engine's search_path via build_asyncpg_connect_args"

View File

@ -145,7 +145,7 @@ async def test_init_engine_postgres_uses_hardened_kwargs() -> None:
await engine_mod.init_engine(backend="postgres", url=url, echo=True, pool_size=12)
create_engine.assert_called_once_with(url, **engine_mod._postgres_engine_kwargs(echo=True, pool_size=12))
bootstrap_schema.assert_awaited_once_with(mock_engine, backend="postgres")
bootstrap_schema.assert_awaited_once_with(mock_engine, backend="postgres", postgres_schema="")
finally:
await engine_mod.close_engine()
@ -174,7 +174,10 @@ async def test_init_engine_postgres_retry_uses_hardened_kwargs() -> None:
assert create_engine.call_args_list == [call(url, **kwargs), call(url, **kwargs)]
auto_create.assert_awaited_once_with(url)
initial_engine.dispose.assert_awaited_once()
assert bootstrap_schema.await_args_list == [call(initial_engine, backend="postgres"), call(retry_engine, backend="postgres")]
assert bootstrap_schema.await_args_list == [
call(initial_engine, backend="postgres", postgres_schema=""),
call(retry_engine, backend="postgres", postgres_schema=""),
]
finally:
await engine_mod.close_engine()

View File

@ -10,7 +10,7 @@ Tests:
import sys
from datetime import UTC, datetime
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
import pytest
@ -71,6 +71,50 @@ class TestDatabaseConfig:
with pytest.raises(ValueError, match="No SQLAlchemy URL"):
_ = c.app_sqlalchemy_url
def test_postgres_schema_default_empty(self):
c = DatabaseConfig()
assert c.postgres_schema == ""
@pytest.mark.parametrize("schema", ["deerflow", "my_schema", "_private", "s", "a" * 63])
def test_postgres_schema_accepts_valid_identifier(self, schema):
c = DatabaseConfig(backend="postgres", postgres_url="postgresql://u:p@h:5432/db", postgres_schema=schema)
assert c.postgres_schema == schema
@pytest.mark.parametrize(
"schema",
[
"1abc",
"a b",
"a;b",
"a-b",
"a" * 64,
'a"b',
"MySchema",
"Orders",
"Public",
# Trailing/leading whitespace must be rejected: a ``$``-anchored
# ``re.match`` accepts a single trailing ``\n``, which would create a
# quoted schema literally named ``deerflow\n`` while the unquoted
# search_path folds to ``deerflow`` and misses it (tables land in
# ``public``). ``re.fullmatch`` on an unanchored pattern rejects it.
"deerflow\n",
"deerflow\t",
"\ndeerflow",
"deerflow ",
],
)
def test_postgres_schema_rejects_invalid_identifier(self, schema):
from pydantic import ValidationError
with pytest.raises(ValidationError):
DatabaseConfig(backend="postgres", postgres_url="postgresql://u:p@h:5432/db", postgres_schema=schema)
def test_postgres_schema_does_not_pollute_url(self):
c = DatabaseConfig(backend="postgres", postgres_url="postgresql://u:p@h:5432/db", postgres_schema="deerflow")
url = c.app_sqlalchemy_url
assert "deerflow" not in url.replace("/db", "")
assert url.startswith("postgresql+asyncpg://")
# -- MemoryRunStore --
@ -276,20 +320,22 @@ class TestBaseToDictMixin:
name: Mapped[str] = mapped_column(String(128))
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'test.db'}")
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
try:
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
sf = async_sessionmaker(engine, expire_on_commit=False)
async with sf() as session:
session.add(_Tmp(id="1", name="hello"))
await session.commit()
obj = await session.get(_Tmp, "1")
sf = async_sessionmaker(engine, expire_on_commit=False)
async with sf() as session:
session.add(_Tmp(id="1", name="hello"))
await session.commit()
obj = await session.get(_Tmp, "1")
assert obj.to_dict() == {"id": "1", "name": "hello"}
assert obj.to_dict(exclude={"name"}) == {"id": "1"}
assert "_Tmp" in repr(obj)
await engine.dispose()
assert obj.to_dict() == {"id": "1", "name": "hello"}
assert obj.to_dict(exclude={"name"}) == {"id": "1"}
assert "_Tmp" in repr(obj)
finally:
await engine.dispose()
Base.metadata.remove(_Tmp.__table__)
# -- Engine lifecycle --
@ -327,3 +373,111 @@ class TestEngineLifecycle:
pytest.raises(ImportError, match="uv sync --all-packages --extra postgres"),
):
await init_engine("postgres", url="postgresql+asyncpg://x:x@localhost/x")
def _make_fake_pg_engine():
"""Build a fake async engine whose begin()/dispose() are awaitable mocks.
Tracks ordering of conn.execute (CREATE SCHEMA) vs conn.run_sync
(create_all) through a shared parent mock's ``mock_calls``.
"""
from unittest.mock import AsyncMock, MagicMock
calls = MagicMock()
conn = MagicMock()
# Both are awaited by the engine code, so they must return awaitables.
calls.execute = AsyncMock()
calls.run_sync = AsyncMock()
conn.execute = calls.execute
conn.run_sync = calls.run_sync
begin_cm = AsyncMock()
begin_cm.__aenter__.return_value = conn
begin_cm.__aexit__.return_value = False
engine = MagicMock()
engine.begin = MagicMock(return_value=begin_cm)
engine.dispose = AsyncMock()
return engine, calls
class TestPostgresSchemaInit:
@pytest.mark.anyio
async def test_passes_search_path_connect_args(self, monkeypatch):
import deerflow.persistence.engine as engine_module
monkeypatch.setitem(sys.modules, "asyncpg", object())
fake_engine, _calls = _make_fake_pg_engine()
captured = {}
def fake_create(_url, **kwargs):
captured.update(kwargs)
return fake_engine
monkeypatch.setattr(engine_module, "create_async_engine", fake_create)
monkeypatch.setattr("deerflow.persistence.bootstrap.bootstrap_schema", AsyncMock())
await engine_module.init_engine(
"postgres",
url="postgresql+asyncpg://u:p@h:5432/db",
postgres_schema="deerflow",
)
assert captured["connect_args"] == {
"command_timeout": engine_module.POSTGRES_COMMAND_TIMEOUT_SECONDS,
"server_settings": {"search_path": "deerflow"},
}
await engine_module.close_engine()
@pytest.mark.anyio
async def test_creates_schema_before_bootstrap(self, monkeypatch):
import deerflow.persistence.engine as engine_module
monkeypatch.setitem(sys.modules, "asyncpg", object())
fake_engine, calls = _make_fake_pg_engine()
monkeypatch.setattr(engine_module, "create_async_engine", lambda url, **kw: fake_engine)
calls.attach_mock(AsyncMock(), "bootstrap_schema")
monkeypatch.setattr("deerflow.persistence.bootstrap.bootstrap_schema", calls.bootstrap_schema)
await engine_module.init_engine(
"postgres",
url="postgresql+asyncpg://u:p@h:5432/db",
postgres_schema="deerflow",
)
names = [c[0] for c in calls.mock_calls]
assert "execute" in names
assert "bootstrap_schema" in names
# CREATE SCHEMA must run before the alembic bootstrap so the
# subsequent create_all / migration DDL lands in the target schema.
assert names.index("execute") < names.index("bootstrap_schema")
# The DDL passed to execute must be a CreateSchema for the target schema.
execute_arg = calls.execute.call_args[0][0]
assert "deerflow" in str(execute_arg)
await engine_module.close_engine()
@pytest.mark.anyio
async def test_empty_schema_skips_connect_args_and_ddl(self, monkeypatch):
import deerflow.persistence.engine as engine_module
monkeypatch.setitem(sys.modules, "asyncpg", object())
fake_engine, calls = _make_fake_pg_engine()
captured = {}
def fake_create(_url, **kwargs):
captured.update(kwargs)
return fake_engine
monkeypatch.setattr(engine_module, "create_async_engine", fake_create)
calls.attach_mock(AsyncMock(), "bootstrap_schema")
monkeypatch.setattr("deerflow.persistence.bootstrap.bootstrap_schema", calls.bootstrap_schema)
await engine_module.init_engine("postgres", url="postgresql+asyncpg://u:p@h:5432/db")
assert captured.get("connect_args", {}) == {
"command_timeout": engine_module.POSTGRES_COMMAND_TIMEOUT_SECONDS,
}
names = [c[0] for c in calls.mock_calls]
assert "execute" not in names # no CREATE SCHEMA
assert "bootstrap_schema" in names # bootstrap still runs
await engine_module.close_engine()

View File

@ -0,0 +1,112 @@
"""Optional live PostgreSQL schema integration tests for issue #3380."""
from __future__ import annotations
import os
import uuid
from types import SimpleNamespace
import pytest
from sqlalchemy import text
from deerflow.config.database_config import DatabaseConfig
from deerflow.persistence.engine import close_engine, get_engine, init_engine_from_config
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
from deerflow.runtime.checkpointer.provider import _resolve_checkpointer_config, _sync_checkpointer_cm
from deerflow.runtime.store.async_provider import make_store
from deerflow.runtime.store.provider import _resolve_store_config, _sync_store_cm
POSTGRES_URL = os.getenv("DEERFLOW_TEST_POSTGRES_URL")
pytestmark = pytest.mark.skipif(
not POSTGRES_URL,
reason="set DEERFLOW_TEST_POSTGRES_URL to run live PostgreSQL schema integration tests",
)
@pytest.mark.anyio
async def test_postgres_schema_places_orm_checkpointer_and_store_tables_together():
"""Verify a real PostgreSQL backend places all persistence tables in one schema."""
schema = f"deerflow_test_{uuid.uuid4().hex[:12]}"
db_config = DatabaseConfig(backend="postgres", postgres_url=POSTGRES_URL or "", postgres_schema=schema)
app_config = SimpleNamespace(checkpointer=None, database=db_config)
await init_engine_from_config(db_config)
engine = get_engine()
assert engine is not None
try:
async with make_checkpointer(app_config) as checkpointer:
assert checkpointer is not None
async with make_store(app_config) as store:
assert store is not None
async with engine.begin() as conn:
rows = (
await conn.execute(
text(
"""
SELECT table_schema, table_name
FROM information_schema.tables
WHERE table_schema IN (:schema, 'public')
ORDER BY table_schema, table_name
"""
),
{"schema": schema},
)
).all()
by_schema = {(row.table_schema, row.table_name) for row in rows}
orm_tables = {"runs", "run_events", "threads_meta", "feedback", "users"}
assert {("public", table) for table in orm_tables}.isdisjoint(by_schema)
assert {(schema, table) for table in orm_tables}.issubset(by_schema)
assert any(table_schema == schema and "checkpoint" in table_name for table_schema, table_name in by_schema)
assert any(table_schema == schema and ("store" in table_name or "migration" in table_name) for table_schema, table_name in by_schema)
finally:
async with engine.begin() as conn:
await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
await close_engine()
def test_sync_postgres_schema_places_checkpointer_and_store_tables_together():
"""Verify the sync (psycopg) path honours search_path via the DSN encoding.
This exercises ``dsn_with_search_path`` against a real psycopg connection,
guarding against regression of the ``%20`` vs ``+`` libpq encoding bug.
"""
import psycopg
schema = f"deerflow_test_{uuid.uuid4().hex[:12]}"
db_config = DatabaseConfig(
backend="postgres",
postgres_url=POSTGRES_URL or "",
postgres_schema=schema,
)
checkpointer_config = _resolve_checkpointer_config(SimpleNamespace(checkpointer=None, database=db_config))
store_config = _resolve_store_config(SimpleNamespace(checkpointer=None, database=db_config))
try:
with _sync_checkpointer_cm(checkpointer_config) as checkpointer:
assert checkpointer is not None
with _sync_store_cm(store_config) as store:
assert store is not None
with psycopg.connect(POSTGRES_URL or "", autocommit=True) as conn:
rows = conn.execute(
"""
SELECT table_schema, table_name
FROM information_schema.tables
WHERE table_schema IN (%s, 'public')
ORDER BY table_schema, table_name
""",
(schema,),
).fetchall()
by_schema = {(table_schema, table_name) for table_schema, table_name in rows}
assert any(table_schema == schema and "checkpoint" in table_name for table_schema, table_name in by_schema)
assert any(table_schema == schema and ("store" in table_name or "migration" in table_name) for table_schema, table_name in by_schema)
# The DeerFlow LangGraph tables must NOT leak into public.
assert not any(table_schema == "public" and ("checkpoint" in table_name or table_name == "store") for table_schema, table_name in by_schema)
finally:
with psycopg.connect(POSTGRES_URL or "", autocommit=True) as conn:
conn.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')

View File

@ -0,0 +1,178 @@
"""Tests for the PostgreSQL schema helpers (Issue #3380)."""
from urllib.parse import parse_qs, urlsplit
import pytest
from deerflow.persistence.postgres_schema import (
build_asyncpg_connect_args,
build_psycopg_options,
create_schema_sql,
dsn_with_search_path,
normalize_libpq_dsn,
)
class TestBuildAsyncpgConnectArgs:
def test_sets_search_path_for_schema(self):
assert build_asyncpg_connect_args("deerflow") == {"server_settings": {"search_path": "deerflow"}}
def test_empty_schema_returns_empty_dict(self):
assert build_asyncpg_connect_args("") == {}
class TestBuildPsycopgOptions:
def test_builds_libpq_options(self):
assert build_psycopg_options("deerflow") == "-c search_path=deerflow"
def test_empty_schema_returns_none(self):
assert build_psycopg_options("") is None
class TestCreateSchemaSql:
def test_builds_create_schema_statement(self):
assert create_schema_sql("deerflow") == 'CREATE SCHEMA IF NOT EXISTS "deerflow"'
def test_empty_schema_returns_none(self):
assert create_schema_sql("") is None
@pytest.mark.parametrize("schema", ['a"; DROP SCHEMA public; --', "MySchema", "a b", "deerflow\n"])
def test_rejects_non_plain_identifier(self, schema):
# Defense-in-depth: the SQL-emitting boundary re-validates so a caller
# that bypasses the pydantic config validator cannot inject.
with pytest.raises(ValueError):
create_schema_sql(schema)
class TestDsnWithSearchPath:
def test_empty_schema_returns_dsn_unchanged(self):
dsn = "postgresql://u:p@h:5432/db"
assert dsn_with_search_path(dsn, "") == dsn
def test_appends_options_query_encoded(self):
dsn = "postgresql://u:p@h:5432/db"
out = dsn_with_search_path(dsn, "deerflow")
# libpq only decodes %XX in URI query values; '+' is NOT treated as a
# space. The space MUST therefore be encoded as %20, never as '+'.
assert "+" not in out
assert "options=-c%20search_path%3Ddeerflow" in out
parts = urlsplit(out)
query = parse_qs(parts.query)
assert query["options"] == ["-c search_path=deerflow"]
def test_merges_with_existing_query(self):
dsn = "postgresql://u:p@h:5432/db?sslmode=require"
out = dsn_with_search_path(dsn, "deerflow")
query = parse_qs(urlsplit(out).query)
assert query["sslmode"] == ["require"]
assert query["options"] == ["-c search_path=deerflow"]
def test_replaces_existing_options_query(self):
dsn = "postgresql://u:p@h:5432/db?options=-c%20search_path%3Dpublic"
out = dsn_with_search_path(dsn, "deerflow")
query = parse_qs(urlsplit(out).query)
assert query["options"] == ["-c search_path=deerflow"]
def test_preserves_existing_options_query(self):
dsn = "postgresql://u:p@h:5432/db?options=-c%20statement_timeout%3D5000"
out = dsn_with_search_path(dsn, "deerflow")
query = parse_qs(urlsplit(out).query)
assert query["options"] == ["-c statement_timeout=5000 -c search_path=deerflow"]
def test_replaces_only_existing_search_path_option(self):
dsn = "postgresql://u:p@h:5432/db?options=-c%20statement_timeout%3D5000%20-c%20search_path%3Dpublic"
out = dsn_with_search_path(dsn, "deerflow")
query = parse_qs(urlsplit(out).query)
assert query["options"] == ["-c statement_timeout=5000 -c search_path=deerflow"]
def test_supports_keyword_dsn(self):
pytest.importorskip("psycopg")
from psycopg.conninfo import conninfo_to_dict
dsn = "host=localhost dbname=deerflow user=postgres"
out = dsn_with_search_path(dsn, "deerflow")
assert conninfo_to_dict(out) == {
"host": "localhost",
"dbname": "deerflow",
"user": "postgres",
"options": "-c search_path=deerflow",
}
def test_preserves_keyword_dsn_options(self):
pytest.importorskip("psycopg")
from psycopg.conninfo import conninfo_to_dict
dsn = "host=localhost dbname=deerflow options='-c statement_timeout=5000'"
out = dsn_with_search_path(dsn, "deerflow")
assert conninfo_to_dict(out)["options"] == "-c statement_timeout=5000 -c search_path=deerflow"
def test_normalizes_sqlalchemy_driver_scheme(self):
# DatabaseConfig.postgres_url may carry a +asyncpg suffix; the libpq DSN
# produced for psycopg must drop the driver and still inject search_path.
dsn = "postgresql+asyncpg://u:p@h:5432/db"
out = dsn_with_search_path(dsn, "deerflow")
parts = urlsplit(out)
assert parts.scheme == "postgresql"
query = parse_qs(parts.query)
assert query["options"] == ["-c search_path=deerflow"]
def test_rejects_non_postgres_url_scheme(self):
try:
dsn_with_search_path("mysql://localhost/db", "deerflow")
except ValueError as exc:
assert "Unsupported PostgreSQL DSN scheme" in str(exc)
else:
raise AssertionError("Expected ValueError")
def test_roundtrip_preserves_host_and_db(self):
dsn = "postgresql://u:p@h:5432/db"
out = dsn_with_search_path(dsn, "deerflow")
parts = urlsplit(out)
assert parts.hostname == "h"
assert parts.port == 5432
assert parts.path == "/db"
def test_preserves_option_value_containing_space(self):
# libpq's options parameter separates args on spaces unless they are
# backslash-escaped. shlex.join would emit single-quotes, which libpq
# treats as literal characters and would corrupt the option. A token
# carrying a space must round-trip as a single backslash-escaped token.
from deerflow.persistence.postgres_schema import _merge_search_path_option
merged = _merge_search_path_option(r"-c application_name=My\ App", "deerflow")
assert "'" not in merged
assert r"application_name=My\ App" in merged
assert merged.endswith("-c search_path=deerflow")
def test_preserves_option_value_containing_tab(self):
# Non-space whitespace (TAB/CR/LF) inside an existing escaped token must
# also be re-escaped on re-join, otherwise libpq re-tokenizes on the bare
# whitespace byte and the round-trip is lossy.
from deerflow.persistence.postgres_schema import (
_merge_search_path_option,
_split_libpq_options,
)
merged = _merge_search_path_option("-c application_name=My\\\tApp", "deerflow")
assert "'" not in merged
# The tab-bearing value must round-trip back to a single token.
tokens = _split_libpq_options(merged)
assert "application_name=My\tApp" in tokens
assert merged.endswith("-c search_path=deerflow")
class TestNormalizeLibpqDsn:
def test_strips_asyncpg_driver_suffix(self):
assert normalize_libpq_dsn("postgresql+asyncpg://u:p@h:5432/db") == "postgresql://u:p@h:5432/db"
def test_leaves_bare_postgres_scheme_unchanged(self):
dsn = "postgresql://u:p@h:5432/db"
assert normalize_libpq_dsn(dsn) == dsn
def test_leaves_keyword_dsn_unchanged(self):
dsn = "host=localhost dbname=deerflow"
assert normalize_libpq_dsn(dsn) == dsn
def test_rejects_non_postgres_scheme(self):
with pytest.raises(ValueError, match="Unsupported PostgreSQL DSN scheme"):
normalize_libpq_dsn("mysql://localhost/db")

View File

@ -1815,6 +1815,9 @@ skill_evolution:
# checkpointer:
# type: postgres
# connection_string: postgresql://user:password@localhost:5432/deerflow
# # Optional: place LangGraph checkpointer/store tables in this schema.
# # Leave empty to use the server default search_path, usually public.
# postgres_schema: deerflow
# ============================================================================
# Database
@ -1866,6 +1869,51 @@ skill_evolution:
# database:
# backend: postgres
# postgres_url: $DATABASE_URL
# # Optional: place DeerFlow ORM tables and LangGraph checkpointer/store
# # tables in this schema. Leave empty to use the server default
# # search_path, usually public.
# postgres_schema: deerflow
#
# postgres_schema only takes effect for new tables. On startup DeerFlow runs
# CREATE SCHEMA IF NOT EXISTS and pins the connection search_path; it never
# moves existing tables, so third-party tables in `public` are left untouched.
#
# Migrating an EXISTING deployment from `public` to a dedicated schema:
# 1. Stop the DeerFlow services.
# 2. Move EVERY DeerFlow-owned table AND the Alembic state into the new schema
# (run as the DB owner). Missing any table strands its rows: after restart
# DeerFlow recreates an empty counterpart in the target schema and the old
# rows stay invisible in `public`. Moving the ORM tables but leaving
# `alembic_version` behind is especially dangerous -- bootstrap then treats
# the target schema as unversioned, re-baselines it, and replays migrations.
# CREATE SCHEMA IF NOT EXISTS deerflow;
# -- Application ORM tables:
# ALTER TABLE public.runs SET SCHEMA deerflow;
# ALTER TABLE public.run_events SET SCHEMA deerflow;
# ALTER TABLE public.threads_meta SET SCHEMA deerflow;
# ALTER TABLE public.feedback SET SCHEMA deerflow;
# ALTER TABLE public.users SET SCHEMA deerflow;
# ALTER TABLE public.agents SET SCHEMA deerflow;
# -- IM channel tables:
# ALTER TABLE public.channel_connections SET SCHEMA deerflow;
# ALTER TABLE public.channel_credentials SET SCHEMA deerflow;
# ALTER TABLE public.channel_oauth_states SET SCHEMA deerflow;
# ALTER TABLE public.channel_conversations SET SCHEMA deerflow;
# -- Scheduled-task tables:
# ALTER TABLE public.scheduled_tasks SET SCHEMA deerflow;
# ALTER TABLE public.scheduled_task_runs SET SCHEMA deerflow;
# -- Alembic migration state (REQUIRED -- see note above):
# ALTER TABLE public.alembic_version SET SCHEMA deerflow;
# The DeerFlow-owned set grows over time; discover any tables the list
# above misses (as the DeerFlow DB owner/role) with:
# SELECT table_name FROM information_schema.tables WHERE table_schema='public';
# LangGraph checkpointer/store table names also vary by version -- list them
# the same way, then ALTER ... SET SCHEMA deerflow for each (e.g. checkpoints,
# checkpoint_blobs, checkpoint_writes, checkpoint_migrations, store,
# store_migrations).
# 3. Set postgres_schema: deerflow here and restart.
# 4. Verify with `SHOW search_path;` and a smoke test (create a run, read
# its history). The `public` schema should gain no new DeerFlow tables.
database:
backend: sqlite
sqlite_dir: .deer-flow/data

View File

@ -68,9 +68,15 @@ The frontend is a stateful chat application. Users create **threads** (conversat
1. Optional composer helpers such as `core/input-polish` can rewrite the local draft before submission, and `core/voice-input` can transcribe browser microphone input into that same local draft; confirmed user input then flows to thread hooks (`core/threads/hooks.ts`) → LangGraph SDK streaming
2. Stream events update thread state (messages, artifacts, todos, goal). The main thread stream uses the LangGraph SDK's `throttle: true` mode so updates received in the same macrotask coalesce before React is notified; do not replace it with a numeric delay without validating the SDK's trailing-debounce behavior on a continuous stream.
File-tool artifact auto-open work must run in an effect with timer cleanup; never schedule timers while rendering streamed `write_file` or `str_replace` updates.
`ThreadState.artifacts` remains the authoritative artifact list. The artifacts provider persists only thread-scoped panel UI state (`open`, selected path, and a refresh bootstrap cache) in session storage; an initial empty stream value must not overwrite that restored state before history finishes loading.
Formal artifact content is refreshed once when the run finishes; transient `write-file:` previews remain message-driven.
3. `useThreadHistory` loads persisted conversation pages from `GET /api/threads/{id}/messages/page`, preserving the backend's thread-global event `seq`; rendering overlays checkpoint/live copies at their matching canonical identities (a summarized checkpoint may contain a protected early input plus a recent tail). Context-compaction rescue diffs every retained visible identity rather than slicing at the first anchor, and keeps a run-scoped ledger of committed visible messages so replacement updates and repeated rolling checkpoint windows cannot erase an already displayed step. The resolver suppresses checkpoint/transient prefixes whose canonical position is still behind an unloaded cursor page instead of collapsing that unknown gap before a recent anchor, then adds optimistic messages without timestamp re-sorting. History invalidation preserves already-loaded pages so their established ordering positions are not discarded.
4. Stop actions call the LangGraph SDK stream stop path; `core/threads/hooks.ts` invalidates current-thread, thread-history, token-usage, and sidebar/search caches immediately and schedules one follow-up refetch because SDK stop may finish via abort + fire-and-forget cancel before backend title finalization commits
5. TanStack Query manages server state; localStorage stores user settings
5. TanStack Query manages server state; localStorage stores user settings. The
Settings > Tools MCP switch calls the targeted `PATCH /api/mcp/config`
mutation, disables switches until that mutation's success refetch completes,
displays the backend error `detail` through a toast, and invalidates
`["mcpConfig"]` only after success.
6. Components subscribe to thread state and render updates
Run duration is run-scoped UI metadata even though the compatibility field `additional_kwargs.turn_duration` is repeated on historical AI messages. `core/messages/run-duration.ts` folds those copies into one display anchored after the run's last visible message group. `MessageList` owns the temporary client-side duration for a just-completed live turn until authoritative history arrives. The duration is total run wall-clock time, not per-message reasoning time; reasoning disclosure and run activity/duration are rendered separately.
@ -88,6 +94,7 @@ Human input requests are a structured message protocol layered on normal chat hi
Tool-calling AI messages can contain user-visible text as well as `tool_calls`. `core/messages/utils.ts` keeps these turns in an `assistant:processing` group, and `components/workspace/messages/message-group.tsx` must render the visible text as a processing step instead of treating the message as only tool metadata. This preserves provider text such as error explanations or "trying another approach" notes during tool-heavy runs.
While the current turn is still loading, a content-only AI message after the latest visible human input also stays in that processing group until the turn settles: a provider may append tool-call chunks to the same message later, and classifying it as a final assistant bubble too early makes the text jump into the steps panel. `MessageGroup` therefore renders processing text even before the first tool call arrives.
The same rule applies after an earlier tool call: a later content-only AI message remains visible after the current last tool-call step while streaming, because that message may itself gain another tool call before the turn settles.
Because the same message is rendered by two different components over its lifetime, reasoning must sit above the answer text in both. `MessageListItem` paints the settled bubble's `<Reasoning>` disclosure above its content, so `MessageGroup` puts the trailing reasoning disclosure above the assistant text that follows it and `convertToSteps` emits a message's reasoning step before its content step — otherwise the two swap places the instant the turn settles (#4576). Assistant text emitted _before_ that reasoning keeps its earlier position; only the answer the reasoning produced moves below it.
Edit-and-rerun is deliberately latest-turn-only. `core/messages/utils.ts::getLatestEditableTurn()` exposes a human turn only when the transcript is idle and the most recent visible turn ends in a terminal assistant message. `core/threads/hooks.ts::editAndRegenerateMessage()` calls `POST /api/threads/{id}/runs/edit-regenerate/prepare`, submits the returned replacement message/checkpoint/metadata through the same LangGraph stream path as regenerate, optimistically hides the superseded message ids, and clears the optimistic replacement once the persisted replacement arrives.

View File

@ -1,19 +1,9 @@
"use client";
import { useRouter, useSearchParams } from "next/navigation";
import { useEffect, useState, useCallback, useRef } from "react";
import { useCallback, useEffect, useRef, useState } from "react";
/**
* Validates the next parameter only allows relative paths starting with /.
*/
function validateNextParam(next: string | null): string {
if (!next) return "/workspace";
if (!next.startsWith("/") || next.startsWith("//")) return "/workspace";
if (next.startsWith("http://") || next.startsWith("https://"))
return "/workspace";
if (next.includes(":")) return "/workspace";
return next;
}
import { resolveAuthNextPath } from "@/core/auth/next-path";
export default function AuthCallbackPage() {
const router = useRouter();
@ -27,7 +17,7 @@ export default function AuthCallbackPage() {
if (calledRef.current) return;
calledRef.current = true;
const next = validateNextParam(searchParams.get("next"));
const next = resolveAuthNextPath(searchParams.get("next"));
try {
const res = await fetch("/api/v1/auth/me", { credentials: "include" });

View File

@ -10,6 +10,7 @@ import { Button } from "@/components/ui/button";
import { FlickeringGrid } from "@/components/ui/flickering-grid";
import { Input } from "@/components/ui/input";
import { useAuth } from "@/core/auth/AuthProvider";
import { resolveAuthNextPath } from "@/core/auth/next-path";
import {
loadRememberLoginPreference,
saveRememberLoginPreference,
@ -22,39 +23,6 @@ import {
import { parseAuthError } from "@/core/auth/types";
import { useI18n } from "@/core/i18n/hooks";
/**
* Validate next parameter
* Prevent open redirect attacks
* Per RFC-001: Only allow relative paths starting with /
*/
function validateNextParam(next: string | null): string | null {
if (!next) {
return null;
}
// Need start with / (relative path)
if (!next.startsWith("/")) {
return null;
}
// Disallow protocol-relative URLs
if (
next.startsWith("//") ||
next.startsWith("http://") ||
next.startsWith("https://")
) {
return null;
}
// Disallow URLs with different protocols (e.g., javascript:, data:, etc)
if (next.includes(":") && !next.startsWith("/")) {
return null;
}
// Valid relative path
return next;
}
export default function LoginPage() {
const router = useRouter();
const searchParams = useSearchParams();
@ -94,7 +62,7 @@ export default function LoginPage() {
// Get next parameter for validated redirect
const nextParam = searchParams.get("next");
const redirectPath = validateNextParam(nextParam) ?? "/workspace";
const redirectPath = resolveAuthNextPath(nextParam);
const regularSignupAllowed = canCreateRegularAccount({
// A failed probe must not expose registration while the system's setup
// state is unknown. Existing users can still sign in normally.

View File

@ -1,7 +1,10 @@
import { usePathname } from "next/navigation";
import {
createContext,
useCallback,
useContext,
useEffect,
useRef,
useState,
type ReactNode,
} from "react";
@ -27,6 +30,46 @@ const ArtifactsContext = createContext<ArtifactsContextType | undefined>(
undefined,
);
const ARTIFACTS_STORAGE_PREFIX = "deerflow:artifacts:v1";
type PersistedArtifactsState = {
artifacts: string[];
selectedArtifact: string | null;
open: boolean;
};
function storageKey(pathname: string) {
return `${ARTIFACTS_STORAGE_PREFIX}:${encodeURIComponent(pathname)}`;
}
function readPersistedState(pathname: string): PersistedArtifactsState | null {
try {
const raw = window.sessionStorage.getItem(storageKey(pathname));
if (!raw) {
return null;
}
const parsed = JSON.parse(raw) as Partial<PersistedArtifactsState>;
if (
!Array.isArray(parsed.artifacts) ||
!parsed.artifacts.every((artifact) => typeof artifact === "string") ||
!(
parsed.selectedArtifact === null ||
typeof parsed.selectedArtifact === "string"
) ||
typeof parsed.open !== "boolean"
) {
return null;
}
return {
artifacts: parsed.artifacts,
selectedArtifact: parsed.selectedArtifact,
open: parsed.open,
};
} catch {
return null;
}
}
interface ArtifactsProviderProps {
children: ReactNode;
}
@ -40,6 +83,36 @@ export function ArtifactsProvider({ children }: ArtifactsProviderProps) {
);
const [autoOpen, setAutoOpen] = useState(true);
const { setOpen: setSidebarOpen } = useSidebar();
const pathname = usePathname();
const hydratedPathRef = useRef<string | null>(null);
useEffect(() => {
if (!pathname) {
return;
}
const persisted = readPersistedState(pathname);
setArtifacts(persisted?.artifacts ?? []);
setSelectedArtifact(persisted?.selectedArtifact ?? null);
setOpen(persisted?.open ?? env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true");
setAutoOpen(true);
setAutoSelect(!persisted?.selectedArtifact);
hydratedPathRef.current = pathname;
}, [pathname]);
useEffect(() => {
if (!pathname || hydratedPathRef.current !== pathname) {
return;
}
try {
window.sessionStorage.setItem(
storageKey(pathname),
JSON.stringify({ artifacts, selectedArtifact, open }),
);
} catch {
// Browser storage can be disabled or full; panel state must keep working.
}
}, [artifacts, open, pathname, selectedArtifact]);
const select = useCallback(
(artifact: string, autoSelect = false) => {

View File

@ -72,11 +72,12 @@ const ChatBox: React.FC<{
if (threadIdRef.current !== threadId) {
threadIdRef.current = threadId;
deselect();
setArtifacts([]);
}
// Update artifacts from the current thread
if (threadArtifacts) {
// An empty initial state must not erase artifacts restored by the provider
// before the persisted thread state has arrived.
if (threadArtifacts && threadArtifacts.length > 0) {
setArtifacts(threadArtifacts);
}

View File

@ -132,6 +132,25 @@ function MessageGroupComponent({
return filteredSteps[filteredSteps.length - 1];
}
}, [lastToolCallStep, steps]);
// Assistant text emitted after the trailing reasoning is the answer that
// reasoning produced, so it renders below the reasoning disclosure. The
// settled assistant bubble always paints reasoning above content, and the
// streaming processing group has to agree or the two swap places the moment
// the turn ends (#4576). Text emitted before that reasoning keeps its
// earlier position.
const belowLastReasoningAssistantTextSteps = useMemo(() => {
if (!lastReasoningStep) {
return [];
}
const index = steps.indexOf(lastReasoningStep);
return steps
.slice(index + 1)
.filter((step) => step.type === "assistantText");
}, [lastReasoningStep, steps]);
const belowLastReasoningSteps = useMemo(
() => new Set<CoTStep>(belowLastReasoningAssistantTextSteps),
[belowLastReasoningAssistantTextSteps],
);
const firstEligibleDebugSummaryStepIndexByMessageId = useMemo(() => {
const firstIndices = new Map<string, number>();
@ -324,7 +343,10 @@ function MessageGroupComponent({
</Button>
)}
{(lastToolCallStep ??
steps.some((step) => step.type === "assistantText")) && (
steps.some(
(step) =>
step.type === "assistantText" && !belowLastReasoningSteps.has(step),
)) && (
<ChainOfThoughtContent className="px-4 pb-2">
{(lastToolCallStep
? showAbove
@ -332,7 +354,11 @@ function MessageGroupComponent({
: aboveLastToolCallSteps.filter(
(step) => step.type === "assistantText",
)
: steps.filter((step) => step.type === "assistantText")
: steps.filter(
(step) =>
step.type === "assistantText" &&
!belowLastReasoningSteps.has(step),
)
).flatMap(renderStep)}
{lastToolCallStep && (
<>
@ -343,7 +369,9 @@ function MessageGroupComponent({
<FlipDisplay uniqueKey={lastToolCallStep.id ?? ""}>
{renderToolCall(lastToolCallStep, { isLast: true })}
</FlipDisplay>
{afterLastToolCallAssistantTextSteps.flatMap(renderStep)}
{afterLastToolCallAssistantTextSteps
.filter((step) => !belowLastReasoningSteps.has(step))
.flatMap(renderStep)}
</>
)}
</ChainOfThoughtContent>
@ -404,6 +432,11 @@ function MessageGroupComponent({
></ChainOfThoughtStep>
</ChainOfThoughtContent>
)}
{belowLastReasoningAssistantTextSteps.length > 0 && (
<ChainOfThoughtContent className="px-4 pb-2">
{belowLastReasoningAssistantTextSteps.flatMap(renderStep)}
</ChainOfThoughtContent>
)}
</>
)}
</ChainOfThought>
@ -947,15 +980,9 @@ function convertToSteps(messages: Message[]): CoTStep[] {
const { browserViews, toolCallResults } = indexToolCallData(messages);
for (const [messageIndex, message] of messages.entries()) {
if (message.type === "ai") {
const content = extractContentFromMessage(message);
if (content) {
steps.push({
id: `${message.id ?? `ai-${messageIndex}`}-content`,
messageId: message.id,
type: "assistantText",
content,
});
}
// Reasoning precedes the answer text it produced, so it is pushed first:
// step order is what the group renders in, and a message carrying both
// would otherwise paint its answer above its own thinking (#4576).
const reasoning = extractReasoningContentFromMessage(message);
if (reasoning) {
const step: CoTReasoningStep = {
@ -966,6 +993,15 @@ function convertToSteps(messages: Message[]): CoTStep[] {
};
steps.push(step);
}
const content = extractContentFromMessage(message);
if (content) {
steps.push({
id: `${message.id ?? `ai-${messageIndex}`}-content`,
messageId: message.id,
type: "assistantText",
content,
});
}
for (const tool_call of message.tool_calls ?? []) {
if (tool_call.name === "task") {
continue;

View File

@ -47,7 +47,7 @@ function MCPServerList({
servers?: Record<string, MCPServerConfig>;
}) {
const { t } = useI18n();
const { mutate: enableMCPServer } = useEnableMCPServer();
const { isPending, mutate: enableMCPServer } = useEnableMCPServer();
const entries = Object.entries(servers ?? {});
if (entries.length === 0) {
return (
@ -73,7 +73,9 @@ function MCPServerList({
<ItemActions>
<Switch
checked={config.enabled}
disabled={env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true"}
disabled={
env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true" || isPending
}
onCheckedChange={(checked) =>
enableMCPServer({ serverName: name, enabled: checked })
}

View File

@ -1,5 +1,5 @@
import { useQuery } from "@tanstack/react-query";
import { useMemo } from "react";
import { useEffect, useMemo, useRef } from "react";
import { useThread } from "@/components/workspace/messages/context";
@ -25,15 +25,27 @@ export function useArtifactContent({
return null;
}, [filepath, isWriteFile, thread]);
const { data, isLoading, error } = useQuery({
const { data, isLoading, error, refetch } = useQuery({
queryKey: ["artifact", filepath, threadId, isMock],
queryFn: () => {
return loadArtifactContent({ filepath, threadId, isMock });
},
enabled,
// Cache artifact content for 5 minutes to avoid repeated fetches (especially for .skill ZIP extraction)
staleTime: 5 * 60 * 1000,
staleTime: 0,
refetchOnWindowFocus: true,
});
// Refetch once when the run settles so edits made during the run are
// visible without a manual reload.
const wasLoadingRef = useRef(thread.isLoading);
useEffect(() => {
const wasLoading = wasLoadingRef.current;
wasLoadingRef.current = thread.isLoading;
if (wasLoading && !thread.isLoading && enabled && !isWriteFile) {
void refetch().catch(() => undefined);
}
}, [enabled, isWriteFile, refetch, thread.isLoading]);
return {
content: isWriteFile ? content : data?.content,
url: isWriteFile ? undefined : data?.url,

View File

@ -19,7 +19,7 @@ export async function loadArtifactContent({
enhancedFilepath = filepath + "/SKILL.md";
}
const url = urlOfArtifact({ filepath: enhancedFilepath, threadId, isMock });
const response = await fetch(url);
const response = await fetch(url, { cache: "no-store" });
const text = await response.text();
return { content: text, url };
}

View File

@ -0,0 +1,26 @@
export const DEFAULT_AUTH_NEXT_PATH = "/workspace";
export function validateAuthNextPath(
nextPath: string | null | undefined,
): string | null {
if (!nextPath) {
return null;
}
if (!nextPath.startsWith("/")) {
return null;
}
if (nextPath.startsWith("//")) {
return null;
}
if (nextPath.includes("\\") || nextPath.includes(":")) {
return null;
}
return nextPath;
}
export function resolveAuthNextPath(
nextPath: string | null | undefined,
fallback = DEFAULT_AUTH_NEXT_PATH,
): string {
return validateAuthNextPath(nextPath) ?? fallback;
}

View File

@ -52,3 +52,26 @@ export async function updateMCPConfig(config: MCPConfig) {
}
return response.json();
}
export async function updateMCPServerState(
serverName: string,
enabled: boolean,
) {
const response = await fetch(`${getBackendBaseURL()}/api/mcp/config`, {
method: "PATCH",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
server_name: serverName,
enabled,
}),
});
if (!response.ok) {
throw new MCPConfigRequestError(
response.status,
await readErrorDetail(response, "Failed to update MCP server"),
);
}
return response.json() as Promise<MCPConfig>;
}

View File

@ -1,6 +1,16 @@
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import {
type QueryClient,
useMutation,
useQuery,
useQueryClient,
} from "@tanstack/react-query";
import { toast } from "sonner";
import { loadMCPConfig, MCPConfigRequestError, updateMCPConfig } from "./api";
import {
loadMCPConfig,
MCPConfigRequestError,
updateMCPServerState,
} from "./api";
export function useMCPConfig() {
const { data, isLoading, error } = useQuery({
@ -12,35 +22,23 @@ export function useMCPConfig() {
return { config: data, isLoading, error };
}
interface EnableMCPServerVariables {
serverName: string;
enabled: boolean;
}
export function getEnableMCPServerMutationOptions(queryClient: QueryClient) {
return {
mutationFn: ({ serverName, enabled }: EnableMCPServerVariables) =>
updateMCPServerState(serverName, enabled),
onSuccess: () => queryClient.invalidateQueries({ queryKey: ["mcpConfig"] }),
onError: (error: Error) => {
toast.error(error.message);
},
};
}
export function useEnableMCPServer() {
const queryClient = useQueryClient();
const { config } = useMCPConfig();
return useMutation({
mutationFn: async ({
serverName,
enabled,
}: {
serverName: string;
enabled: boolean;
}) => {
if (!config) {
throw new Error("MCP config not found");
}
if (!config.mcp_servers[serverName]) {
throw new Error(`MCP server ${serverName} not found`);
}
await updateMCPConfig({
mcp_servers: {
...config.mcp_servers,
[serverName]: {
...config.mcp_servers[serverName],
enabled,
},
},
});
},
onSuccess: () => {
void queryClient.invalidateQueries({ queryKey: ["mcpConfig"] });
},
});
return useMutation(getEnableMCPServerMutationOptions(queryClient));
}

View File

@ -13,6 +13,7 @@ const MARKDOWN_THREAD_ID = "00000000-0000-0000-0000-000000003121";
const MARKDOWN_ANCHOR_THREAD_ID = "00000000-0000-0000-0000-000000003123";
const JSON_THREAD_ID = "00000000-0000-0000-0000-000000003122";
const PRESENTED_THREAD_ID = "00000000-0000-0000-0000-000000003123";
const PERSISTED_PANEL_THREAD_ID = "00000000-0000-0000-0000-000000003125";
const PDF_THREAD_ID = "00000000-0000-0000-0000-000000003124";
function writeFileMessages({
@ -326,6 +327,46 @@ test.describe("Artifact preview stability", () => {
await expect(artifactsPanel.getByText("Presented Report")).toBeVisible();
});
test("restores the artifact panel and selected file after a page refresh", async ({
page,
}) => {
mockLangGraphAPI(page, {
threads: [
{
thread_id: PERSISTED_PANEL_THREAD_ID,
title: "Persisted artifact panel",
messages: presentFilesMessages(),
artifacts: [MARKDOWN_ARTIFACT_PATH],
},
],
});
await page.route(
`**/api/threads/${PERSISTED_PANEL_THREAD_ID}/artifacts/mnt/user-data/outputs/presented-report.md`,
(route) =>
route.fulfill({
status: 200,
contentType: "text/markdown",
body: "# Presented Report\n\nGenerated content",
}),
);
await page.goto(`/workspace/chats/${PERSISTED_PANEL_THREAD_ID}`);
await expect(page.getByText("presented-report.md")).toBeVisible({
timeout: 15_000,
});
await page.getByText("presented-report.md").first().click();
const artifactsPanel = page.locator("#artifacts");
await expect(artifactsPanel.getByText("Presented Report")).toBeVisible();
await page.reload();
await expect(page.getByTestId("artifact-trigger")).toBeVisible();
await expect(
page.locator("#artifacts").getByText("Presented Report"),
).toBeVisible();
});
test("renders sandboxed iframe for a browser-previewable non-code file (urlOfArtifact path)", async ({
page,
}) => {

View File

@ -0,0 +1,180 @@
import { createServer } from "node:http";
import type { AddressInfo } from "node:net";
import { expect, test, type Locator } from "@playwright/test";
import { mockLangGraphAPI } from "./utils/mock-api";
const STREAMING_THREAD_ID = "00000000-0000-0000-0000-000000004576";
const SETTLED_THREAD_ID = "00000000-0000-0000-0000-000000004577";
const RUN_ID = "00000000-0000-0000-0000-000000004578";
const REASONING_TEXT =
"The user asked who I am, so I will list the core capabilities.";
const ANSWER_TEXT = "I am DeerFlow, an open-source super agent.";
const INITIAL_MESSAGES = [
{
type: "human",
id: "msg-human-4576",
content: [{ type: "text", text: "Who are you?" }],
},
];
const SETTLED_AI_MESSAGE = {
type: "ai",
id: "msg-ai-4576-settled",
content: ANSWER_TEXT,
additional_kwargs: { reasoning_content: REASONING_TEXT },
};
/**
* One AI chunk carrying both reasoning and answer text, which is the state the
* bubble is in while a reasoning model streams its answer.
*/
function reasoningStreamFrames() {
const events = [
{
event: "metadata",
data: { run_id: RUN_ID, thread_id: STREAMING_THREAD_ID },
},
{
event: "values",
data: {
messages: [
...INITIAL_MESSAGES,
{
type: "human",
id: "msg-human-4576-follow-up",
content: [{ type: "text", text: "Summarize that briefly" }],
},
],
},
},
{
event: "messages",
data: [
{
content: ANSWER_TEXT,
additional_kwargs: { reasoning_content: REASONING_TEXT },
response_metadata: {},
type: "AIMessageChunk",
name: null,
id: "msg-ai-4576-streaming",
tool_calls: [],
invalid_tool_calls: [],
usage_metadata: null,
tool_call_chunks: [],
chunk_position: null,
},
{},
],
},
];
return events.map(
(event) => `event: ${event.event}\ndata: ${JSON.stringify(event.data)}\n\n`,
);
}
/** Holds the SSE connection open so the turn stays in its streaming state. */
async function startHeldOpenStreamServer() {
const frames = reasoningStreamFrames();
const server = createServer((_request, response) => {
response.writeHead(200, {
"Access-Control-Allow-Origin": "*",
"Cache-Control": "no-cache",
"Content-Type": "text/event-stream",
});
response.write(frames.join(""));
});
await new Promise<void>((resolve, reject) => {
const handleError = (error: Error) => reject(error);
server.once("error", handleError);
server.listen(0, "127.0.0.1", () => {
server.off("error", handleError);
resolve();
});
});
const { port } = server.address() as AddressInfo;
return {
url: `http://127.0.0.1:${port}/runs/stream`,
async close() {
server.closeAllConnections();
await new Promise<void>((resolve, reject) => {
server.close((error) => (error ? reject(error) : resolve()));
});
},
};
}
async function expectRenderedAbove(upper: Locator, lower: Locator) {
await expect(upper).toBeVisible();
await expect(lower).toBeVisible();
const upperBox = await upper.boundingBox();
const lowerBox = await lower.boundingBox();
expect(upperBox).not.toBeNull();
expect(lowerBox).not.toBeNull();
expect(upperBox!.y).toBeLessThan(lowerBox!.y);
}
test("renders reasoning above the answer text while the turn is streaming", async ({
page,
}) => {
const streamServer = await startHeldOpenStreamServer();
mockLangGraphAPI(page, {
threads: [
{
thread_id: STREAMING_THREAD_ID,
title: "Streaming reasoning order",
messages: INITIAL_MESSAGES,
},
],
});
await page.route("**/api/langgraph/threads/*/runs/stream", (route) =>
route.continue({ url: streamServer.url }),
);
try {
await page.goto(`/workspace/chats/${STREAMING_THREAD_ID}`);
const textarea = page.getByPlaceholder(/how can i assist you/i);
await expect(textarea).toBeVisible({ timeout: 15_000 });
await textarea.fill("Summarize that briefly");
await textarea.press("Enter");
// The streaming turn renders inside the chain-of-thought panel, whose
// reasoning disclosure is labelled "Thinking".
await expectRenderedAbove(
page.getByText("Thinking", { exact: true }),
page.getByText(ANSWER_TEXT),
);
} finally {
await streamServer.close();
}
});
test("renders reasoning above the answer text after the turn settles", async ({
page,
}) => {
mockLangGraphAPI(page, {
threads: [
{
thread_id: SETTLED_THREAD_ID,
title: "Settled reasoning order",
messages: [...INITIAL_MESSAGES, SETTLED_AI_MESSAGE],
},
],
});
await page.goto(`/workspace/chats/${SETTLED_THREAD_ID}`);
// The settled turn renders as an assistant bubble, whose reasoning
// disclosure is labelled "Reasoning".
await expectRenderedAbove(
page.getByText("Reasoning", { exact: true }),
page.getByText(ANSWER_TEXT),
);
});

View File

@ -185,6 +185,108 @@ describe("MessageGroup", () => {
expect(timeoutSpy).not.toHaveBeenCalled();
});
it("renders streaming reasoning above the answer text of the same message", () => {
const html = renderGroup(
[
{
id: "ai-1",
type: "ai",
content: "Zephyr answer body.",
additional_kwargs: {
reasoning_content: "The user asked who I am, so I will summarize.",
},
} as Message,
],
{ isLoading: true },
);
expectRenderedInOrder(html, ["Thinking", ">Zephyr</span>"]);
});
it("renders streaming inline think reasoning above the answer text", () => {
const html = renderGroup(
[
{
id: "ai-1",
type: "ai",
content:
"<think>\nThe user only said hello, so I will greet back.\n</think>\n\nZephyr answer body.",
} as Message,
],
{ isLoading: true },
);
expectRenderedInOrder(html, ["Thinking", ">Zephyr</span>"]);
});
it("renders trailing reasoning above the answer text that follows a tool call", () => {
const html = renderGroup(
[
{
id: "ai-1",
type: "ai",
content: "",
tool_calls: [
{
id: "call-1",
name: "read_file",
args: { path: "message-group.tsx" },
},
],
} as Message,
{
id: "tool-1",
type: "tool",
name: "read_file",
tool_call_id: "call-1",
content: "file contents",
} as Message,
{
id: "ai-2",
type: "ai",
content: "Zephyr answer body.",
additional_kwargs: {
reasoning_content: "The file confirms the renderer order.",
},
} as Message,
],
{ isLoading: true },
);
expectRenderedInOrder(html, [
"message-group.tsx",
"Thinking",
">Zephyr</span>",
]);
});
it("keeps assistant text emitted before the trailing reasoning above it", () => {
const html = renderGroup(
[
{
id: "ai-1",
type: "ai",
content: "Quartz interim note.",
} as Message,
{
id: "ai-2",
type: "ai",
content: "Zephyr answer body.",
additional_kwargs: {
reasoning_content: "Now I can write the final answer.",
},
} as Message,
],
{ isLoading: true },
);
expectRenderedInOrder(html, [
">Quartz</span>",
"Thinking",
">Zephyr</span>",
]);
});
it("keeps tool-calling assistant text visible when reasoning is also present", () => {
const html = renderGroup([
{
@ -413,6 +515,15 @@ describe("MessageGroup", () => {
});
});
/** Asserts every needle is present and that they appear in the given order. */
function expectRenderedInOrder(html: string, needles: string[]) {
const indices = needles.map((needle) => html.indexOf(needle));
for (const index of indices) {
expect(index).toBeGreaterThan(-1);
}
expect(indices).toStrictEqual([...indices].sort((a, b) => a - b));
}
function renderGroup(
messages: Message[],
props: Omit<ComponentProps<typeof MessageGroup>, "messages"> = {},

View File

@ -0,0 +1,82 @@
import { afterEach, describe, expect, it, rs } from "@rstest/core";
import { cleanup, fireEvent, render, screen } from "@testing-library/react";
import { ToolSettingsPage } from "@/components/workspace/settings/tool-settings-page";
const mcpMockState = rs.hoisted(() => ({
isPending: false,
mutate: rs.fn(),
}));
rs.mock("@/core/i18n/hooks", () => ({
useI18n: () => ({
t: {
common: { loading: "Loading" },
settings: {
tools: {
title: "Tools",
description: "Manage MCP tools",
adminRequired: "Admin required",
empty: "No tools",
},
},
},
}),
}));
rs.mock("@/core/mcp/hooks", () => ({
useMCPConfig: () => ({
config: {
mcp_servers: {
github: { enabled: true, description: "GitHub tools" },
remote: { enabled: false, description: "Remote tools" },
},
},
isLoading: false,
error: null,
}),
useEnableMCPServer: () => ({
isPending: mcpMockState.isPending,
mutate: mcpMockState.mutate,
}),
}));
rs.mock("@/env", () => ({
env: { NEXT_PUBLIC_STATIC_WEBSITE_ONLY: "false" },
}));
afterEach(() => {
mcpMockState.isPending = false;
mcpMockState.mutate.mockReset();
cleanup();
});
describe("ToolSettingsPage MCP switches", () => {
it("disables every switch while a targeted update is pending", () => {
mcpMockState.isPending = true;
render(<ToolSettingsPage />);
const switches = screen.getAllByRole("switch");
expect(switches).toHaveLength(2);
for (const item of switches) {
expect((item as HTMLButtonElement).disabled).toBe(true);
}
});
it("submits only the selected server state when idle", () => {
render(<ToolSettingsPage />);
const switches = screen.getAllByRole("switch");
const githubSwitch = switches[0];
expect(githubSwitch).toBeDefined();
expect((githubSwitch as HTMLButtonElement).disabled).toBe(false);
fireEvent.click(githubSwitch!);
expect(mcpMockState.mutate).toHaveBeenCalledWith({
serverName: "github",
enabled: false,
});
});
});

View File

@ -0,0 +1,30 @@
import { describe, expect, test } from "@rstest/core";
import {
resolveAuthNextPath,
validateAuthNextPath,
} from "@/core/auth/next-path";
describe("auth next path validation", () => {
test("accepts local absolute paths", () => {
expect(validateAuthNextPath("/workspace")).toBe("/workspace");
expect(validateAuthNextPath("/workspace/chats/new?tab=recent#top")).toBe(
"/workspace/chats/new?tab=recent#top",
);
});
test("rejects external and ambiguous redirects", () => {
expect(validateAuthNextPath(null)).toBeNull();
expect(validateAuthNextPath("workspace")).toBeNull();
expect(validateAuthNextPath("//evil.example")).toBeNull();
expect(validateAuthNextPath("https://evil.example")).toBeNull();
expect(validateAuthNextPath("/:evil")).toBeNull();
expect(validateAuthNextPath("/\\evil.example")).toBeNull();
expect(validateAuthNextPath("/foo\\bar")).toBeNull();
});
test("falls back for unsafe paths", () => {
expect(resolveAuthNextPath("/\\evil.example")).toBe("/workspace");
expect(resolveAuthNextPath("/safe", "/fallback")).toBe("/safe");
});
});

View File

@ -28,6 +28,7 @@ import {
MCPConfigRequestError,
loadMCPConfig,
updateMCPConfig,
updateMCPServerState,
} from "@/core/mcp/api";
const mockedFetch = rs.mocked(fetcher);
@ -116,3 +117,41 @@ describe("updateMCPConfig", () => {
});
});
});
describe("updateMCPServerState", () => {
test("patches only the requested server state", async () => {
mockedFetch.mockResolvedValueOnce(
jsonResponse(200, {
mcp_servers: { github: { enabled: false } },
}),
);
await expect(updateMCPServerState("github", false)).resolves.toEqual({
mcp_servers: { github: { enabled: false } },
});
expect(mockedFetch).toHaveBeenCalledWith("/api/mcp/config", {
method: "PATCH",
headers: {
"Content-Type": "application/json",
},
body: JSON.stringify({
server_name: "github",
enabled: false,
}),
});
});
test("surfaces backend validation detail", async () => {
const detail =
"MCP server 'semantic-scholar' uses disallowed stdio command 's2-mcp-server'.";
mockedFetch.mockResolvedValueOnce(jsonResponse(400, { detail }));
await expect(
updateMCPServerState("semantic-scholar", true),
).rejects.toMatchObject({
name: "MCPConfigRequestError",
status: 400,
message: detail,
});
});
});

View File

@ -1,14 +1,23 @@
import { beforeEach, describe, expect, it, rs } from "@rstest/core";
import { QueryClient } from "@tanstack/react-query";
import { toast } from "sonner";
rs.mock("@/core/api/fetcher", () => ({
fetch: rs.fn(),
}));
rs.mock("sonner", () => ({
toast: {
error: rs.fn(),
},
}));
import { fetch } from "@/core/api/fetcher";
import { MCPConfigRequestError, loadMCPConfig } from "@/core/mcp/api";
import { getEnableMCPServerMutationOptions } from "@/core/mcp/hooks";
const mockedFetch = rs.mocked(fetch);
const mockedToastError = rs.mocked(toast.error);
function makeClient() {
return new QueryClient({
@ -23,6 +32,7 @@ function makeClient() {
describe("useMCPConfig retry policy", () => {
beforeEach(() => {
mockedFetch.mockReset();
mockedToastError.mockReset();
});
it("does not retry when loadMCPConfig throws MCPConfigRequestError (403)", async () => {
@ -82,3 +92,50 @@ describe("useMCPConfig retry policy", () => {
expect(mockedFetch).toHaveBeenCalledTimes(1);
});
});
describe("MCP server state mutation", () => {
beforeEach(() => {
mockedFetch.mockReset();
mockedToastError.mockReset();
});
it("invalidates MCP config after a successful targeted update", async () => {
mockedFetch.mockResolvedValue(
new Response(JSON.stringify({ mcp_servers: {} }), { status: 200 }),
);
const client = makeClient();
const invalidateQueries = rs
.spyOn(client, "invalidateQueries")
.mockResolvedValue();
const mutation = client
.getMutationCache()
.build(client, getEnableMCPServerMutationOptions(client));
await mutation.execute({ serverName: "github", enabled: false });
expect(invalidateQueries).toHaveBeenCalledWith({
queryKey: ["mcpConfig"],
});
expect(mockedToastError).not.toHaveBeenCalled();
});
it("shows the backend error detail when a targeted update fails", async () => {
const detail =
"MCP server 'semantic-scholar' uses disallowed stdio command 's2-mcp-server'.";
mockedFetch.mockResolvedValue(
new Response(JSON.stringify({ detail }), { status: 400 }),
);
const client = makeClient();
const invalidateQueries = rs.spyOn(client, "invalidateQueries");
const mutation = client
.getMutationCache()
.build(client, getEnableMCPServerMutationOptions(client));
await expect(
mutation.execute({ serverName: "semantic-scholar", enabled: true }),
).rejects.toThrow(detail);
expect(mockedToastError).toHaveBeenCalledWith(detail);
expect(invalidateQueries).not.toHaveBeenCalled();
});
});