diff --git a/CHANGELOG.md b/CHANGELOG.md index bd0db7cfa..5e79d4d75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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. diff --git a/README.md b/README.md index eb1827822..642eecea7 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 85fbee4c8..7b947ed96 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -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`) diff --git a/backend/app/gateway/routers/auth.py b/backend/app/gateway/routers/auth.py index 8b490c9cb..f9c1c3c01 100644 --- a/backend/app/gateway/routers/auth.py +++ b/backend/app/gateway/routers/auth.py @@ -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 diff --git a/backend/app/gateway/routers/mcp.py b/backend/app/gateway/routers/mcp.py index 105864c1c..0eaae52c1 100644 --- a/backend/app/gateway/routers/mcp.py +++ b/backend/app/gateway/routers/mcp.py @@ -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)}") diff --git a/backend/app/gateway/routers/skills.py b/backend/app/gateway/routers/skills.py index 437477ce2..68f7da7b5 100644 --- a/backend/app/gateway/routers/skills.py +++ b/backend/app/gateway/routers/skills.py @@ -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() diff --git a/backend/docs/API.md b/backend/docs/API.md index 22c0de562..f6948d27c 100644 --- a/backend/docs/API.md +++ b/backend/docs/API.md @@ -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 diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index ec08aeeb4..1b8a687bc 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -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.""" diff --git a/backend/packages/harness/deerflow/config/app_config.py b/backend/packages/harness/deerflow/config/app_config.py index 24064bd92..ae11a3f95 100644 --- a/backend/packages/harness/deerflow/config/app_config.py +++ b/backend/packages/harness/deerflow/config/app_config.py @@ -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 diff --git a/backend/packages/harness/deerflow/config/checkpointer_config.py b/backend/packages/harness/deerflow/config/checkpointer_config.py index ce10f3434..4d1ef59ff 100644 --- a/backend/packages/harness/deerflow/config/checkpointer_config.py +++ b/backend/packages/harness/deerflow/config/checkpointer_config.py @@ -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. diff --git a/backend/packages/harness/deerflow/config/database_config.py b/backend/packages/harness/deerflow/config/database_config.py index bc95052db..f903eba55 100644 --- a/backend/packages/harness/deerflow/config/database_config.py +++ b/backend/packages/harness/deerflow/config/database_config.py @@ -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) -- diff --git a/backend/packages/harness/deerflow/config/extensions_config.py b/backend/packages/harness/deerflow/config/extensions_config.py index 833ed42e5..05c237a31 100644 --- a/backend/packages/harness/deerflow/config/extensions_config.py +++ b/backend/packages/harness/deerflow/config/extensions_config.py @@ -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. diff --git a/backend/packages/harness/deerflow/config/postgres_schema.py b/backend/packages/harness/deerflow/config/postgres_schema.py new file mode 100644 index 000000000..2a0f0498f --- /dev/null +++ b/backend/packages/harness/deerflow/config/postgres_schema.py @@ -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 ""``, 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 diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index 148f82b49..0072616e4 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -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() diff --git a/backend/packages/harness/deerflow/persistence/bootstrap.py b/backend/packages/harness/deerflow/persistence/bootstrap.py index 92c9caad5..be5e5919a 100644 --- a/backend/packages/harness/deerflow/persistence/bootstrap.py +++ b/backend/packages/harness/deerflow/persistence/bootstrap.py @@ -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: diff --git a/backend/packages/harness/deerflow/persistence/engine.py b/backend/packages/harness/deerflow/persistence/engine.py index e6ac0e17a..1503e82e6 100644 --- a/backend/packages/harness/deerflow/persistence/engine.py +++ b/backend/packages/harness/deerflow/persistence/engine.py @@ -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 "", ) diff --git a/backend/packages/harness/deerflow/persistence/migrations/env.py b/backend/packages/harness/deerflow/persistence/migrations/env.py index 13b436b12..8e4f8c1d1 100644 --- a/backend/packages/harness/deerflow/persistence/migrations/env.py +++ b/backend/packages/harness/deerflow/persistence/migrations/env.py @@ -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 diff --git a/backend/packages/harness/deerflow/persistence/postgres_schema.py b/backend/packages/harness/deerflow/persistence/postgres_schema.py new file mode 100644 index 000000000..f7457b654 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/postgres_schema.py @@ -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=`` 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() diff --git a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py b/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py index be4543492..f48a4992c 100644 --- a/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py +++ b/backend/packages/harness/deerflow/runtime/checkpointer/async_provider.py @@ -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 diff --git a/backend/packages/harness/deerflow/runtime/checkpointer/provider.py b/backend/packages/harness/deerflow/runtime/checkpointer/provider.py index 1545ab03f..e2b260b76 100644 --- a/backend/packages/harness/deerflow/runtime/checkpointer/provider.py +++ b/backend/packages/harness/deerflow/runtime/checkpointer/provider.py @@ -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 diff --git a/backend/packages/harness/deerflow/runtime/store/async_provider.py b/backend/packages/harness/deerflow/runtime/store/async_provider.py index 035a1a9c4..4ef4967f1 100644 --- a/backend/packages/harness/deerflow/runtime/store/async_provider.py +++ b/backend/packages/harness/deerflow/runtime/store/async_provider.py @@ -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 diff --git a/backend/packages/harness/deerflow/runtime/store/provider.py b/backend/packages/harness/deerflow/runtime/store/provider.py index 37269995c..bb93d8400 100644 --- a/backend/packages/harness/deerflow/runtime/store/provider.py +++ b/backend/packages/harness/deerflow/runtime/store/provider.py @@ -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 diff --git a/backend/tests/blocking_io/test_mcp_router.py b/backend/tests/blocking_io/test_mcp_router.py index bbc779029..5ef82f52e 100644 --- a/backend/tests/blocking_io/test_mcp_router.py +++ b/backend/tests/blocking_io/test_mcp_router.py @@ -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']})" diff --git a/backend/tests/blocking_io/test_skills_update_router.py b/backend/tests/blocking_io/test_skills_update_router.py index 2655e57d1..aa01fe44c 100644 --- a/backend/tests/blocking_io/test_skills_update_router.py +++ b/backend/tests/blocking_io/test_skills_update_router.py @@ -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 diff --git a/backend/tests/test_app_config_reload.py b/backend/tests/test_app_config_reload.py index 5c1cc1e84..9baa71ff8 100644 --- a/backend/tests/test_app_config_reload.py +++ b/backend/tests/test_app_config_reload.py @@ -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" diff --git a/backend/tests/test_auth.py b/backend/tests/test_auth.py index b2b6aa1a4..bec069b16 100644 --- a/backend/tests/test_auth.py +++ b/backend/tests/test_auth.py @@ -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 diff --git a/backend/tests/test_checkpointer.py b/backend/tests/test_checkpointer.py index dd95426ed..ba2e4653e 100644 --- a/backend/tests/test_checkpointer.py +++ b/backend/tests/test_checkpointer.py @@ -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 diff --git a/backend/tests/test_extensions_config_atomic_write.py b/backend/tests/test_extensions_config_atomic_write.py new file mode 100644 index 000000000..fef45c165 --- /dev/null +++ b/backend/tests/test_extensions_config_atomic_write.py @@ -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": {}, + } diff --git a/backend/tests/test_mcp_config_secrets.py b/backend/tests/test_mcp_config_secrets.py index 7f2df98a0..af159e467 100644 --- a/backend/tests/test_mcp_config_secrets.py +++ b/backend/tests/test_mcp_config_secrets.py @@ -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 diff --git a/backend/tests/test_mcp_file_migration.py b/backend/tests/test_mcp_file_migration.py index 235be193e..c8216f804 100644 --- a/backend/tests/test_mcp_file_migration.py +++ b/backend/tests/test_mcp_file_migration.py @@ -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") diff --git a/backend/tests/test_persistence_bootstrap_url.py b/backend/tests/test_persistence_bootstrap_url.py index d29799c71..a2c6c2dcd 100644 --- a/backend/tests/test_persistence_bootstrap_url.py +++ b/backend/tests/test_persistence_bootstrap_url.py @@ -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" diff --git a/backend/tests/test_persistence_engine_postgres_config.py b/backend/tests/test_persistence_engine_postgres_config.py index 843919db0..e3c9cf39f 100644 --- a/backend/tests/test_persistence_engine_postgres_config.py +++ b/backend/tests/test_persistence_engine_postgres_config.py @@ -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() diff --git a/backend/tests/test_persistence_scaffold.py b/backend/tests/test_persistence_scaffold.py index ba460b9af..3e831e38c 100644 --- a/backend/tests/test_persistence_scaffold.py +++ b/backend/tests/test_persistence_scaffold.py @@ -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() diff --git a/backend/tests/test_pg_schema_integration.py b/backend/tests/test_pg_schema_integration.py new file mode 100644 index 000000000..7604a4e89 --- /dev/null +++ b/backend/tests/test_pg_schema_integration.py @@ -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') diff --git a/backend/tests/test_postgres_schema_helper.py b/backend/tests/test_postgres_schema_helper.py new file mode 100644 index 000000000..589ca097a --- /dev/null +++ b/backend/tests/test_postgres_schema_helper.py @@ -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") diff --git a/config.example.yaml b/config.example.yaml index a69ab1401..2578beba1 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -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 diff --git a/frontend/AGENTS.md b/frontend/AGENTS.md index c0bc87efc..fb605c627 100644 --- a/frontend/AGENTS.md +++ b/frontend/AGENTS.md @@ -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 `` 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. diff --git a/frontend/src/app/(auth)/auth/callback/page.tsx b/frontend/src/app/(auth)/auth/callback/page.tsx index a48ade103..79a7d6e0e 100644 --- a/frontend/src/app/(auth)/auth/callback/page.tsx +++ b/frontend/src/app/(auth)/auth/callback/page.tsx @@ -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" }); diff --git a/frontend/src/app/(auth)/login/page.tsx b/frontend/src/app/(auth)/login/page.tsx index b7acbc309..de392b113 100644 --- a/frontend/src/app/(auth)/login/page.tsx +++ b/frontend/src/app/(auth)/login/page.tsx @@ -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. diff --git a/frontend/src/components/workspace/artifacts/context.tsx b/frontend/src/components/workspace/artifacts/context.tsx index af9b19ab9..33a236930 100644 --- a/frontend/src/components/workspace/artifacts/context.tsx +++ b/frontend/src/components/workspace/artifacts/context.tsx @@ -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( 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; + 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(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) => { diff --git a/frontend/src/components/workspace/chats/chat-box.tsx b/frontend/src/components/workspace/chats/chat-box.tsx index 077603f5f..2963296eb 100644 --- a/frontend/src/components/workspace/chats/chat-box.tsx +++ b/frontend/src/components/workspace/chats/chat-box.tsx @@ -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); } diff --git a/frontend/src/components/workspace/messages/message-group.tsx b/frontend/src/components/workspace/messages/message-group.tsx index bc55194be..59cb969c7 100644 --- a/frontend/src/components/workspace/messages/message-group.tsx +++ b/frontend/src/components/workspace/messages/message-group.tsx @@ -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(belowLastReasoningAssistantTextSteps), + [belowLastReasoningAssistantTextSteps], + ); const firstEligibleDebugSummaryStepIndexByMessageId = useMemo(() => { const firstIndices = new Map(); @@ -324,7 +343,10 @@ function MessageGroupComponent({ )} {(lastToolCallStep ?? - steps.some((step) => step.type === "assistantText")) && ( + steps.some( + (step) => + step.type === "assistantText" && !belowLastReasoningSteps.has(step), + )) && ( {(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({ {renderToolCall(lastToolCallStep, { isLast: true })} - {afterLastToolCallAssistantTextSteps.flatMap(renderStep)} + {afterLastToolCallAssistantTextSteps + .filter((step) => !belowLastReasoningSteps.has(step)) + .flatMap(renderStep)} )} @@ -404,6 +432,11 @@ function MessageGroupComponent({ > )} + {belowLastReasoningAssistantTextSteps.length > 0 && ( + + {belowLastReasoningAssistantTextSteps.flatMap(renderStep)} + + )} )} @@ -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; diff --git a/frontend/src/components/workspace/settings/tool-settings-page.tsx b/frontend/src/components/workspace/settings/tool-settings-page.tsx index ce5746373..7572dfe30 100644 --- a/frontend/src/components/workspace/settings/tool-settings-page.tsx +++ b/frontend/src/components/workspace/settings/tool-settings-page.tsx @@ -47,7 +47,7 @@ function MCPServerList({ servers?: Record; }) { 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({ enableMCPServer({ serverName: name, enabled: checked }) } diff --git a/frontend/src/core/artifacts/hooks.ts b/frontend/src/core/artifacts/hooks.ts index df6a68a37..e66c3bee1 100644 --- a/frontend/src/core/artifacts/hooks.ts +++ b/frontend/src/core/artifacts/hooks.ts @@ -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, diff --git a/frontend/src/core/artifacts/loader.ts b/frontend/src/core/artifacts/loader.ts index 9584507a0..d935e1237 100644 --- a/frontend/src/core/artifacts/loader.ts +++ b/frontend/src/core/artifacts/loader.ts @@ -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 }; } diff --git a/frontend/src/core/auth/next-path.ts b/frontend/src/core/auth/next-path.ts new file mode 100644 index 000000000..41a96b263 --- /dev/null +++ b/frontend/src/core/auth/next-path.ts @@ -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; +} diff --git a/frontend/src/core/mcp/api.ts b/frontend/src/core/mcp/api.ts index fe057d7c4..fb9ea17ba 100644 --- a/frontend/src/core/mcp/api.ts +++ b/frontend/src/core/mcp/api.ts @@ -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; +} diff --git a/frontend/src/core/mcp/hooks.ts b/frontend/src/core/mcp/hooks.ts index 695392af2..12f3ef406 100644 --- a/frontend/src/core/mcp/hooks.ts +++ b/frontend/src/core/mcp/hooks.ts @@ -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)); } diff --git a/frontend/tests/e2e/artifact-preview.spec.ts b/frontend/tests/e2e/artifact-preview.spec.ts index 7d078c949..e7845ec25 100644 --- a/frontend/tests/e2e/artifact-preview.spec.ts +++ b/frontend/tests/e2e/artifact-preview.spec.ts @@ -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, }) => { diff --git a/frontend/tests/e2e/streaming-reasoning-order.spec.ts b/frontend/tests/e2e/streaming-reasoning-order.spec.ts new file mode 100644 index 000000000..5902a2ad6 --- /dev/null +++ b/frontend/tests/e2e/streaming-reasoning-order.spec.ts @@ -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((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((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), + ); +}); diff --git a/frontend/tests/unit/components/workspace/messages/message-group.test.ts b/frontend/tests/unit/components/workspace/messages/message-group.test.ts index 3746438b2..85ff47678 100644 --- a/frontend/tests/unit/components/workspace/messages/message-group.test.ts +++ b/frontend/tests/unit/components/workspace/messages/message-group.test.ts @@ -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"]); + }); + + it("renders streaming inline think reasoning above the answer text", () => { + const html = renderGroup( + [ + { + id: "ai-1", + type: "ai", + content: + "\nThe user only said hello, so I will greet back.\n\n\nZephyr answer body.", + } as Message, + ], + { isLoading: true }, + ); + + expectRenderedInOrder(html, ["Thinking", ">Zephyr"]); + }); + + 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", + ]); + }); + + 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", + "Thinking", + ">Zephyr", + ]); + }); + 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, "messages"> = {}, diff --git a/frontend/tests/unit/components/workspace/settings/tool-settings-page.dom.test.tsx b/frontend/tests/unit/components/workspace/settings/tool-settings-page.dom.test.tsx new file mode 100644 index 000000000..d25e99a06 --- /dev/null +++ b/frontend/tests/unit/components/workspace/settings/tool-settings-page.dom.test.tsx @@ -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(); + + 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(); + + 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, + }); + }); +}); diff --git a/frontend/tests/unit/core/auth/next-path.test.ts b/frontend/tests/unit/core/auth/next-path.test.ts new file mode 100644 index 000000000..1d9a06e29 --- /dev/null +++ b/frontend/tests/unit/core/auth/next-path.test.ts @@ -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"); + }); +}); diff --git a/frontend/tests/unit/core/mcp/api.test.ts b/frontend/tests/unit/core/mcp/api.test.ts index 813490d3c..9bcf949ea 100644 --- a/frontend/tests/unit/core/mcp/api.test.ts +++ b/frontend/tests/unit/core/mcp/api.test.ts @@ -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, + }); + }); +}); diff --git a/frontend/tests/unit/core/mcp/hooks.test.ts b/frontend/tests/unit/core/mcp/hooks.test.ts index 35e661103..5047800d5 100644 --- a/frontend/tests/unit/core/mcp/hooks.test.ts +++ b/frontend/tests/unit/core/mcp/hooks.test.ts @@ -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(); + }); +});