From a046d19ce13be77e77e6bff31a9c47c8b6581f10 Mon Sep 17 00:00:00 2001 From: Administrator Date: Sat, 27 Jun 2026 18:55:26 +0800 Subject: [PATCH] fix: improve model concurrency limiting --- frontend-web/src/core/threads/hooks.ts | 43 +- .../pages/ModelConcurrencySettingsPage.tsx | 4 +- .../backend/app/gateway/services.py | 1527 +++++++++-------- .../harness/deerflow/runtime/__init__.py | 103 +- .../harness/deerflow/runtime/runs/__init__.py | 33 +- .../harness/deerflow/runtime/runs/manager.py | 636 +++---- .../backend/tests/test_model_concurrency.py | 80 +- 7 files changed, 1266 insertions(+), 1160 deletions(-) diff --git a/frontend-web/src/core/threads/hooks.ts b/frontend-web/src/core/threads/hooks.ts index e3f725d8..03c37857 100644 --- a/frontend-web/src/core/threads/hooks.ts +++ b/frontend-web/src/core/threads/hooks.ts @@ -224,15 +224,24 @@ function isRunNotFoundError(message: string): boolean { * 「already has an active run」。匹配 HTTP 409、英文兜底文案，或后端已中文化 * 的提示，统一映射成一句友好的中文。 */ -function isActiveRunConflictError(message: string): boolean { - return ( - /\b409\b/.test(message) || - /already has an active run/i.test(message) || - /还在生成回答/.test(message) - ); -} - -function buildCurrentUserLabel(): string { +function isActiveRunConflictError(message: string): boolean { + return ( + /\b409\b/.test(message) || + /already has an active run/i.test(message) || + /还在生成回答/.test(message) + ); +} + +function isModelConcurrencyLimitError(message: string): boolean { + return ( + /\b429\b/.test(message) || + /并发已达到上限/.test(message) || + /使用人数较多/.test(message) || + /too many requests/i.test(message) + ); +} + +function buildCurrentUserLabel(): string { const auth = getStoredAuth(); if (!auth) { return "未登录"; @@ -570,12 +579,16 @@ export function useThreadStream({ } // 上一轮还没结束又发了新问题（后端 409）：给一句明确的中文提示，而不是 // 把原始的英文/HTTP 错误抛给用户。 - if (isActiveRunConflictError(message)) { - toast.info("当前对话还在生成回答，请等待本轮回复结束，或先点「停止」再发送新的问题。"); - return; - } - toast.error(message); - }, + if (isActiveRunConflictError(message)) { + toast.info("当前对话还在生成回答，请等待本轮回复结束，或先点「停止」再发送新的问题。"); + return; + } + if (isModelConcurrencyLimitError(message)) { + toast.info("当前使用人数较多，请稍等一下再试。"); + return; + } + toast.error(message); + }, onFinish(state) { listeners.current.onFinish?.(state.values); void queryClient.invalidateQueries({ queryKey: ["threads", "search"] }); diff --git a/frontend-web/src/pages/ModelConcurrencySettingsPage.tsx b/frontend-web/src/pages/ModelConcurrencySettingsPage.tsx index af060286..d8c59569 100644 --- a/frontend-web/src/pages/ModelConcurrencySettingsPage.tsx +++ b/frontend-web/src/pages/ModelConcurrencySettingsPage.tsx @@ -85,7 +85,7 @@ export default function ModelConcurrencySettingsPage({ className={ embedded ? "flex w-full flex-col gap-6" - : "mx-auto flex w-full max-w-3xl flex-col gap-6 px-6 py-10" + : "mx-auto flex h-full min-h-0 w-full max-w-3xl flex-col gap-6 overflow-y-auto px-6 py-10" } >
@@ -93,7 +93,7 @@ export default function ModelConcurrencySettingsPage({ 模型并发限制

- 按模型限制同时进行中的问答数。达到上限后，新请求会收到“当前并发已达到上限，请稍后再试”的提示。 + 按模型限制同时进行中的问答数。达到上限后，新请求会收到“当前使用人数较多，请稍等一下再试”的提示。

diff --git a/offline-backend-20260512/backend/app/gateway/services.py b/offline-backend-20260512/backend/app/gateway/services.py index 5a0d0252..cdda6c57 100644 --- a/offline-backend-20260512/backend/app/gateway/services.py +++ b/offline-backend-20260512/backend/app/gateway/services.py @@ -1,400 +1,397 @@ -"""Run lifecycle service layer. - -Centralizes the business logic for creating runs, formatting SSE -frames, and consuming stream bridge events. Router modules -(``thread_runs``, ``runs``) are thin HTTP handlers that delegate here. -""" - -from __future__ import annotations - -import asyncio -import json -import logging -import re -import time -from collections.abc import Mapping -from typing import Any - +"""Run lifecycle service layer. + +Centralizes the business logic for creating runs, formatting SSE +frames, and consuming stream bridge events. Router modules +(``thread_runs``, ``runs``) are thin HTTP handlers that delegate here. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import re +import time +from collections.abc import Mapping +from typing import Any + from fastapi import HTTPException, Request from langchain_core.messages import HumanMessage from app.gateway.deps import get_run_context, get_run_manager, get_stream_bridge from app.gateway.utils import sanitize_log_param from deerflow.config import get_app_config -from deerflow.config.agents_config import load_agent_config +from deerflow.config.agents_config import load_agent_config, validate_agent_id from deerflow.config.system_settings import load_system_settings from deerflow.runtime import ( - END_SENTINEL, - HEARTBEAT_SENTINEL, - ConflictError, - DisconnectMode, - RunManager, - RunRecord, - RunStatus, - StreamBridge, - UnsupportedStrategyError, - run_agent, -) - -logger = logging.getLogger(__name__) - - -# --------------------------------------------------------------------------- -# SSE formatting -# --------------------------------------------------------------------------- - - -async def immediate_end_sse(): - """Yield a single terminal ``end`` SSE frame and stop. - - Used when a run exists in the persistent store but not in the in-memory - registry — typically after a Gateway restart, or once a finished run's - record has been cleaned up. There is no live stream to join, so instead of - returning 404 (which the SDK's ``joinStream`` surfaces as a hard error and - never clears its reconnect key for) we stream a clean end. The SDK treats - this as success, clears ``lg:stream:{threadId}``, and the conversation - falls back to history rendering — killing the repeating switch-back 404. - """ - yield format_sse("end", None) - - -def format_sse(event: str, data: Any, *, event_id: str | None = None) -> str: - """Format a single SSE frame. - - Field order: ``event:`` -> ``data:`` -> ``id:`` (optional) -> blank line. - This matches the LangGraph Platform wire format consumed by the - ``useStream`` React hook and the Python ``langgraph-sdk`` SSE decoder. - """ - payload = json.dumps(data, default=str, ensure_ascii=False) - parts = [f"event: {event}", f"data: {payload}"] - if event_id: - parts.append(f"id: {event_id}") - parts.append("") - parts.append("") - return "\n".join(parts) - - -# --------------------------------------------------------------------------- -# Input / config helpers -# --------------------------------------------------------------------------- - - -def normalize_stream_modes(raw: list[str] | str | None) -> list[str]: - """Normalize the stream_mode parameter to a list. - - Default matches what ``useStream`` expects: values + messages-tuple. - """ - if raw is None: - return ["values"] - if isinstance(raw, str): - return [raw] - return raw if raw else ["values"] - - -def normalize_input(raw_input: dict[str, Any] | None) -> dict[str, Any]: - """Convert LangGraph Platform input format to LangChain state dict.""" - if raw_input is None: - return {} - messages = raw_input.get("messages") - if messages and isinstance(messages, list): - converted = [] - for msg in messages: - if isinstance(msg, dict): - content = msg.get("content", "") - # Preserve additional_kwargs (e.g. uploaded `files`, `client_ts`, - # `prompt_prefix`) and optional id/name so they round-trip through - # the checkpoint to the UI and to UploadsMiddleware. Dropping them - # here previously made uploaded files invisible in the chat and - # broke the per-message ask-time timestamp. - msg_kwargs: dict[str, Any] = {"content": content} - extra = msg.get("additional_kwargs") - if isinstance(extra, dict) and extra: - msg_kwargs["additional_kwargs"] = extra - if msg.get("id"): - msg_kwargs["id"] = msg["id"] - if msg.get("name"): - msg_kwargs["name"] = msg["name"] - # TODO: handle other message types (system, ai, tool); for now all - # inbound dict messages are treated as human turns. - converted.append(HumanMessage(**msg_kwargs)) - else: - converted.append(msg) - return {**raw_input, "messages": converted} - return raw_input - - -_DEFAULT_ASSISTANT_ID = "lead_agent" - -# Built-in LangGraph graph IDs registered in langgraph.json. -# These are NOT stored in the agents DB and must bypass the custom-agent access check. -_BUILTIN_GRAPH_IDS: frozenset[str] = frozenset({"lead_agent", "canvas_agent", "ai_writing"}) - -_RAG_CITATION_PROMPT = """当前对话已开启参考文献模式。 -如果本轮调用了搜索、知识库、检索类技能或工具，请基于本轮所有检索结果整体回答。 -如果你已经知道需要多个检索词，请优先在同一步里并行调用检索工具，等这些检索都返回后再回答，不要查一个检索词就先写一版答案。 -你必须把同一轮内多次检索得到的资料视为同一个 referenceBatch，按系统稍后提供的合并批次从 [1] 开始连续编号，不要让每个检索词各自从 [1] 重新编号。 -回答正文中的关键结论必须标注对应编号，例如 [1]、[2]，多个来源可写为 [1][3][5]。 -不要引用不存在的编号，不要使用技能内部自带的编号，不要把参考文献完整列表堆在正文末尾。 -如果没有任何可用检索结果，请直接说明未检索到可靠参考资料。""" - - -def _prepend_prompt_prefix(graph_input: dict[str, Any], prefix: str) -> dict[str, Any]: - """Return ``graph_input`` with the prefix *tagged* onto its first HumanMessage. - - The prefix is stored in ``additional_kwargs['prompt_prefix']`` rather than - concatenated into ``content``. This keeps the persisted/displayed message - clean — the user never sees the prefix — while ``PromptPrefixMiddleware`` - injects it into the content of the request actually sent to the LLM. - - Leaves the input unchanged when there are no messages or the first message - is not a human turn. - """ - if not prefix: - return graph_input - messages = graph_input.get("messages") - if not isinstance(messages, list) or not messages: - return graph_input - new_messages = list(messages) - for idx, msg in enumerate(new_messages): - if isinstance(msg, HumanMessage): - extra = dict(getattr(msg, "additional_kwargs", None) or {}) - extra["prompt_prefix"] = prefix - new_messages[idx] = msg.model_copy(update={"additional_kwargs": extra}) - break - return {**graph_input, "messages": new_messages} - - -def _prompt_prefix_opted_in(body: Any) -> bool: - """Detect whether the caller asked for the prompt prefix to be applied.""" - context = getattr(body, "context", None) - if isinstance(context, Mapping) and context.get("prompt_prefix_enabled"): - return True - metadata = getattr(body, "metadata", None) - if isinstance(metadata, Mapping) and metadata.get("prompt_prefix_enabled"): - return True - return False - - -def _rag_mode_requested(body: Any) -> bool: - context = getattr(body, "context", None) - if isinstance(context, Mapping) and context.get("rag_mode_enabled"): - return True - metadata = getattr(body, "metadata", None) - if isinstance(metadata, Mapping) and ( - metadata.get("rag_mode_enabled") or metadata.get("rag_mode") - ): - return True - return False - - -def _thread_rag_mode_enabled(record: Any) -> bool: - if not isinstance(record, Mapping): - return False - metadata = record.get("metadata") - return isinstance(metadata, Mapping) and bool( - metadata.get("rag_mode_enabled") or metadata.get("rag_mode") - ) - - -def _append_invisible_prompt(graph_input: dict[str, Any], prompt: str) -> dict[str, Any]: - messages = graph_input.get("messages") - if not isinstance(messages, list) or not messages: - return graph_input - new_messages = list(messages) - for idx, msg in enumerate(new_messages): - if not isinstance(msg, HumanMessage): - continue - extra = dict(getattr(msg, "additional_kwargs", None) or {}) - existing = str(extra.get("prompt_prefix") or "").strip() - extra["prompt_prefix"] = f"{existing}\n\n{prompt}".strip() if existing else prompt - new_messages[idx] = msg.model_copy(update={"additional_kwargs": extra}) - break - return {**graph_input, "messages": new_messages} - - -# Whitelist of run-context keys that the langgraph-compat layer forwards from -# ``body.context`` into the run config. ``config["context"]`` exists in -# LangGraph >=0.6, but these values must be written to both ``configurable`` -# (for legacy ``_get_runtime_config`` consumers) and ``context`` because -# LangGraph >=1.1.9 no longer makes ``ToolRuntime.context`` fall back to -# ``configurable`` for consumers like ``setup_agent``. -_CONTEXT_CONFIGURABLE_KEYS: frozenset[str] = frozenset( - { - "agent_id", - "model_name", - "mode", - "thinking_enabled", - # 强制关闭思考(内网 vLLM/Qwen 模型只声明 supports_thinking 时,单靠 - # thinking_enabled=False 不会注入 enable_thinking=false)。透传到 run config - # 后由 lead agent 交给 create_chat_model(force_disable_thinking=...)。 - "thinking_force_disabled", - "reasoning_effort", - "rag_mode_enabled", - "writing_mode", - "writing_artifact_path", - "canvas_artifact_title", - "canvas_artifact_type", - "is_plan_mode", - "subagent_enabled", - "max_concurrent_subagents", - "agent_name", - "is_bootstrap", - "is_writing_setup", - "valid_article_types", - "writing_setup_quick", - "excluded_tools", - "skill_stop_names", - "skill_stop_message", - # 记忆开关:必须在白名单里才会从 body.context 透传到 run config, - # 否则 agent 端 cfg.get(...) 永远取默认值。圆桌等临时任务靠这两个键 - # 关掉个人记忆召回 / builtin 记忆注入(见 multi_agent._run_payload)。 - "memory_recall_disabled", - "memory_injection_enabled", - # 上下文压缩开关:由聊天框下方的按钮控制(默认关闭),透传到 run config 后 - # 在 _build_middlewares 里决定是否挂载 SummarizationMiddleware。 - "summarization_enabled", - } -) - - -def merge_run_context_overrides(config: dict[str, Any], context: Mapping[str, Any] | None) -> None: - """Merge whitelisted keys from ``body.context`` into both ``config['configurable']`` - and ``config['context']`` so they are visible to legacy configurable readers and - to LangGraph ``ToolRuntime.context`` consumers (e.g. the ``setup_agent`` tool — - see issue #2677).""" - if not context: - return - configurable = config.setdefault("configurable", {}) - runtime_context = config.setdefault("context", {}) - for key in _CONTEXT_CONFIGURABLE_KEYS: - if key in context: - if isinstance(configurable, dict): - configurable.setdefault(key, context[key]) - if isinstance(runtime_context, dict): - runtime_context.setdefault(key, context[key]) - - -def resolve_agent_factory(assistant_id: str | None): - """Resolve the agent factory callable from config. - - Built-in graphs (``lead_agent``, ``canvas_agent``, ``ai_writing``) are dispatched - directly to their own factory. Custom agents share the ``lead_agent`` factory - with ``agent_id`` injected into configurable/context by :func:`build_run_config`. - """ - if assistant_id == "canvas_agent": - from deerflow.agents.canvas_agent.agent import make_canvas_agent - - return make_canvas_agent - - if assistant_id == "ai_writing": - from deerflow.agents.ai_writing.graph import make_ai_writing_graph - - return make_ai_writing_graph - - from deerflow.agents.lead_agent.agent import make_lead_agent - - return make_lead_agent - - -def build_run_config( - thread_id: str, - request_config: dict[str, Any] | None, - metadata: dict[str, Any] | None, - *, - assistant_id: str | None = None, -) -> dict[str, Any]: - """Build a RunnableConfig dict for the agent. - - When *assistant_id* refers to a custom agent (anything other than - ``"lead_agent"`` / ``None``), the id is forwarded as ``agent_id`` in - whichever runtime options container is active: ``context`` for - LangGraph >= 0.6.0 requests, otherwise ``configurable``. - ``make_lead_agent`` reads this key to load the matching - ``agents//SOUL.md`` and per-agent config — without it the agent - silently runs as the default lead agent. - - This mirrors the channel manager's ``_resolve_run_params`` logic so that - the LangGraph Platform-compatible HTTP API and the IM channel path behave - identically. - """ - config: dict[str, Any] = {"recursion_limit": 250} - if request_config: - # LangGraph >= 0.6.0 introduced ``context`` as the preferred way to - # pass thread-level data and rejects requests that include both - # ``configurable`` and ``context``. If the caller already sends - # ``context``, honour it and skip our own ``configurable`` dict. - if "context" in request_config: - if "configurable" in request_config: - logger.warning( - "build_run_config: client sent both 'context' and 'configurable'; preferring 'context' (LangGraph >= 0.6.0). thread_id=%s, caller_configurable keys=%s", - thread_id, - list(request_config.get("configurable", {}).keys()), - ) - context_value = request_config["context"] - if context_value is None: - context = {} - elif isinstance(context_value, Mapping): - context = dict(context_value) - else: - raise ValueError("request config 'context' must be a mapping or null.") - config["context"] = context - else: - configurable = {"thread_id": thread_id} - configurable.update(request_config.get("configurable", {})) - config["configurable"] = configurable - for k, v in request_config.items(): - if k not in ("configurable", "context"): - config[k] = v - else: - config["configurable"] = {"thread_id": thread_id} - - # Inject custom agent id when the caller specified a non-default assistant. - # Built-in graphs (lead_agent, canvas_agent) are dispatched via their own - # factory and must NOT receive an agent_id override. - if assistant_id and assistant_id not in _BUILTIN_GRAPH_IDS: - normalized = assistant_id.strip() - if not normalized or not re.fullmatch(r"[A-Za-z0-9_-]+", normalized): - raise ValueError(f"Invalid assistant_id {assistant_id!r}: must contain only letters, digits, hyphens, and underscores.") - if "configurable" in config: - target = config["configurable"] - elif "context" in config: - target = config["context"] - else: - target = config.setdefault("configurable", {}) - if target is not None and "agent_id" not in target: - target["agent_id"] = normalized - if metadata: - config.setdefault("metadata", {}).update(metadata) - return config - - + END_SENTINEL, + HEARTBEAT_SENTINEL, + ConflictError, + DisconnectMode, + ModelConcurrencyLimitError, + RunManager, + RunRecord, + RunStatus, + StreamBridge, + UnsupportedStrategyError, + run_agent, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# SSE formatting +# --------------------------------------------------------------------------- + + +async def immediate_end_sse(): + """Yield a single terminal ``end`` SSE frame and stop. + + Used when a run exists in the persistent store but not in the in-memory + registry — typically after a Gateway restart, or once a finished run's + record has been cleaned up. There is no live stream to join, so instead of + returning 404 (which the SDK's ``joinStream`` surfaces as a hard error and + never clears its reconnect key for) we stream a clean end. The SDK treats + this as success, clears ``lg:stream:{threadId}``, and the conversation + falls back to history rendering — killing the repeating switch-back 404. + """ + yield format_sse("end", None) + + +def format_sse(event: str, data: Any, *, event_id: str | None = None) -> str: + """Format a single SSE frame. + + Field order: ``event:`` -> ``data:`` -> ``id:`` (optional) -> blank line. + This matches the LangGraph Platform wire format consumed by the + ``useStream`` React hook and the Python ``langgraph-sdk`` SSE decoder. + """ + payload = json.dumps(data, default=str, ensure_ascii=False) + parts = [f"event: {event}", f"data: {payload}"] + if event_id: + parts.append(f"id: {event_id}") + parts.append("") + parts.append("") + return "\n".join(parts) + + +# --------------------------------------------------------------------------- +# Input / config helpers +# --------------------------------------------------------------------------- + + +def normalize_stream_modes(raw: list[str] | str | None) -> list[str]: + """Normalize the stream_mode parameter to a list. + + Default matches what ``useStream`` expects: values + messages-tuple. + """ + if raw is None: + return ["values"] + if isinstance(raw, str): + return [raw] + return raw if raw else ["values"] + + +def normalize_input(raw_input: dict[str, Any] | None) -> dict[str, Any]: + """Convert LangGraph Platform input format to LangChain state dict.""" + if raw_input is None: + return {} + messages = raw_input.get("messages") + if messages and isinstance(messages, list): + converted = [] + for msg in messages: + if isinstance(msg, dict): + content = msg.get("content", "") + # Preserve additional_kwargs (e.g. uploaded `files`, `client_ts`, + # `prompt_prefix`) and optional id/name so they round-trip through + # the checkpoint to the UI and to UploadsMiddleware. Dropping them + # here previously made uploaded files invisible in the chat and + # broke the per-message ask-time timestamp. + msg_kwargs: dict[str, Any] = {"content": content} + extra = msg.get("additional_kwargs") + if isinstance(extra, dict) and extra: + msg_kwargs["additional_kwargs"] = extra + if msg.get("id"): + msg_kwargs["id"] = msg["id"] + if msg.get("name"): + msg_kwargs["name"] = msg["name"] + # TODO: handle other message types (system, ai, tool); for now all + # inbound dict messages are treated as human turns. + converted.append(HumanMessage(**msg_kwargs)) + else: + converted.append(msg) + return {**raw_input, "messages": converted} + return raw_input + + +_DEFAULT_ASSISTANT_ID = "lead_agent" + +# Built-in LangGraph graph IDs registered in langgraph.json. +# These are NOT stored in the agents DB and must bypass the custom-agent access check. +_BUILTIN_GRAPH_IDS: frozenset[str] = frozenset({"lead_agent", "canvas_agent", "ai_writing"}) + +_RAG_CITATION_PROMPT = """当前对话已开启参考文献模式。 +如果本轮调用了搜索、知识库、检索类技能或工具，请基于本轮所有检索结果整体回答。 +如果你已经知道需要多个检索词，请优先在同一步里并行调用检索工具，等这些检索都返回后再回答，不要查一个检索词就先写一版答案。 +你必须把同一轮内多次检索得到的资料视为同一个 referenceBatch，按系统稍后提供的合并批次从 [1] 开始连续编号，不要让每个检索词各自从 [1] 重新编号。 +回答正文中的关键结论必须标注对应编号，例如 [1]、[2]，多个来源可写为 [1][3][5]。 +不要引用不存在的编号，不要使用技能内部自带的编号，不要把参考文献完整列表堆在正文末尾。 +如果没有任何可用检索结果，请直接说明未检索到可靠参考资料。""" + + +def _prepend_prompt_prefix(graph_input: dict[str, Any], prefix: str) -> dict[str, Any]: + """Return ``graph_input`` with the prefix *tagged* onto its first HumanMessage. + + The prefix is stored in ``additional_kwargs['prompt_prefix']`` rather than + concatenated into ``content``. This keeps the persisted/displayed message + clean — the user never sees the prefix — while ``PromptPrefixMiddleware`` + injects it into the content of the request actually sent to the LLM. + + Leaves the input unchanged when there are no messages or the first message + is not a human turn. + """ + if not prefix: + return graph_input + messages = graph_input.get("messages") + if not isinstance(messages, list) or not messages: + return graph_input + new_messages = list(messages) + for idx, msg in enumerate(new_messages): + if isinstance(msg, HumanMessage): + extra = dict(getattr(msg, "additional_kwargs", None) or {}) + extra["prompt_prefix"] = prefix + new_messages[idx] = msg.model_copy(update={"additional_kwargs": extra}) + break + return {**graph_input, "messages": new_messages} + + +def _prompt_prefix_opted_in(body: Any) -> bool: + """Detect whether the caller asked for the prompt prefix to be applied.""" + context = getattr(body, "context", None) + if isinstance(context, Mapping) and context.get("prompt_prefix_enabled"): + return True + metadata = getattr(body, "metadata", None) + if isinstance(metadata, Mapping) and metadata.get("prompt_prefix_enabled"): + return True + return False + + +def _rag_mode_requested(body: Any) -> bool: + context = getattr(body, "context", None) + if isinstance(context, Mapping) and context.get("rag_mode_enabled"): + return True + metadata = getattr(body, "metadata", None) + if isinstance(metadata, Mapping) and (metadata.get("rag_mode_enabled") or metadata.get("rag_mode")): + return True + return False + + +def _thread_rag_mode_enabled(record: Any) -> bool: + if not isinstance(record, Mapping): + return False + metadata = record.get("metadata") + return isinstance(metadata, Mapping) and bool(metadata.get("rag_mode_enabled") or metadata.get("rag_mode")) + + +def _append_invisible_prompt(graph_input: dict[str, Any], prompt: str) -> dict[str, Any]: + messages = graph_input.get("messages") + if not isinstance(messages, list) or not messages: + return graph_input + new_messages = list(messages) + for idx, msg in enumerate(new_messages): + if not isinstance(msg, HumanMessage): + continue + extra = dict(getattr(msg, "additional_kwargs", None) or {}) + existing = str(extra.get("prompt_prefix") or "").strip() + extra["prompt_prefix"] = f"{existing}\n\n{prompt}".strip() if existing else prompt + new_messages[idx] = msg.model_copy(update={"additional_kwargs": extra}) + break + return {**graph_input, "messages": new_messages} + + +# Whitelist of run-context keys that the langgraph-compat layer forwards from +# ``body.context`` into the run config. ``config["context"]`` exists in +# LangGraph >=0.6, but these values must be written to both ``configurable`` +# (for legacy ``_get_runtime_config`` consumers) and ``context`` because +# LangGraph >=1.1.9 no longer makes ``ToolRuntime.context`` fall back to +# ``configurable`` for consumers like ``setup_agent``. +_CONTEXT_CONFIGURABLE_KEYS: frozenset[str] = frozenset( + { + "agent_id", + "model_name", + "mode", + "thinking_enabled", + # 强制关闭思考(内网 vLLM/Qwen 模型只声明 supports_thinking 时,单靠 + # thinking_enabled=False 不会注入 enable_thinking=false)。透传到 run config + # 后由 lead agent 交给 create_chat_model(force_disable_thinking=...)。 + "thinking_force_disabled", + "reasoning_effort", + "rag_mode_enabled", + "writing_mode", + "writing_artifact_path", + "canvas_artifact_title", + "canvas_artifact_type", + "is_plan_mode", + "subagent_enabled", + "max_concurrent_subagents", + "agent_name", + "is_bootstrap", + "is_writing_setup", + "valid_article_types", + "writing_setup_quick", + "excluded_tools", + "skill_stop_names", + "skill_stop_message", + # 记忆开关:必须在白名单里才会从 body.context 透传到 run config, + # 否则 agent 端 cfg.get(...) 永远取默认值。圆桌等临时任务靠这两个键 + # 关掉个人记忆召回 / builtin 记忆注入(见 multi_agent._run_payload)。 + "memory_recall_disabled", + "memory_injection_enabled", + # 上下文压缩开关:由聊天框下方的按钮控制(默认关闭),透传到 run config 后 + # 在 _build_middlewares 里决定是否挂载 SummarizationMiddleware。 + "summarization_enabled", + } +) + + +def merge_run_context_overrides(config: dict[str, Any], context: Mapping[str, Any] | None) -> None: + """Merge whitelisted keys from ``body.context`` into both ``config['configurable']`` + and ``config['context']`` so they are visible to legacy configurable readers and + to LangGraph ``ToolRuntime.context`` consumers (e.g. the ``setup_agent`` tool — + see issue #2677).""" + if not context: + return + configurable = config.setdefault("configurable", {}) + runtime_context = config.setdefault("context", {}) + for key in _CONTEXT_CONFIGURABLE_KEYS: + if key in context: + if isinstance(configurable, dict): + configurable.setdefault(key, context[key]) + if isinstance(runtime_context, dict): + runtime_context.setdefault(key, context[key]) + + +def resolve_agent_factory(assistant_id: str | None): + """Resolve the agent factory callable from config. + + Built-in graphs (``lead_agent``, ``canvas_agent``, ``ai_writing``) are dispatched + directly to their own factory. Custom agents share the ``lead_agent`` factory + with ``agent_id`` injected into configurable/context by :func:`build_run_config`. + """ + if assistant_id == "canvas_agent": + from deerflow.agents.canvas_agent.agent import make_canvas_agent + + return make_canvas_agent + + if assistant_id == "ai_writing": + from deerflow.agents.ai_writing.graph import make_ai_writing_graph + + return make_ai_writing_graph + + from deerflow.agents.lead_agent.agent import make_lead_agent + + return make_lead_agent + + +def build_run_config( + thread_id: str, + request_config: dict[str, Any] | None, + metadata: dict[str, Any] | None, + *, + assistant_id: str | None = None, +) -> dict[str, Any]: + """Build a RunnableConfig dict for the agent. + + When *assistant_id* refers to a custom agent (anything other than + ``"lead_agent"`` / ``None``), the id is forwarded as ``agent_id`` in + whichever runtime options container is active: ``context`` for + LangGraph >= 0.6.0 requests, otherwise ``configurable``. + ``make_lead_agent`` reads this key to load the matching + ``agents//SOUL.md`` and per-agent config — without it the agent + silently runs as the default lead agent. + + This mirrors the channel manager's ``_resolve_run_params`` logic so that + the LangGraph Platform-compatible HTTP API and the IM channel path behave + identically. + """ + config: dict[str, Any] = {"recursion_limit": 250} + if request_config: + # LangGraph >= 0.6.0 introduced ``context`` as the preferred way to + # pass thread-level data and rejects requests that include both + # ``configurable`` and ``context``. If the caller already sends + # ``context``, honour it and skip our own ``configurable`` dict. + if "context" in request_config: + if "configurable" in request_config: + logger.warning( + "build_run_config: client sent both 'context' and 'configurable'; preferring 'context' (LangGraph >= 0.6.0). thread_id=%s, caller_configurable keys=%s", + thread_id, + list(request_config.get("configurable", {}).keys()), + ) + context_value = request_config["context"] + if context_value is None: + context = {} + elif isinstance(context_value, Mapping): + context = dict(context_value) + else: + raise ValueError("request config 'context' must be a mapping or null.") + config["context"] = context + else: + configurable = {"thread_id": thread_id} + configurable.update(request_config.get("configurable", {})) + config["configurable"] = configurable + for k, v in request_config.items(): + if k not in ("configurable", "context"): + config[k] = v + else: + config["configurable"] = {"thread_id": thread_id} + + # Inject custom agent id when the caller specified a non-default assistant. + # Built-in graphs (lead_agent, canvas_agent) are dispatched via their own + # factory and must NOT receive an agent_id override. + if assistant_id and assistant_id not in _BUILTIN_GRAPH_IDS: + normalized = assistant_id.strip() + if not normalized or not re.fullmatch(r"[A-Za-z0-9_-]+", normalized): + raise ValueError(f"Invalid assistant_id {assistant_id!r}: must contain only letters, digits, hyphens, and underscores.") + if "configurable" in config: + target = config["configurable"] + elif "context" in config: + target = config["context"] + else: + target = config.setdefault("configurable", {}) + if target is not None and "agent_id" not in target: + target["agent_id"] = normalized + if metadata: + config.setdefault("metadata", {}).update(metadata) + return config + + def _collect_requested_agent_ids(body: Any) -> set[str]: - """Collect agent ids from assistant_id/config/context for access checks.""" - agent_ids: set[str] = set() - if body.assistant_id and body.assistant_id not in _BUILTIN_GRAPH_IDS: - agent_ids.add(str(body.assistant_id)) - - def collect(container: Mapping[str, Any] | None) -> None: - if not container: - return - for key in ("agent_id", "agent_name"): - value = container.get(key) - if value: - agent_ids.add(str(value)) - - context = getattr(body, "context", None) - if isinstance(context, Mapping): - collect(context) - - config = body.config or {} - if isinstance(config, Mapping): - nested_context = config.get("context") - nested_configurable = config.get("configurable") - if isinstance(nested_context, Mapping): - collect(nested_context) - if isinstance(nested_configurable, Mapping): - collect(nested_configurable) - + """Collect agent ids from assistant_id/config/context for access checks.""" + agent_ids: set[str] = set() + if body.assistant_id and body.assistant_id not in _BUILTIN_GRAPH_IDS: + agent_ids.add(str(body.assistant_id)) + + def collect(container: Mapping[str, Any] | None) -> None: + if not container: + return + for key in ("agent_id", "agent_name"): + value = container.get(key) + if value: + agent_ids.add(str(value)) + + context = getattr(body, "context", None) + if isinstance(context, Mapping): + collect(context) + + config = body.config or {} + if isinstance(config, Mapping): + nested_context = config.get("context") + nested_configurable = config.get("configurable") + if isinstance(nested_context, Mapping): + collect(nested_context) + if isinstance(nested_configurable, Mapping): + collect(nested_configurable) + return agent_ids @@ -420,6 +417,46 @@ def _extract_requested_model_name(body: Any) -> str | None: return None +def _extract_effective_agent_id(body: Any) -> str | None: + """Best-effort read of the custom agent id that the lead graph will run.""" + candidates: list[str] = [] + + def collect(container: Mapping[str, Any] | None) -> None: + if not container: + return + for key in ("agent_id", "agent_name"): + value = str(container.get(key) or "").strip() + if value: + candidates.append(value) + + context = getattr(body, "context", None) + if isinstance(context, Mapping): + collect(context) + + config = getattr(body, "config", None) + if isinstance(config, Mapping): + nested_context = config.get("context") + nested_configurable = config.get("configurable") + if isinstance(nested_context, Mapping): + collect(nested_context) + if isinstance(nested_configurable, Mapping): + collect(nested_configurable) + + assistant_id = str(getattr(body, "assistant_id", "") or "").strip() + if assistant_id and assistant_id not in _BUILTIN_GRAPH_IDS: + candidates.append(assistant_id) + + for agent_id in candidates: + try: + normalized = validate_agent_id(agent_id) + except ValueError: + logger.debug("Ignoring invalid agent id during model resolution: %s", agent_id) + continue + if normalized: + return normalized + return None + + def _resolve_effective_run_model_name(body: Any) -> str: """Resolve the model that this run will effectively use for concurrency limiting.""" app_config = get_app_config() @@ -438,14 +475,14 @@ def _resolve_effective_run_model_name(body: Any) -> str: ) return default_model_name - assistant_id = str(getattr(body, "assistant_id", "") or "").strip() - if assistant_id and assistant_id not in _BUILTIN_GRAPH_IDS: + agent_id = _extract_effective_agent_id(body) + if agent_id: try: - agent_config = load_agent_config(assistant_id) + agent_config = load_agent_config(agent_id) except FileNotFoundError: agent_config = None except Exception: - logger.debug("Failed to load agent config for model resolution: %s", assistant_id, exc_info=True) + logger.debug("Failed to load agent config for model resolution: %s", agent_id, exc_info=True) agent_config = None agent_model_name = str(getattr(agent_config, "model", "") or "").strip() if agent_model_name: @@ -454,216 +491,204 @@ def _resolve_effective_run_model_name(body: Any) -> str: logger.warning( "Agent model '%s' for agent '%s' not found; falling back to default model '%s'.", agent_model_name, - assistant_id, + agent_id, default_model_name, ) return default_model_name -async def _enforce_model_concurrency_limit(run_mgr: RunManager, model_name: str) -> None: - """Reject the run when the selected model has reached its concurrent-run limit.""" +def _get_model_concurrency_limit(model_name: str) -> int | None: + """Return the configured limit for a model, or ``None`` when limiting is off.""" settings = load_system_settings().model_concurrency if not settings.enabled: - return + return None limit = settings.model_limits.get(model_name, settings.default_limit) try: limit = int(limit) except (TypeError, ValueError): limit = int(settings.default_limit) - limit = max(1, limit) - - active = await run_mgr.list_active() - current = sum( - 1 - for record in active - if str((record.metadata or {}).get("model_name") or "").strip() == model_name - ) - if current >= limit: - raise HTTPException( - status_code=429, - detail=f"当前模型 {model_name} 并发已达到上限（{limit}），请稍后再试。", - ) + return max(1, limit) async def _sync_legacy_agent_records(store) -> None: - """Import existing folder-backed agents as built-in records. - - Heavy: ``list_custom_agents()`` scans ``.deer-flow/agents/`` and reads - ``config.yaml`` + ``SOUL.md`` per directory, then ``store.ensure_builtin`` - does one DB round-trip per agent. On a machine with N agent dirs this is - O(N) file IO + O(N) DB calls — historically the worst per-request hotspot - in ``start_run`` (observed 30s+ on machines with 100+ residual coordinator - shells + slow MySQL). **Never call this directly from request handlers** — - use :func:`_maybe_sync_legacy_agent_records` which applies the same TTL - gating as ``routers/agents._maybe_sync_legacy_agents``. - """ - from deerflow.config.agents_config import list_custom_agents - - for agent_cfg in list_custom_agents(): - agent_id = agent_cfg.id or agent_cfg.name - if not agent_id: - continue - await store.ensure_builtin( - { - "id": agent_id, - "name": agent_cfg.name, - "description": agent_cfg.description or "", - "published": True, - } - ) - - -# Throttle per-request legacy-sync. Matches ``routers/agents._LEGACY_SYNC_TTL_SECONDS`` -# semantics — the real sync runs once at startup (lifespan); this is a safety -# net for environments that bypass the lifespan hook or for freshly dropped -# agent directories. Lifespan registers itself via the same state by calling -# the agents router's ``_mark_legacy_sync_fresh()`` after its boot-time sync, -# so this module's TTL state is independent (a separate gate kept here so we -# don't import from ``routers/agents`` and create a circular import). -_LEGACY_SYNC_TTL_SECONDS = 3600.0 -_legacy_sync_state: dict[str, Any] = {"last_at": 0.0, "lock": None} - - -# Anything slower than this triggers the WARNING-level breakdown in -# ``start_run``. 500ms gives plenty of headroom for cold caches + a few DB -# round-trips on healthy systems while still catching the pathological -# multi-second cases that brought us here. -_START_RUN_SLOW_THRESHOLD_SECONDS = 0.5 - - -def _mark_legacy_sync_fresh() -> None: - """Tell this module's TTL gate that a sync just completed elsewhere. - - Called by the Gateway lifespan hook after it runs the boot-time legacy - agent sync (``routers/agents._sync_legacy_agents``). Without this, the - first request after startup would race past the TTL gate here and trigger - a redundant rescan even though the data is already fresh. - """ - _legacy_sync_state["last_at"] = time.monotonic() - - -async def _maybe_sync_legacy_agent_records(store) -> None: - """TTL-gated wrapper around :func:`_sync_legacy_agent_records`. - - Steady-state behavior: returns in microseconds inside the TTL window — no - filesystem walk, no DB round-trips. Outside the TTL window we serialize - callers under an :class:`asyncio.Lock` so a burst of concurrent - ``/runs/stream`` requests fans into a single rescan instead of N parallel - fan-outs. Pattern mirrors ``routers/agents._maybe_sync_legacy_agents``. - """ - now = time.monotonic() - if now - float(_legacy_sync_state["last_at"]) < _LEGACY_SYNC_TTL_SECONDS: - return - lock = _legacy_sync_state["lock"] - if lock is None: - lock = asyncio.Lock() - _legacy_sync_state["lock"] = lock - async with lock: - # Re-check under the lock — a sibling coroutine may have just refreshed. - if time.monotonic() - float(_legacy_sync_state["last_at"]) < _LEGACY_SYNC_TTL_SECONDS: - return - t_sync_start = time.perf_counter() - await _sync_legacy_agent_records(store) - elapsed = time.perf_counter() - t_sync_start - _legacy_sync_state["last_at"] = time.monotonic() - # When this fires it's a "first request of the hour" warmup — log it - # so operators can see the spike isn't an isolated bug. Anything - # over 2s usually means a lot of agent directories + slow IO/DB. - logger.info( - "[start-run-timing] legacy-agent sync warmed cache in %.3fs (next sync in %.0fs)", - elapsed, - _LEGACY_SYNC_TTL_SECONDS, - ) - - -async def _enforce_agent_access(body: Any, request: Request) -> None: - agent_ids = _collect_requested_agent_ids(body) - if not agent_ids: - return - - from app.gateway.deps import get_agent_store - from deerflow.config.agents_config import validate_agent_id - from deerflow.runtime.user_context import get_effective_user_id - - store = get_agent_store(request) - await _maybe_sync_legacy_agent_records(store) - user_id = get_effective_user_id() - # Admins see every agent in the admin console listing (?admin=true), so - # they may also start a run against any of them. - # ``request`` may be an in-process virtual request (SimpleNamespace) from - # the roundtable job gateway / scheduled-task runtime — those carry no - # ``state`` attribute, so the whole chain must be getattr-guarded - # (a bare ``request.state`` raises AttributeError and kills the run). - is_admin = getattr(getattr(getattr(request, "state", None), "user", None), "system_role", None) == "admin" - for agent_id in agent_ids: - try: - normalized = validate_agent_id(agent_id) - except ValueError as exc: - raise HTTPException(status_code=422, detail=str(exc)) from exc - if normalized is None: - continue - if await store.get_visible(normalized, user_id) is not None: - continue - if is_admin and await store.get_any(normalized) is not None: - continue - raise HTTPException(status_code=404, detail=f"Agent {normalized} not found") - - -# --------------------------------------------------------------------------- -# Run lifecycle -# --------------------------------------------------------------------------- - - -async def start_run( - body: Any, - thread_id: str, - request: Request, - *, - enforce_access: bool = True, -) -> RunRecord: - """Create a RunRecord and launch the background agent task. - - Parameters - ---------- - body : RunCreateRequest - The validated request body (typed as Any to avoid circular import - with the router module that defines the Pydantic model). - thread_id : str - Target thread. - request : Request - FastAPI request — used to retrieve singletons from ``app.state``. - - Timing observability - -------------------- - Each phase records ``perf_counter`` deltas. When the total wallclock from - function entry to ``run_agent`` task creation exceeds - ``_START_RUN_SLOW_THRESHOLD_SECONDS`` we emit a ``[start-run-timing]`` - breakdown at WARNING level. Stays at DEBUG when fast, so steady-state - logs are not polluted. Investigate spikes by grep'ing - ``[start-run-timing]`` together with ``[upstream-timing]`` (multi-agent - router) and ``[leader-timing] / [special-timing]`` (圆桌 leader/special). - """ - t_entry = time.perf_counter() - bridge = get_stream_bridge(request) - run_mgr = get_run_manager(request) - run_ctx = get_run_context(request) - - disconnect = DisconnectMode.cancel if body.on_disconnect == "cancel" else DisconnectMode.continue_ - - t_pre_enforce = time.perf_counter() - # ``enforce_access=False`` is for intentionally open/unauthenticated entry - # points (e.g. /api/open/3qfx/ask) that resolve the agent themselves from an - # admin-configured mapping rather than from a per-user-visible assistant_id. - # Those run under the "default" bucket, where a user-owned unpublished agent - # would otherwise 404 the access check. + """Import existing folder-backed agents as built-in records. + + Heavy: ``list_custom_agents()`` scans ``.deer-flow/agents/`` and reads + ``config.yaml`` + ``SOUL.md`` per directory, then ``store.ensure_builtin`` + does one DB round-trip per agent. On a machine with N agent dirs this is + O(N) file IO + O(N) DB calls — historically the worst per-request hotspot + in ``start_run`` (observed 30s+ on machines with 100+ residual coordinator + shells + slow MySQL). **Never call this directly from request handlers** — + use :func:`_maybe_sync_legacy_agent_records` which applies the same TTL + gating as ``routers/agents._maybe_sync_legacy_agents``. + """ + from deerflow.config.agents_config import list_custom_agents + + for agent_cfg in list_custom_agents(): + agent_id = agent_cfg.id or agent_cfg.name + if not agent_id: + continue + await store.ensure_builtin( + { + "id": agent_id, + "name": agent_cfg.name, + "description": agent_cfg.description or "", + "published": True, + } + ) + + +# Throttle per-request legacy-sync. Matches ``routers/agents._LEGACY_SYNC_TTL_SECONDS`` +# semantics — the real sync runs once at startup (lifespan); this is a safety +# net for environments that bypass the lifespan hook or for freshly dropped +# agent directories. Lifespan registers itself via the same state by calling +# the agents router's ``_mark_legacy_sync_fresh()`` after its boot-time sync, +# so this module's TTL state is independent (a separate gate kept here so we +# don't import from ``routers/agents`` and create a circular import). +_LEGACY_SYNC_TTL_SECONDS = 3600.0 +_legacy_sync_state: dict[str, Any] = {"last_at": 0.0, "lock": None} + + +# Anything slower than this triggers the WARNING-level breakdown in +# ``start_run``. 500ms gives plenty of headroom for cold caches + a few DB +# round-trips on healthy systems while still catching the pathological +# multi-second cases that brought us here. +_START_RUN_SLOW_THRESHOLD_SECONDS = 0.5 + + +def _mark_legacy_sync_fresh() -> None: + """Tell this module's TTL gate that a sync just completed elsewhere. + + Called by the Gateway lifespan hook after it runs the boot-time legacy + agent sync (``routers/agents._sync_legacy_agents``). Without this, the + first request after startup would race past the TTL gate here and trigger + a redundant rescan even though the data is already fresh. + """ + _legacy_sync_state["last_at"] = time.monotonic() + + +async def _maybe_sync_legacy_agent_records(store) -> None: + """TTL-gated wrapper around :func:`_sync_legacy_agent_records`. + + Steady-state behavior: returns in microseconds inside the TTL window — no + filesystem walk, no DB round-trips. Outside the TTL window we serialize + callers under an :class:`asyncio.Lock` so a burst of concurrent + ``/runs/stream`` requests fans into a single rescan instead of N parallel + fan-outs. Pattern mirrors ``routers/agents._maybe_sync_legacy_agents``. + """ + now = time.monotonic() + if now - float(_legacy_sync_state["last_at"]) < _LEGACY_SYNC_TTL_SECONDS: + return + lock = _legacy_sync_state["lock"] + if lock is None: + lock = asyncio.Lock() + _legacy_sync_state["lock"] = lock + async with lock: + # Re-check under the lock — a sibling coroutine may have just refreshed. + if time.monotonic() - float(_legacy_sync_state["last_at"]) < _LEGACY_SYNC_TTL_SECONDS: + return + t_sync_start = time.perf_counter() + await _sync_legacy_agent_records(store) + elapsed = time.perf_counter() - t_sync_start + _legacy_sync_state["last_at"] = time.monotonic() + # When this fires it's a "first request of the hour" warmup — log it + # so operators can see the spike isn't an isolated bug. Anything + # over 2s usually means a lot of agent directories + slow IO/DB. + logger.info( + "[start-run-timing] legacy-agent sync warmed cache in %.3fs (next sync in %.0fs)", + elapsed, + _LEGACY_SYNC_TTL_SECONDS, + ) + + +async def _enforce_agent_access(body: Any, request: Request) -> None: + agent_ids = _collect_requested_agent_ids(body) + if not agent_ids: + return + + from app.gateway.deps import get_agent_store + from deerflow.config.agents_config import validate_agent_id + from deerflow.runtime.user_context import get_effective_user_id + + store = get_agent_store(request) + await _maybe_sync_legacy_agent_records(store) + user_id = get_effective_user_id() + # Admins see every agent in the admin console listing (?admin=true), so + # they may also start a run against any of them. + # ``request`` may be an in-process virtual request (SimpleNamespace) from + # the roundtable job gateway / scheduled-task runtime — those carry no + # ``state`` attribute, so the whole chain must be getattr-guarded + # (a bare ``request.state`` raises AttributeError and kills the run). + is_admin = getattr(getattr(getattr(request, "state", None), "user", None), "system_role", None) == "admin" + for agent_id in agent_ids: + try: + normalized = validate_agent_id(agent_id) + except ValueError as exc: + raise HTTPException(status_code=422, detail=str(exc)) from exc + if normalized is None: + continue + if await store.get_visible(normalized, user_id) is not None: + continue + if is_admin and await store.get_any(normalized) is not None: + continue + raise HTTPException(status_code=404, detail=f"Agent {normalized} not found") + + +# --------------------------------------------------------------------------- +# Run lifecycle +# --------------------------------------------------------------------------- + + +async def start_run( + body: Any, + thread_id: str, + request: Request, + *, + enforce_access: bool = True, +) -> RunRecord: + """Create a RunRecord and launch the background agent task. + + Parameters + ---------- + body : RunCreateRequest + The validated request body (typed as Any to avoid circular import + with the router module that defines the Pydantic model). + thread_id : str + Target thread. + request : Request + FastAPI request — used to retrieve singletons from ``app.state``. + + Timing observability + -------------------- + Each phase records ``perf_counter`` deltas. When the total wallclock from + function entry to ``run_agent`` task creation exceeds + ``_START_RUN_SLOW_THRESHOLD_SECONDS`` we emit a ``[start-run-timing]`` + breakdown at WARNING level. Stays at DEBUG when fast, so steady-state + logs are not polluted. Investigate spikes by grep'ing + ``[start-run-timing]`` together with ``[upstream-timing]`` (multi-agent + router) and ``[leader-timing] / [special-timing]`` (圆桌 leader/special). + """ + t_entry = time.perf_counter() + bridge = get_stream_bridge(request) + run_mgr = get_run_manager(request) + run_ctx = get_run_context(request) + + disconnect = DisconnectMode.cancel if body.on_disconnect == "cancel" else DisconnectMode.continue_ + + t_pre_enforce = time.perf_counter() + # ``enforce_access=False`` is for intentionally open/unauthenticated entry + # points (e.g. /api/open/3qfx/ask) that resolve the agent themselves from an + # admin-configured mapping rather than from a per-user-visible assistant_id. + # Those run under the "default" bucket, where a user-owned unpublished agent + # would otherwise 404 the access check. if enforce_access: await _enforce_agent_access(body, request) t_post_enforce = time.perf_counter() effective_model_name = _resolve_effective_run_model_name(body) - await _enforce_model_concurrency_limit(run_mgr, effective_model_name) + model_concurrency_limit = _get_model_concurrency_limit(effective_model_name) run_metadata = dict(body.metadata or {}) run_metadata["model_name"] = effective_model_name @@ -676,174 +701,176 @@ async def start_run( metadata=run_metadata, kwargs={"input": body.input, "config": body.config}, multitask_strategy=body.multitask_strategy, + model_concurrency_limit=(effective_model_name, model_concurrency_limit) if model_concurrency_limit is not None else None, ) - except ConflictError as exc: - # 同一会话上一轮回答还没结束，又提交了新问题 → 拒绝。给用户中文提示， - # 而不是把内部英文「Thread ... already has an active run」直接抛到前端。 - raise HTTPException( - status_code=409, - detail="当前对话还在生成回答，请等待本轮回复结束，或先点「停止」再发送新的问题。", - ) from exc - except UnsupportedStrategyError as exc: - raise HTTPException(status_code=501, detail=str(exc)) from exc - t_post_create = time.perf_counter() - - # Upsert thread metadata so the thread appears in /threads/search, - # even for threads that were never explicitly created via POST /threads - # (e.g. stateless runs). - requested_rag_mode = _rag_mode_requested(body) - effective_rag_mode = requested_rag_mode - try: - existing = await run_ctx.thread_store.get(thread_id) - effective_rag_mode = requested_rag_mode or _thread_rag_mode_enabled(existing) - if existing is None: - metadata = dict(body.metadata or {}) - if effective_rag_mode: - metadata["rag_mode_enabled"] = True - metadata["rag_mode"] = True - await run_ctx.thread_store.create( - thread_id, - assistant_id=body.assistant_id, - metadata=metadata, - ) - else: - await run_ctx.thread_store.update_status(thread_id, "running") - if requested_rag_mode and not _thread_rag_mode_enabled(existing): - await run_ctx.thread_store.update_metadata( - thread_id, - {"rag_mode_enabled": True, "rag_mode": True}, - ) - effective_rag_mode = True - # SDK 的 ``client.threads.create()`` 不传 assistant_id —— 先建出来空的 thread， - # 真正提交 run 时才知道用哪个 assistant。这里回填一次，避免历史列表 / 按 - # assistant 筛选时拿不到关联。 - if not existing.get("assistant_id") and body.assistant_id: - await run_ctx.thread_store.update_assistant_id(thread_id, body.assistant_id) - except Exception: - logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id)) - t_post_thread_upsert = time.perf_counter() - - agent_factory = resolve_agent_factory(body.assistant_id) - graph_input = normalize_input(body.input) - - # Admin-configurable prompt prefix: when the caller opted in via - # context.prompt_prefix_enabled, tag the human message with the configured - # prefix (stored in additional_kwargs, not concatenated into content, so the - # user never sees it — PromptPrefixMiddleware injects it for the LLM only). - # The frontend only sends the flag on the opening turn, so this naturally - # applies once. We intentionally do NOT gate on thread-row existence: the - # thread_meta row is frequently pre-created, which previously skipped the - # prefix for genuine first turns. - if _prompt_prefix_opted_in(body): - prefix_settings = load_system_settings().prompt_prefix - if prefix_settings.enabled and prefix_settings.prompt_prefix.strip(): - graph_input = _prepend_prompt_prefix(graph_input, prefix_settings.prompt_prefix.strip()) - if effective_rag_mode: - graph_input = _append_invisible_prompt(graph_input, _RAG_CITATION_PROMPT) - config = build_run_config(thread_id, body.config, body.metadata, assistant_id=body.assistant_id) - - # Merge DeerFlow-specific context overrides into both ``configurable`` and ``context``. - # The ``context`` field is a custom extension for the langgraph-compat layer - # that carries agent configuration (model_name, thinking_enabled, etc.). - # Only agent-relevant keys are forwarded; unknown keys (e.g. thread_id) are ignored. - body_context = dict(getattr(body, "context", None) or {}) - if effective_rag_mode: - body_context["rag_mode_enabled"] = True - # Promote whitelisted top-level body fields into the context map so - # callers can pass them either nested under ``context`` or as flat - # top-level fields (``excluded_tools``, ``skill_stop_names``, …). - for top_level_key in ("excluded_tools", "skill_stop_names", "skill_stop_message"): - value = getattr(body, top_level_key, None) - if value is not None and top_level_key not in body_context: - body_context[top_level_key] = value - merge_run_context_overrides(config, body_context or None) - - stream_modes = normalize_stream_modes(body.stream_mode) - - t_pre_task = time.perf_counter() - task = asyncio.create_task( - run_agent( - bridge, - run_mgr, - record, - ctx=run_ctx, - agent_factory=agent_factory, - graph_input=graph_input, - config=config, - stream_modes=stream_modes, - stream_subgraphs=body.stream_subgraphs, - interrupt_before=body.interrupt_before, - interrupt_after=body.interrupt_after, - # 透传 SDK ``Command``：resume / update / goto。Pregel 会据此从最近 - # checkpoint 处恢复，而不是把 graph_input 当新一轮输入跑。 - command=getattr(body, "command", None), - ) - ) - record.task = task - - # Title sync is handled by worker.py's finally block which reads the - # title from the checkpoint and calls thread_store.update_display_name - # after the run completes. - - total = t_pre_task - t_entry - if total >= _START_RUN_SLOW_THRESHOLD_SECONDS: - # Slow path: emit a phase breakdown so on-call can identify the - # culprit without re-running with extra instrumentation. The most - # common cause is a stale ``_maybe_sync_legacy_agent_records`` TTL - # window firing on a machine with hundreds of agent directories. - logger.warning( - "[start-run-timing] thread=%s total=%.3fs enforce_access=%.3fs " - "run_create=%.3fs thread_upsert=%.3fs build_payload=%.3fs", - sanitize_log_param(thread_id), - total, - t_post_enforce - t_pre_enforce, - t_post_create - t_post_enforce, - t_post_thread_upsert - t_post_create, - t_pre_task - t_post_thread_upsert, - ) - else: - logger.debug( - "[start-run-timing] thread=%s total=%.3fs enforce=%.3fs create=%.3fs upsert=%.3fs build=%.3fs", - sanitize_log_param(thread_id), - total, - t_post_enforce - t_pre_enforce, - t_post_create - t_post_enforce, - t_post_thread_upsert - t_post_create, - t_pre_task - t_post_thread_upsert, - ) - - return record - - -async def sse_consumer( - bridge: StreamBridge, - record: RunRecord, - request: Request, - run_mgr: RunManager, -): - """Async generator that yields SSE frames from the bridge. - - The ``finally`` block implements ``on_disconnect`` semantics: - - ``cancel``: abort the background task on client disconnect. - - ``continue``: let the task run; events are discarded. - """ - last_event_id = request.headers.get("Last-Event-ID") - try: - async for entry in bridge.subscribe(record.run_id, last_event_id=last_event_id): - if await request.is_disconnected(): - break - - if entry is HEARTBEAT_SENTINEL: - yield ": heartbeat\n\n" - continue - - if entry is END_SENTINEL: - yield format_sse("end", None, event_id=entry.id or None) - return - - yield format_sse(entry.event, entry.data, event_id=entry.id or None) - - finally: - if record.status in (RunStatus.pending, RunStatus.running): - if record.on_disconnect == DisconnectMode.cancel: - await run_mgr.cancel(record.run_id) + except ModelConcurrencyLimitError as exc: + raise HTTPException(status_code=429, detail="当前使用人数较多，请稍等一下再试。") from exc + except ConflictError as exc: + # 同一会话上一轮回答还没结束，又提交了新问题 → 拒绝。给用户中文提示， + # 而不是把内部英文「Thread ... already has an active run」直接抛到前端。 + raise HTTPException( + status_code=409, + detail="当前对话还在生成回答，请等待本轮回复结束，或先点「停止」再发送新的问题。", + ) from exc + except UnsupportedStrategyError as exc: + raise HTTPException(status_code=501, detail=str(exc)) from exc + t_post_create = time.perf_counter() + + # Upsert thread metadata so the thread appears in /threads/search, + # even for threads that were never explicitly created via POST /threads + # (e.g. stateless runs). + requested_rag_mode = _rag_mode_requested(body) + effective_rag_mode = requested_rag_mode + try: + existing = await run_ctx.thread_store.get(thread_id) + effective_rag_mode = requested_rag_mode or _thread_rag_mode_enabled(existing) + if existing is None: + metadata = dict(body.metadata or {}) + if effective_rag_mode: + metadata["rag_mode_enabled"] = True + metadata["rag_mode"] = True + await run_ctx.thread_store.create( + thread_id, + assistant_id=body.assistant_id, + metadata=metadata, + ) + else: + await run_ctx.thread_store.update_status(thread_id, "running") + if requested_rag_mode and not _thread_rag_mode_enabled(existing): + await run_ctx.thread_store.update_metadata( + thread_id, + {"rag_mode_enabled": True, "rag_mode": True}, + ) + effective_rag_mode = True + # SDK 的 ``client.threads.create()`` 不传 assistant_id —— 先建出来空的 thread， + # 真正提交 run 时才知道用哪个 assistant。这里回填一次，避免历史列表 / 按 + # assistant 筛选时拿不到关联。 + if not existing.get("assistant_id") and body.assistant_id: + await run_ctx.thread_store.update_assistant_id(thread_id, body.assistant_id) + except Exception: + logger.warning("Failed to upsert thread_meta for %s (non-fatal)", sanitize_log_param(thread_id)) + t_post_thread_upsert = time.perf_counter() + + agent_factory = resolve_agent_factory(body.assistant_id) + graph_input = normalize_input(body.input) + + # Admin-configurable prompt prefix: when the caller opted in via + # context.prompt_prefix_enabled, tag the human message with the configured + # prefix (stored in additional_kwargs, not concatenated into content, so the + # user never sees it — PromptPrefixMiddleware injects it for the LLM only). + # The frontend only sends the flag on the opening turn, so this naturally + # applies once. We intentionally do NOT gate on thread-row existence: the + # thread_meta row is frequently pre-created, which previously skipped the + # prefix for genuine first turns. + if _prompt_prefix_opted_in(body): + prefix_settings = load_system_settings().prompt_prefix + if prefix_settings.enabled and prefix_settings.prompt_prefix.strip(): + graph_input = _prepend_prompt_prefix(graph_input, prefix_settings.prompt_prefix.strip()) + if effective_rag_mode: + graph_input = _append_invisible_prompt(graph_input, _RAG_CITATION_PROMPT) + config = build_run_config(thread_id, body.config, body.metadata, assistant_id=body.assistant_id) + + # Merge DeerFlow-specific context overrides into both ``configurable`` and ``context``. + # The ``context`` field is a custom extension for the langgraph-compat layer + # that carries agent configuration (model_name, thinking_enabled, etc.). + # Only agent-relevant keys are forwarded; unknown keys (e.g. thread_id) are ignored. + body_context = dict(getattr(body, "context", None) or {}) + if effective_rag_mode: + body_context["rag_mode_enabled"] = True + # Promote whitelisted top-level body fields into the context map so + # callers can pass them either nested under ``context`` or as flat + # top-level fields (``excluded_tools``, ``skill_stop_names``, …). + for top_level_key in ("excluded_tools", "skill_stop_names", "skill_stop_message"): + value = getattr(body, top_level_key, None) + if value is not None and top_level_key not in body_context: + body_context[top_level_key] = value + merge_run_context_overrides(config, body_context or None) + + stream_modes = normalize_stream_modes(body.stream_mode) + + t_pre_task = time.perf_counter() + task = asyncio.create_task( + run_agent( + bridge, + run_mgr, + record, + ctx=run_ctx, + agent_factory=agent_factory, + graph_input=graph_input, + config=config, + stream_modes=stream_modes, + stream_subgraphs=body.stream_subgraphs, + interrupt_before=body.interrupt_before, + interrupt_after=body.interrupt_after, + # 透传 SDK ``Command``：resume / update / goto。Pregel 会据此从最近 + # checkpoint 处恢复，而不是把 graph_input 当新一轮输入跑。 + command=getattr(body, "command", None), + ) + ) + record.task = task + + # Title sync is handled by worker.py's finally block which reads the + # title from the checkpoint and calls thread_store.update_display_name + # after the run completes. + + total = t_pre_task - t_entry + if total >= _START_RUN_SLOW_THRESHOLD_SECONDS: + # Slow path: emit a phase breakdown so on-call can identify the + # culprit without re-running with extra instrumentation. The most + # common cause is a stale ``_maybe_sync_legacy_agent_records`` TTL + # window firing on a machine with hundreds of agent directories. + logger.warning( + "[start-run-timing] thread=%s total=%.3fs enforce_access=%.3fs run_create=%.3fs thread_upsert=%.3fs build_payload=%.3fs", + sanitize_log_param(thread_id), + total, + t_post_enforce - t_pre_enforce, + t_post_create - t_post_enforce, + t_post_thread_upsert - t_post_create, + t_pre_task - t_post_thread_upsert, + ) + else: + logger.debug( + "[start-run-timing] thread=%s total=%.3fs enforce=%.3fs create=%.3fs upsert=%.3fs build=%.3fs", + sanitize_log_param(thread_id), + total, + t_post_enforce - t_pre_enforce, + t_post_create - t_post_enforce, + t_post_thread_upsert - t_post_create, + t_pre_task - t_post_thread_upsert, + ) + + return record + + +async def sse_consumer( + bridge: StreamBridge, + record: RunRecord, + request: Request, + run_mgr: RunManager, +): + """Async generator that yields SSE frames from the bridge. + + The ``finally`` block implements ``on_disconnect`` semantics: + - ``cancel``: abort the background task on client disconnect. + - ``continue``: let the task run; events are discarded. + """ + last_event_id = request.headers.get("Last-Event-ID") + try: + async for entry in bridge.subscribe(record.run_id, last_event_id=last_event_id): + if await request.is_disconnected(): + break + + if entry is HEARTBEAT_SENTINEL: + yield ": heartbeat\n\n" + continue + + if entry is END_SENTINEL: + yield format_sse("end", None, event_id=entry.id or None) + return + + yield format_sse(entry.event, entry.data, event_id=entry.id or None) + + finally: + if record.status in (RunStatus.pending, RunStatus.running): + if record.on_disconnect == DisconnectMode.cancel: + await run_mgr.cancel(record.run_id) diff --git a/offline-backend-20260512/backend/packages/harness/deerflow/runtime/__init__.py b/offline-backend-20260512/backend/packages/harness/deerflow/runtime/__init__.py index a6b87e59..63ea0860 100644 --- a/offline-backend-20260512/backend/packages/harness/deerflow/runtime/__init__.py +++ b/offline-backend-20260512/backend/packages/harness/deerflow/runtime/__init__.py @@ -1,46 +1,57 @@ -"""LangGraph-compatible runtime — runs, streaming, and lifecycle management. - -Re-exports the public API of :mod:`~deerflow.runtime.runs` and -:mod:`~deerflow.runtime.stream_bridge` so that consumers can import -directly from ``deerflow.runtime``. -""" - -from .checkpointer import checkpointer_context, get_checkpointer, make_checkpointer, reset_checkpointer -from .runs import ConflictError, DisconnectMode, RunContext, RunManager, RunRecord, RunStatus, UnsupportedStrategyError, run_agent -from .serialization import serialize, serialize_channel_values, serialize_lc_object, serialize_messages_tuple -from .store import get_store, make_store, reset_store, store_context -from .stream_bridge import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamBridge, StreamEvent, make_stream_bridge - -__all__ = [ - # checkpointer - "checkpointer_context", - "get_checkpointer", - "make_checkpointer", - "reset_checkpointer", - # runs - "ConflictError", - "DisconnectMode", - "RunContext", - "RunManager", - "RunRecord", - "RunStatus", - "UnsupportedStrategyError", - "run_agent", - # serialization - "serialize", - "serialize_channel_values", - "serialize_lc_object", - "serialize_messages_tuple", - # store - "get_store", - "make_store", - "reset_store", - "store_context", - # stream_bridge - "END_SENTINEL", - "HEARTBEAT_SENTINEL", - "MemoryStreamBridge", - "StreamBridge", - "StreamEvent", - "make_stream_bridge", -] +"""LangGraph-compatible runtime — runs, streaming, and lifecycle management. + +Re-exports the public API of :mod:`~deerflow.runtime.runs` and +:mod:`~deerflow.runtime.stream_bridge` so that consumers can import +directly from ``deerflow.runtime``. +""" + +from .checkpointer import checkpointer_context, get_checkpointer, make_checkpointer, reset_checkpointer +from .runs import ( + ConflictError, + DisconnectMode, + ModelConcurrencyLimitError, + RunContext, + RunManager, + RunRecord, + RunStatus, + UnsupportedStrategyError, + run_agent, +) +from .serialization import serialize, serialize_channel_values, serialize_lc_object, serialize_messages_tuple +from .store import get_store, make_store, reset_store, store_context +from .stream_bridge import END_SENTINEL, HEARTBEAT_SENTINEL, MemoryStreamBridge, StreamBridge, StreamEvent, make_stream_bridge + +__all__ = [ + # checkpointer + "checkpointer_context", + "get_checkpointer", + "make_checkpointer", + "reset_checkpointer", + # runs + "ConflictError", + "DisconnectMode", + "ModelConcurrencyLimitError", + "RunContext", + "RunManager", + "RunRecord", + "RunStatus", + "UnsupportedStrategyError", + "run_agent", + # serialization + "serialize", + "serialize_channel_values", + "serialize_lc_object", + "serialize_messages_tuple", + # store + "get_store", + "make_store", + "reset_store", + "store_context", + # stream_bridge + "END_SENTINEL", + "HEARTBEAT_SENTINEL", + "MemoryStreamBridge", + "StreamBridge", + "StreamEvent", + "make_stream_bridge", +] diff --git a/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/__init__.py b/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/__init__.py index da5fc787..318b58c5 100644 --- a/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/__init__.py +++ b/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/__init__.py @@ -1,16 +1,17 @@ -"""Run lifecycle management for LangGraph Platform API compatibility.""" - -from .manager import ConflictError, RunManager, RunRecord, UnsupportedStrategyError -from .schemas import DisconnectMode, RunStatus -from .worker import RunContext, run_agent - -__all__ = [ - "ConflictError", - "DisconnectMode", - "RunContext", - "RunManager", - "RunRecord", - "RunStatus", - "UnsupportedStrategyError", - "run_agent", -] +"""Run lifecycle management for LangGraph Platform API compatibility.""" + +from .manager import ConflictError, ModelConcurrencyLimitError, RunManager, RunRecord, UnsupportedStrategyError +from .schemas import DisconnectMode, RunStatus +from .worker import RunContext, run_agent + +__all__ = [ + "ConflictError", + "DisconnectMode", + "ModelConcurrencyLimitError", + "RunContext", + "RunManager", + "RunRecord", + "RunStatus", + "UnsupportedStrategyError", + "run_agent", +] diff --git a/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/manager.py b/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/manager.py index a0783d87..b8b0a6c6 100644 --- a/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/offline-backend-20260512/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -1,310 +1,326 @@ -"""In-memory run registry with optional persistent RunStore backing.""" - -from __future__ import annotations - -import asyncio -import logging -import uuid -from dataclasses import dataclass, field -from typing import TYPE_CHECKING - -from deerflow.utils.time import now_iso as _now_iso - -from .schemas import DisconnectMode, RunStatus - -if TYPE_CHECKING: - from deerflow.runtime.runs.store.base import RunStore - -logger = logging.getLogger(__name__) - - -@dataclass -class RunRecord: - """Mutable record for a single run.""" - - run_id: str - thread_id: str - assistant_id: str | None - status: RunStatus - on_disconnect: DisconnectMode - multitask_strategy: str = "reject" - metadata: dict = field(default_factory=dict) - kwargs: dict = field(default_factory=dict) - created_at: str = "" - updated_at: str = "" - task: asyncio.Task | None = field(default=None, repr=False) - abort_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False) - abort_action: str = "interrupt" - error: str | None = None - - -class RunManager: - """In-memory run registry with optional persistent RunStore backing. - - All mutations are protected by an asyncio lock. When a ``store`` is - provided, serializable metadata is also persisted to the store so - that run history survives process restarts. - """ - - def __init__(self, store: RunStore | None = None) -> None: - self._runs: dict[str, RunRecord] = {} - self._lock = asyncio.Lock() - self._store = store - - async def _persist_to_store(self, record: RunRecord) -> None: - """Best-effort persist run record to backing store.""" - if self._store is None: - return - try: - await self._store.put( - record.run_id, - thread_id=record.thread_id, - assistant_id=record.assistant_id, - status=record.status.value, - multitask_strategy=record.multitask_strategy, - metadata=record.metadata or {}, - kwargs=record.kwargs or {}, - created_at=record.created_at, - ) - except Exception: - logger.warning("Failed to persist run %s to store", record.run_id, exc_info=True) - - async def update_run_completion(self, run_id: str, **kwargs) -> None: - """Persist token usage and completion data to the backing store.""" - if self._store is not None: - try: - await self._store.update_run_completion(run_id, **kwargs) - except Exception: - logger.warning("Failed to persist run completion for %s", run_id, exc_info=True) - - async def create( - self, - thread_id: str, - assistant_id: str | None = None, - *, - on_disconnect: DisconnectMode = DisconnectMode.cancel, - metadata: dict | None = None, - kwargs: dict | None = None, - multitask_strategy: str = "reject", - ) -> RunRecord: - """Create a new pending run and register it.""" - run_id = str(uuid.uuid4()) - now = _now_iso() - record = RunRecord( - run_id=run_id, - thread_id=thread_id, - assistant_id=assistant_id, - status=RunStatus.pending, - on_disconnect=on_disconnect, - multitask_strategy=multitask_strategy, - metadata=metadata or {}, - kwargs=kwargs or {}, - created_at=now, - updated_at=now, - ) - async with self._lock: - self._runs[run_id] = record - await self._persist_to_store(record) - logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) - return record - - def get(self, run_id: str) -> RunRecord | None: - """Return a run record by ID, or ``None``.""" - return self._runs.get(run_id) - - def _record_from_store(self, data: dict) -> RunRecord: - """Rebuild a :class:`RunRecord` from a persisted store row. - - The reconstructed record has no live ``task``/``abort_event`` wired to a - running coroutine — it represents *history*, not an in-flight run. - """ - try: - status = RunStatus(data.get("status") or RunStatus.success.value) - except ValueError: - status = RunStatus.success - return RunRecord( - run_id=data["run_id"], - thread_id=data["thread_id"], - assistant_id=data.get("assistant_id"), - status=status, - # Disconnect mode is irrelevant for a store-only record (no task to - # cancel); ``continue_`` keeps sse_consumer's finally block a no-op. - on_disconnect=DisconnectMode.continue_, - multitask_strategy=data.get("multitask_strategy", "reject"), - metadata=data.get("metadata") or {}, - kwargs=data.get("kwargs") or {}, - created_at=data.get("created_at") or "", - updated_at=data.get("updated_at") or "", - error=data.get("error"), - ) - - async def get_persisted(self, run_id: str) -> RunRecord | None: - """Load a run from the backing store (process-restart fallback). - - The in-memory registry is wiped on every Gateway restart, but the - ``RunStore`` row survives. When :meth:`get` misses, callers fall back - here so a stale ``joinStream``/``GET run`` can report the run's final - status instead of 404ing. Returns ``None`` when there is no store, no - matching row, or the lookup fails (e.g. no user context). - """ - if self._store is None: - return None - try: - data = await self._store.get(run_id) - except Exception: - logger.warning("Failed to load run %s from store", run_id, exc_info=True) - return None - if not data: - return None - return self._record_from_store(data) - - async def list_by_thread(self, thread_id: str) -> list[RunRecord]: - """Return all runs for a given thread, newest first.""" - async with self._lock: - # Dict insertion order matches creation order, so reversing it gives - # us deterministic newest-first results even when timestamps tie. - return [r for r in self._runs.values() if r.thread_id == thread_id] - - async def set_status(self, run_id: str, status: RunStatus, *, error: str | None = None) -> None: - """Transition a run to a new status.""" - async with self._lock: - record = self._runs.get(run_id) - if record is None: - logger.warning("set_status called for unknown run %s", run_id) - return - record.status = status - record.updated_at = _now_iso() - if error is not None: - record.error = error - if self._store is not None: - try: - await self._store.update_status(run_id, status.value, error=error) - except Exception: - logger.warning("Failed to persist status update for run %s", run_id, exc_info=True) - logger.info("Run %s -> %s", run_id, status.value) - - async def cancel(self, run_id: str, *, action: str = "interrupt") -> bool: - """Request cancellation of a run. - - Args: - run_id: The run ID to cancel. - action: "interrupt" keeps checkpoint, "rollback" reverts to pre-run state. - - Sets the abort event with the action reason and cancels the asyncio task. - Returns ``True`` if the run was in-flight and cancellation was initiated. - """ - async with self._lock: - record = self._runs.get(run_id) - if record is None: - return False - if record.status not in (RunStatus.pending, RunStatus.running): - return False - record.abort_action = action - record.abort_event.set() - if record.task is not None and not record.task.done(): - record.task.cancel() - record.status = RunStatus.interrupted - record.updated_at = _now_iso() - logger.info("Run %s cancelled (action=%s)", run_id, action) - return True - - async def create_or_reject( - self, - thread_id: str, - assistant_id: str | None = None, - *, - on_disconnect: DisconnectMode = DisconnectMode.cancel, - metadata: dict | None = None, - kwargs: dict | None = None, - multitask_strategy: str = "reject", - ) -> RunRecord: - """Atomically check for inflight runs and create a new one. - - For ``reject`` strategy, raises ``ConflictError`` if thread - already has a pending/running run. For ``interrupt``/``rollback``, - cancels inflight runs before creating. - - This method holds the lock across both the check and the insert, - eliminating the TOCTOU race in separate ``has_inflight`` + ``create``. - """ - run_id = str(uuid.uuid4()) - now = _now_iso() - - _supported_strategies = ("reject", "interrupt", "rollback") - - async with self._lock: - if multitask_strategy not in _supported_strategies: - raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}") - - inflight = [r for r in self._runs.values() if r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running)] - - if multitask_strategy == "reject" and inflight: - raise ConflictError(f"Thread {thread_id} already has an active run") - - if multitask_strategy in ("interrupt", "rollback") and inflight: - for r in inflight: - r.abort_action = multitask_strategy - r.abort_event.set() - if r.task is not None and not r.task.done(): - r.task.cancel() - r.status = RunStatus.interrupted - r.updated_at = now - logger.info( - "Cancelled %d inflight run(s) on thread %s (strategy=%s)", - len(inflight), - thread_id, - multitask_strategy, - ) - - record = RunRecord( - run_id=run_id, - thread_id=thread_id, - assistant_id=assistant_id, - status=RunStatus.pending, - on_disconnect=on_disconnect, - multitask_strategy=multitask_strategy, - metadata=metadata or {}, - kwargs=kwargs or {}, - created_at=now, - updated_at=now, - ) - self._runs[run_id] = record - - await self._persist_to_store(record) - logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) - return record - - async def has_inflight(self, thread_id: str) -> bool: - """Return ``True`` if *thread_id* has a pending or running run.""" - async with self._lock: - return any(r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running) for r in self._runs.values()) - - async def list_active(self) -> list[RunRecord]: - """Return every in-flight run (status pending/running), oldest first. - - Powers the admin concurrency monitor: each record represents one live - call into the model runtime (interactive chat, agent chat, roundtable - seat, …), so ``len(list_active())`` is the current concurrent-call count - and each record can be cancelled to relieve model pressure. Completed - records linger in ``_runs`` until :meth:`cleanup`, so we filter by status - rather than returning the whole registry. - """ - async with self._lock: - # Dict insertion order == creation order, so this is oldest-first. - return [r for r in self._runs.values() if r.status in (RunStatus.pending, RunStatus.running)] - - async def cleanup(self, run_id: str, *, delay: float = 300) -> None: - """Remove a run record after an optional delay.""" - if delay > 0: - await asyncio.sleep(delay) - async with self._lock: - self._runs.pop(run_id, None) - logger.debug("Run record %s cleaned up", run_id) - - -class ConflictError(Exception): - """Raised when multitask_strategy=reject and thread has inflight runs.""" - - -class UnsupportedStrategyError(Exception): - """Raised when a multitask_strategy value is not yet implemented.""" +"""In-memory run registry with optional persistent RunStore backing.""" + +from __future__ import annotations + +import asyncio +import logging +import uuid +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from deerflow.utils.time import now_iso as _now_iso + +from .schemas import DisconnectMode, RunStatus + +if TYPE_CHECKING: + from deerflow.runtime.runs.store.base import RunStore + +logger = logging.getLogger(__name__) + + +@dataclass +class RunRecord: + """Mutable record for a single run.""" + + run_id: str + thread_id: str + assistant_id: str | None + status: RunStatus + on_disconnect: DisconnectMode + multitask_strategy: str = "reject" + metadata: dict = field(default_factory=dict) + kwargs: dict = field(default_factory=dict) + created_at: str = "" + updated_at: str = "" + task: asyncio.Task | None = field(default=None, repr=False) + abort_event: asyncio.Event = field(default_factory=asyncio.Event, repr=False) + abort_action: str = "interrupt" + error: str | None = None + + +class RunManager: + """In-memory run registry with optional persistent RunStore backing. + + All mutations are protected by an asyncio lock. When a ``store`` is + provided, serializable metadata is also persisted to the store so + that run history survives process restarts. + """ + + def __init__(self, store: RunStore | None = None) -> None: + self._runs: dict[str, RunRecord] = {} + self._lock = asyncio.Lock() + self._store = store + + async def _persist_to_store(self, record: RunRecord) -> None: + """Best-effort persist run record to backing store.""" + if self._store is None: + return + try: + await self._store.put( + record.run_id, + thread_id=record.thread_id, + assistant_id=record.assistant_id, + status=record.status.value, + multitask_strategy=record.multitask_strategy, + metadata=record.metadata or {}, + kwargs=record.kwargs or {}, + created_at=record.created_at, + ) + except Exception: + logger.warning("Failed to persist run %s to store", record.run_id, exc_info=True) + + async def update_run_completion(self, run_id: str, **kwargs) -> None: + """Persist token usage and completion data to the backing store.""" + if self._store is not None: + try: + await self._store.update_run_completion(run_id, **kwargs) + except Exception: + logger.warning("Failed to persist run completion for %s", run_id, exc_info=True) + + async def create( + self, + thread_id: str, + assistant_id: str | None = None, + *, + on_disconnect: DisconnectMode = DisconnectMode.cancel, + metadata: dict | None = None, + kwargs: dict | None = None, + multitask_strategy: str = "reject", + ) -> RunRecord: + """Create a new pending run and register it.""" + run_id = str(uuid.uuid4()) + now = _now_iso() + record = RunRecord( + run_id=run_id, + thread_id=thread_id, + assistant_id=assistant_id, + status=RunStatus.pending, + on_disconnect=on_disconnect, + multitask_strategy=multitask_strategy, + metadata=metadata or {}, + kwargs=kwargs or {}, + created_at=now, + updated_at=now, + ) + async with self._lock: + self._runs[run_id] = record + await self._persist_to_store(record) + logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) + return record + + def get(self, run_id: str) -> RunRecord | None: + """Return a run record by ID, or ``None``.""" + return self._runs.get(run_id) + + def _record_from_store(self, data: dict) -> RunRecord: + """Rebuild a :class:`RunRecord` from a persisted store row. + + The reconstructed record has no live ``task``/``abort_event`` wired to a + running coroutine — it represents *history*, not an in-flight run. + """ + try: + status = RunStatus(data.get("status") or RunStatus.success.value) + except ValueError: + status = RunStatus.success + return RunRecord( + run_id=data["run_id"], + thread_id=data["thread_id"], + assistant_id=data.get("assistant_id"), + status=status, + # Disconnect mode is irrelevant for a store-only record (no task to + # cancel); ``continue_`` keeps sse_consumer's finally block a no-op. + on_disconnect=DisconnectMode.continue_, + multitask_strategy=data.get("multitask_strategy", "reject"), + metadata=data.get("metadata") or {}, + kwargs=data.get("kwargs") or {}, + created_at=data.get("created_at") or "", + updated_at=data.get("updated_at") or "", + error=data.get("error"), + ) + + async def get_persisted(self, run_id: str) -> RunRecord | None: + """Load a run from the backing store (process-restart fallback). + + The in-memory registry is wiped on every Gateway restart, but the + ``RunStore`` row survives. When :meth:`get` misses, callers fall back + here so a stale ``joinStream``/``GET run`` can report the run's final + status instead of 404ing. Returns ``None`` when there is no store, no + matching row, or the lookup fails (e.g. no user context). + """ + if self._store is None: + return None + try: + data = await self._store.get(run_id) + except Exception: + logger.warning("Failed to load run %s from store", run_id, exc_info=True) + return None + if not data: + return None + return self._record_from_store(data) + + async def list_by_thread(self, thread_id: str) -> list[RunRecord]: + """Return all runs for a given thread, newest first.""" + async with self._lock: + # Dict insertion order matches creation order, so reversing it gives + # us deterministic newest-first results even when timestamps tie. + return [r for r in self._runs.values() if r.thread_id == thread_id] + + async def set_status(self, run_id: str, status: RunStatus, *, error: str | None = None) -> None: + """Transition a run to a new status.""" + async with self._lock: + record = self._runs.get(run_id) + if record is None: + logger.warning("set_status called for unknown run %s", run_id) + return + record.status = status + record.updated_at = _now_iso() + if error is not None: + record.error = error + if self._store is not None: + try: + await self._store.update_status(run_id, status.value, error=error) + except Exception: + logger.warning("Failed to persist status update for run %s", run_id, exc_info=True) + logger.info("Run %s -> %s", run_id, status.value) + + async def cancel(self, run_id: str, *, action: str = "interrupt") -> bool: + """Request cancellation of a run. + + Args: + run_id: The run ID to cancel. + action: "interrupt" keeps checkpoint, "rollback" reverts to pre-run state. + + Sets the abort event with the action reason and cancels the asyncio task. + Returns ``True`` if the run was in-flight and cancellation was initiated. + """ + async with self._lock: + record = self._runs.get(run_id) + if record is None: + return False + if record.status not in (RunStatus.pending, RunStatus.running): + return False + record.abort_action = action + record.abort_event.set() + if record.task is not None and not record.task.done(): + record.task.cancel() + record.status = RunStatus.interrupted + record.updated_at = _now_iso() + logger.info("Run %s cancelled (action=%s)", run_id, action) + return True + + async def create_or_reject( + self, + thread_id: str, + assistant_id: str | None = None, + *, + on_disconnect: DisconnectMode = DisconnectMode.cancel, + metadata: dict | None = None, + kwargs: dict | None = None, + multitask_strategy: str = "reject", + model_concurrency_limit: tuple[str, int] | None = None, + ) -> RunRecord: + """Atomically check for inflight runs and create a new one. + + For ``reject`` strategy, raises ``ConflictError`` if thread + already has a pending/running run. For ``interrupt``/``rollback``, + cancels inflight runs before creating. + + This method holds the lock across both the check and the insert, + eliminating the TOCTOU race in separate ``has_inflight`` + ``create``. + """ + run_id = str(uuid.uuid4()) + now = _now_iso() + + _supported_strategies = ("reject", "interrupt", "rollback") + + async with self._lock: + if multitask_strategy not in _supported_strategies: + raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}") + + inflight = [r for r in self._runs.values() if r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running)] + + if multitask_strategy == "reject" and inflight: + raise ConflictError(f"Thread {thread_id} already has an active run") + + if multitask_strategy in ("interrupt", "rollback") and inflight: + for r in inflight: + r.abort_action = multitask_strategy + r.abort_event.set() + if r.task is not None and not r.task.done(): + r.task.cancel() + r.status = RunStatus.interrupted + r.updated_at = now + logger.info( + "Cancelled %d inflight run(s) on thread %s (strategy=%s)", + len(inflight), + thread_id, + multitask_strategy, + ) + + if model_concurrency_limit is not None: + model_name, limit = model_concurrency_limit + active_for_model = sum(1 for r in self._runs.values() if r.status in (RunStatus.pending, RunStatus.running) and str((r.metadata or {}).get("model_name") or "").strip() == model_name) + if active_for_model >= max(1, int(limit)): + raise ModelConcurrencyLimitError(model_name, limit) + + record = RunRecord( + run_id=run_id, + thread_id=thread_id, + assistant_id=assistant_id, + status=RunStatus.pending, + on_disconnect=on_disconnect, + multitask_strategy=multitask_strategy, + metadata=metadata or {}, + kwargs=kwargs or {}, + created_at=now, + updated_at=now, + ) + self._runs[run_id] = record + + await self._persist_to_store(record) + logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id) + return record + + async def has_inflight(self, thread_id: str) -> bool: + """Return ``True`` if *thread_id* has a pending or running run.""" + async with self._lock: + return any(r.thread_id == thread_id and r.status in (RunStatus.pending, RunStatus.running) for r in self._runs.values()) + + async def list_active(self) -> list[RunRecord]: + """Return every in-flight run (status pending/running), oldest first. + + Powers the admin concurrency monitor: each record represents one live + call into the model runtime (interactive chat, agent chat, roundtable + seat, …), so ``len(list_active())`` is the current concurrent-call count + and each record can be cancelled to relieve model pressure. Completed + records linger in ``_runs`` until :meth:`cleanup`, so we filter by status + rather than returning the whole registry. + """ + async with self._lock: + # Dict insertion order == creation order, so this is oldest-first. + return [r for r in self._runs.values() if r.status in (RunStatus.pending, RunStatus.running)] + + async def cleanup(self, run_id: str, *, delay: float = 300) -> None: + """Remove a run record after an optional delay.""" + if delay > 0: + await asyncio.sleep(delay) + async with self._lock: + self._runs.pop(run_id, None) + logger.debug("Run record %s cleaned up", run_id) + + +class ConflictError(Exception): + """Raised when multitask_strategy=reject and thread has inflight runs.""" + + +class ModelConcurrencyLimitError(Exception): + """Raised when a model-specific in-flight run limit has been reached.""" + + def __init__(self, model_name: str, limit: int) -> None: + self.model_name = model_name + self.limit = limit + super().__init__(f"Model {model_name} has reached its concurrency limit") + + +class UnsupportedStrategyError(Exception): + """Raised when a multitask_strategy value is not yet implemented.""" diff --git a/offline-backend-20260512/backend/tests/test_model_concurrency.py b/offline-backend-20260512/backend/tests/test_model_concurrency.py index 0c14f59a..916c0bbf 100644 --- a/offline-backend-20260512/backend/tests/test_model_concurrency.py +++ b/offline-backend-20260512/backend/tests/test_model_concurrency.py @@ -7,15 +7,14 @@ from pathlib import Path from types import SimpleNamespace import pytest -from fastapi import HTTPException from app.gateway.routers import system_settings as ss_router from app.gateway.services import ( - _enforce_model_concurrency_limit, + _get_model_concurrency_limit, _resolve_effective_run_model_name, ) from deerflow.config import system_settings as ss -from deerflow.runtime.runs.manager import RunManager +from deerflow.runtime.runs.manager import ModelConcurrencyLimitError, RunManager @pytest.fixture @@ -77,9 +76,7 @@ def test_resolve_effective_run_model_name_prefers_requested_model( ) -> None: app_config = SimpleNamespace( models=[SimpleNamespace(name="default-model")], - get_model_config=lambda name: {"name": name} - if name in {"default-model", "glm5", "agent-model"} - else None, + get_model_config=lambda name: {"name": name} if name in {"default-model", "glm5", "agent-model"} else None, ) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr( @@ -100,9 +97,7 @@ def test_resolve_effective_run_model_name_uses_agent_default_when_request_missin ) -> None: app_config = SimpleNamespace( models=[SimpleNamespace(name="default-model")], - get_model_config=lambda name: {"name": name} - if name in {"default-model", "agent-model"} - else None, + get_model_config=lambda name: {"name": name} if name in {"default-model", "agent-model"} else None, ) monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) monkeypatch.setattr( @@ -118,8 +113,50 @@ def test_resolve_effective_run_model_name_uses_agent_default_when_request_missin assert _resolve_effective_run_model_name(body) == "agent-model" +def test_resolve_effective_run_model_name_uses_context_agent_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + app_config = SimpleNamespace( + models=[SimpleNamespace(name="default-model")], + get_model_config=lambda name: {"name": name} if name in {"default-model", "agent-model"} else None, + ) + monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) + monkeypatch.setattr( + "app.gateway.services.load_agent_config", + lambda agent_id: SimpleNamespace(model="agent-model") if agent_id == "custom-agent" else SimpleNamespace(model="default-model"), + ) + + body = SimpleNamespace( + assistant_id="lead_agent", + context={"agent_id": "custom-agent"}, + config=None, + ) + assert _resolve_effective_run_model_name(body) == "agent-model" + + +def test_resolve_effective_run_model_name_context_agent_overrides_assistant_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + app_config = SimpleNamespace( + models=[SimpleNamespace(name="default-model")], + get_model_config=lambda name: {"name": name} if name in {"default-model", "context-model", "assistant-model"} else None, + ) + monkeypatch.setattr("app.gateway.services.get_app_config", lambda: app_config) + monkeypatch.setattr( + "app.gateway.services.load_agent_config", + lambda agent_id: SimpleNamespace(model="context-model" if agent_id == "context-agent" else "assistant-model"), + ) + + body = SimpleNamespace( + assistant_id="assistant-agent", + context={"agent_id": "context-agent"}, + config=None, + ) + assert _resolve_effective_run_model_name(body) == "context-model" + + @pytest.mark.asyncio -async def test_enforce_model_concurrency_limit_rejects_when_limit_reached( +async def test_run_manager_rejects_model_limit_atomically( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( @@ -132,18 +169,21 @@ async def test_enforce_model_concurrency_limit_rejects_when_limit_reached( ) ), ) + assert _get_model_concurrency_limit("glm5") == 2 + run_mgr = RunManager() - await run_mgr.create("t1", metadata={"model_name": "glm5"}) - await run_mgr.create("t2", metadata={"model_name": "glm5"}) + await run_mgr.create_or_reject("t1", metadata={"model_name": "glm5"}) + await run_mgr.create_or_reject("t2", metadata={"model_name": "glm5"}) - with pytest.raises(HTTPException) as exc: - await _enforce_model_concurrency_limit(run_mgr, "glm5") - assert exc.value.status_code == 429 - assert "并发已达到上限" in str(exc.value.detail) + with pytest.raises(ModelConcurrencyLimitError): + await run_mgr.create_or_reject( + "t3", + metadata={"model_name": "glm5"}, + model_concurrency_limit=("glm5", 2), + ) -@pytest.mark.asyncio -async def test_enforce_model_concurrency_limit_skips_when_disabled( +def test_get_model_concurrency_limit_skips_when_disabled( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( @@ -156,6 +196,4 @@ async def test_enforce_model_concurrency_limit_skips_when_disabled( ) ), ) - run_mgr = RunManager() - await run_mgr.create("t1", metadata={"model_name": "glm5"}) - await _enforce_model_concurrency_limit(run_mgr, "glm5") + assert _get_model_concurrency_limit("glm5") is None