diff --git a/app/agent/__init__.py b/app/agent/__init__.py index d24e651a8..1f2533435 100644 --- a/app/agent/__init__.py +++ b/app/agent/__init__.py @@ -66,6 +66,7 @@ from app.agent.tools.impl.mcp import ( create_external_mcp_tools, select_legacy_mcp_tools, ) +from app.agent.tools.impl.query_system_settings import QuerySystemSettingsTool from app.chain import ChainBase from app.core.config import settings from app.core.event import eventmanager @@ -301,6 +302,22 @@ AGENT_CHAT_TITLE_PROMPT = ( ) AGENT_CHAT_TITLE_MAX_LENGTH = 36 AGENT_CHAT_TITLE_MAX_CJK_CHARS = 18 +SECRET_CONFIRMATION_TTL = timedelta(minutes=5) +SECRET_CONFIRM_TEXT = "确认" +SECRET_CANCEL_TEXT = "取消" + + +@dataclass +class _PendingSecretConfirmation: + """保存当前会话中一次待确认的敏感设置读取。""" + + tool: QuerySystemSettingsTool + arguments: Dict[str, Any] + created_at: datetime + user_id: str + channel: str + source: str + original_chat_id: str class MoviePilotAgent: @@ -322,6 +339,7 @@ class MoviePilotAgent: replay_mode: ReplyMode = ReplyMode.DISPATCH, allow_message_tools: bool = True, output_callback: Optional[Callable[[str], None]] = None, + protected_output_callback: Optional[Callable[[str], None]] = None, ): self.session_id = session_id self.user_id = user_id @@ -333,7 +351,9 @@ class MoviePilotAgent: self.reply_mode = replay_mode self.allow_message_tools = allow_message_tools self.output_callback = output_callback + self.protected_output_callback = protected_output_callback self._tool_context: Dict[str, object] = {} + self._pending_secret_confirmation: Optional[_PendingSecretConfirmation] = None self._streamed_output = "" self._session_usage = _SessionUsageSnapshot() self._llm_runtime_config: Optional[Dict[str, Any]] = None @@ -742,11 +762,169 @@ class MoviePilotAgent: "reply_mode": None, "should_dispatch_reply": should_dispatch_reply, "is_admin": await self._is_system_admin_context(), + "require_secret_confirmation": True, + "secret_confirmation_handler": self._register_secret_confirmation, # 工具回调消息需要发回原会话（群聊@机器人时按钮选择等卡片不能发到私聊）， # 后台任务无渠道上下文时置空，交由通知链广播。 "original_chat_id": None if self.is_background else self.original_chat_id, } + def set_protected_output_callback( + self, + protected_output_callback: Optional[Callable[[str], None]], + ) -> None: + """更新仅供当前请求接收的受保护文本输出回调。""" + self.protected_output_callback = protected_output_callback + + def has_pending_secret_confirmation(self) -> bool: + """判断当前会话是否存在仍在有效期内的敏感设置确认。""" + pending = self._pending_secret_confirmation + if not pending: + return False + if datetime.now() - pending.created_at <= SECRET_CONFIRMATION_TTL: + return True + self._pending_secret_confirmation = None + return False + + def _can_confirm_secret_read(self) -> bool: + """判断当前渠道能否把密钥结果直接交付给原用户。""" + if self.channel == MessageChannel.WebAgent.value: + return callable(self.protected_output_callback) + return self.channel in { + MessageChannel.Telegram.value, + MessageChannel.Feishu.value, + } + + async def _register_secret_confirmation( + self, + tool: QuerySystemSettingsTool, + arguments: Dict[str, Any], + ) -> str: + """校验并冻结一次待用户确认的敏感设置读取。""" + if self.has_pending_secret_confirmation(): + return "当前会话已有待确认的敏感设置读取，请先回复“确认”或“取消”。" + if not self._can_confirm_secret_read(): + self._pending_secret_confirmation = None + return "当前入口不支持安全交付敏感设置，未执行读取。" + + if not isinstance(tool, QuerySystemSettingsTool): + self._pending_secret_confirmation = None + return "当前工具不支持敏感设置确认。" + + args_schema = tool.args_schema + if args_schema is None: + self._pending_secret_confirmation = None + return "敏感设置读取参数无法校验，未执行读取。" + try: + validated_arguments = args_schema.model_validate(arguments).model_dump() + except Exception: + self._pending_secret_confirmation = None + return "敏感设置读取参数无效，未执行读取。" + if validated_arguments.get("show_secrets") is not True: + self._pending_secret_confirmation = None + return "当前操作不需要敏感设置确认。" + + permission_result = None + if not await self._is_system_admin_context(): + permission_result = await tool._check_permission() + if permission_result: + self._pending_secret_confirmation = None + return permission_result + + self._pending_secret_confirmation = _PendingSecretConfirmation( + tool=tool, + arguments=validated_arguments, + created_at=datetime.now(), + user_id=str(self.user_id or ""), + channel=str(self.channel or ""), + source=str(self.source or ""), + original_chat_id=str(self.original_chat_id or ""), + ) + target = validated_arguments.get("setting_key") or ( + validated_arguments.get("group") or "all" + ) + confirmation_message = ( + f"即将读取系统设置 {target} 的未脱敏值。" + "结果会直接发送给您，不会交给模型或写入对话历史。" + "请在 5 分钟内回复“确认”继续，或回复“取消”放弃。" + ) + if self.channel == MessageChannel.WebAgent.value: + self._emit_output(confirmation_message) + else: + await self.send_agent_message(confirmation_message) + self._tool_context["user_reply_sent"] = True + return confirmation_message + + async def _deliver_protected_output(self, content: str) -> None: + """绕过模型与会话历史，把敏感结果直接交付给当前用户。""" + if callable(self.protected_output_callback): + try: + self.protected_output_callback(content) + except Exception as e: + logger.error(f"受保护输出回调失败: {e}") + return + + await AgentChain().async_post_message( + Notification( + channel=self.channel, + source=self.source, + mtype=NotificationType.Agent, + userid=self.user_id, + username=self.username, + original_message_id=self.original_message_id, + original_chat_id=self.original_chat_id, + text=content, + save_history=False, + ) + ) + + async def _handle_secret_confirmation_control( + self, + message: str, + images: Optional[List[str]], + files: Optional[List[dict]], + has_audio_input: bool, + ) -> Optional[str]: + """在进入模型前消费当前会话的确认或取消文本。""" + command = str(message or "").strip() + if command not in {SECRET_CONFIRM_TEXT, SECRET_CANCEL_TEXT}: + return None + if images or files or has_audio_input: + return None + + pending = self._pending_secret_confirmation + if not pending: + return None + if ( + pending.user_id != str(self.user_id or "") + or pending.channel != str(self.channel or "") + or pending.source != str(self.source or "") + or pending.original_chat_id != str(self.original_chat_id or "") + ): + return None + if datetime.now() - pending.created_at > SECRET_CONFIRMATION_TTL: + self._pending_secret_confirmation = None + message_text = "敏感设置读取确认已过期，请重新发起。" + await self._deliver_protected_output(message_text) + return message_text + self._pending_secret_confirmation = None + if command == SECRET_CANCEL_TEXT: + message_text = "已取消敏感设置读取。" + await self._deliver_protected_output(message_text) + return message_text + + permission_result = await pending.tool._check_permission() + if permission_result: + await self._deliver_protected_output(permission_result) + return permission_result + + if not self._can_confirm_secret_read(): + return "当前入口不支持安全交付敏感设置，未执行读取。" + + result = await pending.tool._run_confirmed(**pending.arguments) + await self._deliver_protected_output(result) + return "敏感设置确认已处理。" + def _build_policy_context(self) -> ToolPolicyContext: """根据宿主入口建立模型参数无法伪造的策略上下文。""" if not self.has_message_context: @@ -1304,6 +1482,7 @@ class MoviePilotAgent: "reply_mode": None, "should_dispatch_reply": False, "is_admin": bool(self._tool_context.get("is_admin")), + "require_secret_confirmation": True, }, allow_message_tools=False, ) @@ -1505,6 +1684,7 @@ class MoviePilotAgent: AgentPolicyMiddleware( context=policy_context, catalog=tool_catalog, + tools=tools, ), # Skills skills_middleware, @@ -1595,6 +1775,15 @@ class MoviePilotAgent: ) self._streamed_output = "" + confirmation_result = await self._handle_secret_confirmation_control( + message=message, + images=images, + files=files, + has_audio_input=has_audio_input, + ) + if confirmation_result is not None: + return confirmation_result + # 获取历史消息 messages = list(memory_manager.get_agent_messages( session_id=self.session_id, user_id=self.user_id @@ -1950,6 +2139,8 @@ class MoviePilotAgent: """ 清理智能体资源 """ + self._pending_secret_confirmation = None + self.protected_output_callback = None self._compiled_agent_bundle = None logger.info(f"MoviePilot智能体已清理: session_id={self.session_id}") @@ -1975,6 +2166,7 @@ class _MessageTask: reply_mode: ReplyMode = ReplyMode.DISPATCH allow_message_tools: bool = True output_callback: Optional[Callable[[str], None]] = None + protected_output_callback: Optional[Callable[[str], None]] = None notification_callback: Optional[Callable[[Any], None]] = None agent_factory: Optional[Callable[..., MoviePilotAgent]] = None completion_future: Optional[asyncio.Future] = None @@ -2029,6 +2221,29 @@ class AgentManager: ) return status + def matches_secret_confirmation( + self, + session_id: str, + user_id: str, + channel: Optional[str] = None, + source: Optional[str] = None, + original_chat_id: Optional[str] = None, + ) -> bool: + """判断指定用户是否可继续当前会话的敏感设置确认。""" + agent = self.active_agents.get(session_id) + pending = agent._pending_secret_confirmation if agent else None + return bool( + agent + and pending + and str(agent.user_id) == str(user_id) + and (channel is None or pending.channel == str(channel)) + and (source is None or pending.source == str(source)) + and ( + original_chat_id is None + or pending.original_chat_id == str(original_chat_id) + ) + ) + async def initialize(self): """ 初始化管理器 @@ -2130,6 +2345,7 @@ class AgentManager: reply_mode: ReplyMode = ReplyMode.DISPATCH, allow_message_tools: bool = True, output_callback: Optional[Callable[[str], None]] = None, + protected_output_callback: Optional[Callable[[str], None]] = None, notification_callback: Optional[Callable[[Any], None]] = None, agent_factory: Optional[Callable[..., MoviePilotAgent]] = None, wait_for_completion: bool = False, @@ -2156,6 +2372,7 @@ class AgentManager: reply_mode=reply_mode, allow_message_tools=allow_message_tools, output_callback=output_callback, + protected_output_callback=protected_output_callback, notification_callback=notification_callback, agent_factory=agent_factory, completion_future=completion_future, @@ -2291,6 +2508,7 @@ class AgentManager: "replay_mode": task.reply_mode, "allow_message_tools": task.allow_message_tools, "output_callback": task.output_callback, + "protected_output_callback": task.protected_output_callback, } if task.notification_callback is not None and task.agent_factory: agent_kwargs["notification_callback"] = task.notification_callback @@ -2312,6 +2530,7 @@ class AgentManager: agent.set_output_callback(task.output_callback) else: agent.output_callback = task.output_callback + agent.set_protected_output_callback(task.protected_output_callback) if task.notification_callback is not None and hasattr(agent, "set_notification_callback"): agent.set_notification_callback(task.notification_callback) diff --git a/app/agent/middleware/policy.py b/app/agent/middleware/policy.py index 64dcdaa6d..f66d2af09 100644 --- a/app/agent/middleware/policy.py +++ b/app/agent/middleware/policy.py @@ -3,15 +3,18 @@ from collections.abc import Awaitable, Callable from typing import Any -from langchain.agents.middleware import AgentMiddleware, ToolCallRequest +from langchain.agents.middleware import AgentMiddleware, ToolCallRequest, hook_config +from langchain_core.messages import AIMessage, ToolMessage from app.agent.policy import ( DEFAULT_TOOL_POLICY_ORCHESTRATOR, AgentToolPolicyOrchestrator, ToolPolicyContext, + ToolOrigin, call_policy_hook, ) from app.agent.tools.catalog import ToolCatalogSnapshot +from app.agent.tools.impl.query_system_settings import QuerySystemSettingsTool class AgentPolicyMiddleware(AgentMiddleware): @@ -27,11 +30,68 @@ class AgentPolicyMiddleware(AgentMiddleware): context: ToolPolicyContext, orchestrator: AgentToolPolicyOrchestrator = DEFAULT_TOOL_POLICY_ORCHESTRATOR, catalog: ToolCatalogSnapshot | None = None, + tools: list[Any] | None = None, ) -> None: """绑定宿主可信上下文和共享策略编排器。""" self.context = context self.orchestrator = orchestrator self.catalog = catalog + self._tools = { + tool.name: tool + for tool in (tools or []) + if getattr(tool, "name", None) + } + + @hook_config(can_jump_to=["end"]) + async def aafter_model(self, state: dict[str, Any], runtime: Any) -> Any: + """在 ToolNode 前暂停需要用户确认的敏感设置读取。""" + messages = state.get("messages") or [] + if not messages or not isinstance(messages[-1], AIMessage): + return None + + tool_calls = messages[-1].tool_calls or [] + sensitive_call = None + sensitive_tool = None + for tool_call in tool_calls: + arguments = tool_call.get("args") + tool = self._tools.get(tool_call.get("name")) + if ( + isinstance(tool, QuerySystemSettingsTool) + and isinstance(arguments, dict) + and arguments.get("show_secrets") is True + ): + sensitive_call = tool_call + sensitive_tool = tool + break + if sensitive_call is None or sensitive_tool is None: + return None + + confirmation_handler = ( + self.context.agent_context.get("secret_confirmation_handler") + if self.context.origin is ToolOrigin.AGENT_INTERACTIVE + else None + ) + if not callable(confirmation_handler): + confirmation_message = "当前入口不支持敏感设置确认，未执行任何工具。" + else: + confirmation_message = await confirmation_handler( + sensitive_tool, + sensitive_call.get("args") or {}, + ) + + paused_messages = [ + ToolMessage( + content=( + "本轮工具调用已暂停，未执行任何操作；" + "请等待用户确认或取消敏感设置读取。" + ), + tool_call_id=str(tool_call.get("id") or ""), + name=str(tool_call.get("name") or "unknown"), + ) + for tool_call in tool_calls + ] + paused_messages.append(AIMessage(content=confirmation_message)) + return {"messages": paused_messages, "jump_to": "end"} async def awrap_tool_call( self, diff --git a/app/agent/tools/impl/query_system_settings.py b/app/agent/tools/impl/query_system_settings.py index 43ac2a589..19e2609ff 100644 --- a/app/agent/tools/impl/query_system_settings.py +++ b/app/agent/tools/impl/query_system_settings.py @@ -3,7 +3,7 @@ import json from typing import Optional, Type -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, PrivateAttr from app.agent.tools.base import MoviePilotTool from app.agent.tools.tags import ToolTag @@ -64,6 +64,8 @@ class QuerySystemSettingsInput(BaseModel): class QuerySystemSettingsTool(MoviePilotTool): + """查询系统设置，并隔离 Agent 确认后的密钥读取入口。""" + name: str = "query_system_settings" tags: list[str] = [ ToolTag.Read, @@ -78,6 +80,15 @@ class QuerySystemSettingsTool(MoviePilotTool): ) require_admin: bool = True args_schema: Type[BaseModel] = QuerySystemSettingsInput + _secret_read_confirmed: bool = PrivateAttr(default=False) + + async def _run_confirmed(self, **kwargs) -> str: + """仅供宿主在消费有效确认后执行一次未脱敏读取。""" + self._secret_read_confirmed = True + try: + return await self.run(**kwargs) + finally: + self._secret_read_confirmed = False def get_tool_message(self, **kwargs) -> Optional[str]: """根据查询参数生成友好的提示消息。""" @@ -136,12 +147,25 @@ class QuerySystemSettingsTool(MoviePilotTool): show_secrets: Optional[bool] = False, **kwargs, ) -> str: + """查询系统设置，并在 Agent 场景中阻止未经确认的密钥明文读取。""" logger.info( f"执行工具: {self.name}, setting_key={setting_key}, " f"group={group}, keyword={keyword}" ) try: + if ( + show_secrets is True + and self._agent_context.get("require_secret_confirmation") + and not self._secret_read_confirmed + ): + return json.dumps( + { + "success": False, + "message": "读取敏感设置前需要用户确认，本次未执行。", + }, + ensure_ascii=False, + ) if setting_key: spec = resolve_setting_spec(setting_key) if not spec: diff --git a/app/api/endpoints/agent.py b/app/api/endpoints/agent.py index b5e038f60..364ea6f7a 100644 --- a/app/api/endpoints/agent.py +++ b/app/api/endpoints/agent.py @@ -608,6 +608,11 @@ def _build_web_agent_sse( :param locale: 当前请求语言 :return: 符合 SSE 格式的字符串 """ + if event_type == "interaction-protected": + return ( + "event: interaction-protected\n" + f"data: {json.dumps(data or {}, ensure_ascii=False)}\n\n" + ) payload = {"type": event_type, **(data or {})} message = payload.get("message") if event_type == "error" and isinstance(message, str): @@ -1936,6 +1941,37 @@ async def web_agent_stream( prompt = payload.text.strip() locale = LocaleHelper.get_locale_from_request(request) display_prompt = (payload.display_text or payload.text).strip() + session_id = _build_web_agent_session_id(current_user, payload.session_id) + is_secret_confirmation_candidate = ( + prompt in {"确认", "取消"} + and not payload.images + and not payload.audio_refs + and not payload.files + ) + is_secret_confirmation_control = ( + is_secret_confirmation_candidate + and agent_manager.matches_secret_confirmation( + session_id, + str(current_user.id), + channel=MessageChannel.WebAgent.value, + source=WEB_AGENT_SOURCE, + original_chat_id=str(payload.original_chat_id or ""), + ) + ) + protected_transport_supported = ( + getattr(request, "headers", {}).get("X-MoviePilot-Agent-Interaction") == "1" + ) + if is_secret_confirmation_control and not protected_transport_supported: + return StreamingResponse( + iter([ + _build_web_agent_sse( + "error", + {"message": "当前客户端不支持安全交付敏感设置，未执行操作。"}, + locale=locale, + ) + ]), + media_type="text/event-stream", + ) is_traditional_message = ( _is_web_agent_traditional_message(prompt) or _has_web_agent_traditional_interaction(str(current_user.id)) @@ -1966,7 +2002,6 @@ async def web_agent_stream( media_type="text/event-stream", ) - session_id = _build_web_agent_session_id(current_user, payload.session_id) user_attachments = _build_web_agent_input_attachments( images=payload.images or [], files=[ @@ -2100,7 +2135,6 @@ async def web_agent_stream( media_type="text/event-stream", ) - session_id = _build_web_agent_session_id(current_user, payload.session_id) MessageChain().bind_user_session(str(current_user.id), session_id) event_publisher = _WebAgentEventPublisher() user_attachments = _build_web_agent_input_attachments( @@ -2112,7 +2146,7 @@ async def web_agent_stream( audio_refs=payload.audio_refs or [], ) display_messages = [] - if payload.echo_user: + if payload.echo_user and not is_secret_confirmation_control: user_display_message = MoviePilotAgent.build_display_message( role="user", content=display_prompt or prompt, @@ -2143,6 +2177,15 @@ async def web_agent_stream( _apply_web_agent_display_event(item, assistant_display_message) event_publisher.publish(item) + def protected_output_callback(content: str) -> None: + """将敏感文本封装为不进入普通展示快照的命名 SSE 事件。""" + event_publisher.publish( + { + "type": "interaction-protected", + "content": content, + } + ) + async def event_generator() -> AsyncIterator[str]: """ 生成前端 Agent SSE 事件。 @@ -2172,6 +2215,11 @@ async def web_agent_stream( reply_mode=ReplyMode.CAPTURE_ONLY, allow_message_tools=True, output_callback=output_callback, + protected_output_callback=( + protected_output_callback + if protected_transport_supported + else None + ), notification_callback=notification_callback, agent_factory=_WebAgentMoviePilotAgent, wait_for_completion=True, @@ -2189,13 +2237,14 @@ async def web_agent_stream( _apply_web_agent_display_event(done_event, assistant_display_message) # 终态先进入 SSE 队列，避免展示快照落库延迟前端结束动画。 event_publisher.publish(done_event) - await run_in_threadpool( - _save_web_agent_display_snapshot, - session_id=session_id, - current_user=current_user, - messages=display_messages, - client_session_id=payload.session_id or session_id, - ) + if not is_secret_confirmation_control: + await run_in_threadpool( + _save_web_agent_display_snapshot, + session_id=session_id, + current_user=current_user, + messages=display_messages, + client_session_id=payload.session_id or session_id, + ) task = asyncio.create_task(run_agent()) _WEB_AGENT_BACKGROUND_TASKS.add(task) @@ -2241,8 +2290,7 @@ async def web_agent_stream( except asyncio.CancelledError: pass await event_publisher.aclose() - # 客户端退到后台导致 SSE 断开时，保留后台 Agent 继续执行；完成后会保存展示快照， - # 前端恢复可见时可通过会话详情接口拉取最终状态。 + # 客户端断线后保留 Agent 继续执行；发布器关闭后不再接受受保护结果。 return StreamingResponse( event_generator(), @@ -2251,5 +2299,10 @@ async def web_agent_stream( "Cache-Control": "no-cache, no-transform", "Connection": "keep-alive", "X-Accel-Buffering": "no", + **( + {"X-MoviePilot-Agent-Control": "secret-confirmation"} + if is_secret_confirmation_control + else {} + ), }, ) diff --git a/app/chain/message.py b/app/chain/message.py index 5f5b4f763..3627262c4 100644 --- a/app/chain/message.py +++ b/app/chain/message.py @@ -209,6 +209,21 @@ class MessageChain(ChainBase): ) return + if self._handle_secret_confirmation_control( + channel=channel, + source=source, + userid=userid, + username=username, + text=text, + original_message_id=original_message_id, + original_chat_id=original_chat_id, + images=images, + audio_refs=audio_refs, + files=files, + has_audio_input=has_audio_input, + ): + return + if self._handle_plugin_input_interaction( channel=channel, source=source, @@ -277,6 +292,54 @@ class MessageChain(ChainBase): original_chat_id=original_chat_id, ) + def _handle_secret_confirmation_control( + self, + channel: MessageChannel, + source: str, + userid: Union[str, int], + username: str, + text: Optional[str], + original_message_id: Optional[Union[str, int]] = None, + original_chat_id: Optional[str] = None, + images: Optional[List[CommingMessage.MessageImage]] = None, + audio_refs: Optional[List[str]] = None, + files: Optional[List[CommingMessage.MessageAttachment]] = None, + has_audio_input: bool = False, + ) -> bool: + """将 TG/飞书中的确认控制文本交回所属 Agent 会话。""" + if channel not in {MessageChannel.Telegram, MessageChannel.Feishu}: + return False + if str(text or "").strip() not in {"确认", "取消"}: + return False + if images or audio_refs or files or has_audio_input: + return False + + session_info = self._user_sessions.get(userid) + if not session_info: + return False + session_id, _ = session_info + if not agent_manager.matches_secret_confirmation( + session_id, + str(userid), + channel=channel.value, + source=source, + original_chat_id=original_chat_id, + ): + return False + return self._handle_ai_message( + text=str(text).strip(), + channel=channel, + source=source, + userid=userid, + username=username, + original_message_id=original_message_id, + original_chat_id=original_chat_id, + images=images, + files=files, + session_id=session_id, + has_audio_input=has_audio_input, + ) + def _handle_message_core( self, channel: MessageChannel, diff --git a/tests/test_agent_interaction.py b/tests/test_agent_interaction.py index 03aa8148e..d3cb643f6 100644 --- a/tests/test_agent_interaction.py +++ b/tests/test_agent_interaction.py @@ -1,5 +1,6 @@ import asyncio import unittest +from datetime import datetime from unittest.mock import AsyncMock, Mock, patch from app.agent.prompt import prompt_manager @@ -253,3 +254,47 @@ class TestAgentInteraction(unittest.TestCase): ) handle_ai_message.assert_called_once() + + def test_secret_confirmation_preempts_plugin_interaction_on_message_channels(self): + """TG/飞书确认必须回到已有 Agent 会话，不被其它输入会话消费。""" + chain = MessageChain() + MessageChain._user_sessions["10001"] = ("session-secret", datetime.now()) + + try: + for channel in (MessageChannel.Telegram, MessageChannel.Feishu): + with patch( + "app.chain.message.agent_manager.matches_secret_confirmation", + return_value=True, + ), patch.object( + chain, + "_handle_ai_message", + return_value=True, + ) as handle_ai_message, patch.object( + chain, + "_handle_plugin_input_interaction", + ) as handle_plugin_interaction, patch.object( + chain, + "_mark_message_processing_started", + ) as mark_processing_started: + chain.handle_message( + channel=channel, + source=f"{channel.value}-test", + userid="10001", + username="tester", + text="确认", + original_message_id="message-1", + original_chat_id="chat-1", + images=None, + audio_refs=None, + files=None, + ) + + handle_ai_message.assert_called_once() + handle_plugin_interaction.assert_not_called() + mark_processing_started.assert_not_called() + self.assertEqual( + handle_ai_message.call_args.kwargs["session_id"], + "session-secret", + ) + finally: + MessageChain._user_sessions.clear() diff --git a/tests/test_agent_secret_confirmation.py b/tests/test_agent_secret_confirmation.py new file mode 100644 index 000000000..945fd0588 --- /dev/null +++ b/tests/test_agent_secret_confirmation.py @@ -0,0 +1,395 @@ +"""Agent 敏感系统设置读取的宿主确认测试。""" + +import asyncio +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, patch + +from langchain.agents import create_agent +from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel +from langchain_core.messages import AIMessage, HumanMessage + +from app.agent import MoviePilotAgent, ReplyMode, agent_manager +from app.agent.middleware.policy import AgentPolicyMiddleware +from app.agent.policy import AuthSource, PrincipalType, ToolOrigin, ToolPolicyContext +from app.agent.tools.impl.query_system_settings import QuerySystemSettingsTool +from app.schemas.types import MessageChannel + + +class _ToolCallingFakeModel(FakeMessagesListChatModel): + """允许 LangChain 为固定响应假模型绑定本地工具。""" + + def bind_tools(self, tools, **kwargs): + """保留固定响应行为，仅声明测试模型支持工具绑定。""" + return self + + +def _policy_context(agent_context: dict) -> ToolPolicyContext: + """构造可注入确认处理器的交互式策略上下文。""" + return ToolPolicyContext( + session_id="session-secret", + user_id="user-secret", + origin=ToolOrigin.AGENT_INTERACTIVE, + principal_type=PrincipalType.HUMAN, + auth_source=AuthSource.CHANNEL, + agent_context=agent_context, + channel=MessageChannel.Telegram.value, + source="telegram-test", + ) + + +def test_after_model_pauses_secret_setting_read_before_tool_node() -> None: + """首次敏感读取必须结束当前图执行，并闭合整批 tool call。""" + tool = QuerySystemSettingsTool( + session_id="session-secret", + user_id="user-secret", + ) + confirmation_handler = AsyncMock(return_value="请回复“确认”继续，回复“取消”放弃。") + middleware = AgentPolicyMiddleware( + context=_policy_context( + {"secret_confirmation_handler": confirmation_handler} + ), + tools=[tool], + ) + state = { + "messages": [ + AIMessage( + content="", + tool_calls=[ + { + "name": tool.name, + "args": { + "setting_key": "TMDB_API_KEY", + "show_secrets": True, + }, + "id": "secret-call", + }, + { + "name": "query_schedulers", + "args": {}, + "id": "ordinary-call", + }, + ], + ) + ] + } + + result = asyncio.run(middleware.aafter_model(state, runtime=None)) + + assert result["jump_to"] == "end" + assert isinstance(result["messages"][-1], AIMessage) + assert "确认" in result["messages"][-1].content + confirmation_handler.assert_awaited_once() + assert confirmation_handler.await_args.args[0] is tool + assert confirmation_handler.await_args.args[1]["show_secrets"] is True + + +def test_after_model_keeps_redacted_setting_read_on_normal_tool_path() -> None: + """普通设置读取不得增加确认步骤。""" + tool = QuerySystemSettingsTool( + session_id="session-secret", + user_id="user-secret", + ) + confirmation_handler = AsyncMock() + middleware = AgentPolicyMiddleware( + context=_policy_context( + {"secret_confirmation_handler": confirmation_handler} + ), + tools=[tool], + ) + state = { + "messages": [ + AIMessage( + content="", + tool_calls=[ + { + "name": tool.name, + "args": { + "setting_key": "TMDB_API_KEY", + "show_secrets": False, + }, + "id": "redacted-call", + } + ], + ) + ] + } + + assert asyncio.run(middleware.aafter_model(state, runtime=None)) is None + confirmation_handler.assert_not_awaited() + + +def test_real_agent_graph_stops_before_secret_tool_execution() -> None: + """真实 Agent 图必须在 ToolNode 前结束，不得执行敏感读取。""" + tool = QuerySystemSettingsTool( + session_id="session-secret", + user_id="user-secret", + ) + confirmation_handler = AsyncMock(return_value="请回复“确认”继续。") + context = _policy_context( + {"secret_confirmation_handler": confirmation_handler} + ) + model = _ToolCallingFakeModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "name": tool.name, + "args": { + "setting_key": "TMDB_API_KEY", + "show_secrets": True, + }, + "id": "secret-call", + } + ], + ) + ] + ) + graph = create_agent( + model=model, + tools=[tool], + middleware=[AgentPolicyMiddleware(context=context, tools=[tool])], + ) + + with patch.object( + QuerySystemSettingsTool, + "_load_setting_value", + side_effect=AssertionError("敏感工具不应执行"), + ) as load_value: + result = asyncio.run( + graph.ainvoke({"messages": [HumanMessage(content="读取密钥")]}) + ) + + load_value.assert_not_called() + confirmation_handler.assert_awaited_once() + assert isinstance(result["messages"][-1], AIMessage) + assert "确认" in result["messages"][-1].content + + +def test_confirm_executes_once_without_model_or_history() -> None: + """确认应直接执行冻结参数，结果不经过模型和 Agent 历史。""" + secret_marker = "confirmed-secret-marker" + protected_output = [] + ordinary_output = [] + agent = MoviePilotAgent( + session_id="session-secret", + user_id="1", + channel=MessageChannel.WebAgent.value, + source="web-agent", + username="admin", + replay_mode=ReplyMode.CAPTURE_ONLY, + output_callback=ordinary_output.append, + protected_output_callback=protected_output.append, + ) + tool = QuerySystemSettingsTool(session_id="session-secret", user_id="1") + tool.set_message_attr( + channel=MessageChannel.WebAgent.value, + source="web-agent", + username="admin", + ) + tool.set_agent_context(agent._tool_context) + + async def scenario() -> tuple[str, str]: + prompt = await agent._register_secret_confirmation( + tool, + {"setting_key": "TMDB_API_KEY", "show_secrets": True}, + ) + result = await agent.process("确认") + return prompt, result + + with ( + patch.object(agent, "_is_system_admin_context", new=AsyncMock(return_value=True)), + patch.object(agent, "_execute_agent", new=AsyncMock()) as execute_agent, + patch.object(agent, "_save_display_history_messages") as save_display, + patch("app.agent.memory_manager.save_agent_messages") as save_messages, + patch.object( + QuerySystemSettingsTool, + "_load_setting_value", + return_value=secret_marker, + ) as load_value, + ): + prompt, result = asyncio.run(scenario()) + + assert "确认" in prompt + assert ordinary_output == [prompt] + assert result == "敏感设置确认已处理。" + assert len(protected_output) == 1 + assert secret_marker in protected_output[0] + load_value.assert_called_once() + execute_agent.assert_not_awaited() + save_display.assert_not_called() + save_messages.assert_not_called() + assert agent.has_pending_secret_confirmation() is False + + +def test_cancel_clears_pending_without_executing_tool() -> None: + """取消只消费当前 pending，不读取任何设置值。""" + protected_output = [] + agent = MoviePilotAgent( + session_id="session-secret", + user_id="1", + channel=MessageChannel.WebAgent.value, + source="web-agent", + username="admin", + replay_mode=ReplyMode.CAPTURE_ONLY, + protected_output_callback=protected_output.append, + ) + tool = QuerySystemSettingsTool(session_id="session-secret", user_id="1") + + async def scenario() -> str: + await agent._register_secret_confirmation( + tool, + {"setting_key": "TMDB_API_KEY", "show_secrets": True}, + ) + return await agent.process("取消") + + with patch.object( + agent, + "_is_system_admin_context", + new=AsyncMock(return_value=True), + ), patch.object( + QuerySystemSettingsTool, + "_load_setting_value", + ) as load_value: + result = asyncio.run(scenario()) + + assert result == "已取消敏感设置读取。" + assert protected_output == ["已取消敏感设置读取。"] + load_value.assert_not_called() + assert agent.has_pending_secret_confirmation() is False + + +def test_expired_confirmation_reaches_agent_expiry_receipt() -> None: + """入口不得提前清除过期 pending，否则确认文本会被误送给模型。""" + protected_output = [] + agent = MoviePilotAgent( + session_id="session-expired-secret", + user_id="1", + channel=MessageChannel.WebAgent.value, + source="web-agent", + username="admin", + replay_mode=ReplyMode.CAPTURE_ONLY, + protected_output_callback=protected_output.append, + ) + tool = QuerySystemSettingsTool(session_id=agent.session_id, user_id="1") + + async def scenario() -> str: + await agent._register_secret_confirmation( + tool, + {"setting_key": "TMDB_API_KEY", "show_secrets": True}, + ) + agent._pending_secret_confirmation.created_at = ( + datetime.now() - timedelta(minutes=6) + ) + agent_manager.active_agents[agent.session_id] = agent + try: + assert agent_manager.matches_secret_confirmation( + agent.session_id, + "1", + channel=MessageChannel.WebAgent.value, + source="web-agent", + original_chat_id="", + ) + return await agent.process("确认") + finally: + agent_manager.active_agents.pop(agent.session_id, None) + + with patch.object( + agent, + "_is_system_admin_context", + new=AsyncMock(return_value=True), + ), patch.object(agent, "_execute_agent", new=AsyncMock()) as execute_agent: + result = asyncio.run(scenario()) + + assert result == "敏感设置读取确认已过期，请重新发起。" + assert protected_output == [result] + execute_agent.assert_not_awaited() + + +def test_background_agent_refuses_secret_read_without_pending() -> None: + """后台 Agent 没有用户确认通道时必须直接拒绝明文读取。""" + agent = MoviePilotAgent(session_id="background-secret", user_id="system") + tool = QuerySystemSettingsTool(session_id="background-secret", user_id="system") + + async def scenario() -> str: + context = await agent._build_tool_context(should_dispatch_reply=False) + tool.set_agent_context(context) + return await tool.run(setting_key="TMDB_API_KEY", show_secrets=True) + + with patch.object(QuerySystemSettingsTool, "_load_setting_value") as load_value: + result = asyncio.run(scenario()) + + assert "确认" in result + load_value.assert_not_called() + + +def test_message_channel_receives_confirmation_prompt_once() -> None: + """TG/飞书应由宿主直接发送确认提示，不依赖图状态转成渠道输出。""" + agent = MoviePilotAgent( + session_id="session-secret", + user_id="1", + channel=MessageChannel.Telegram.value, + source="telegram-main", + username="admin", + original_chat_id="chat-1", + ) + tool = QuerySystemSettingsTool(session_id="session-secret", user_id="1") + + async def scenario() -> str: + return await agent._register_secret_confirmation( + tool, + {"setting_key": "TMDB_API_KEY", "show_secrets": True}, + ) + + with patch.object( + agent, + "_is_system_admin_context", + new=AsyncMock(return_value=True), + ), patch.object( + agent, + "send_agent_message", + new=AsyncMock(), + ) as send_message: + prompt = asyncio.run(scenario()) + + send_message.assert_awaited_once_with(prompt) + assert agent._tool_context["user_reply_sent"] is True + + +def test_pending_secret_read_keeps_original_owner_and_action() -> None: + """新请求不得覆盖 pending，错误交付目标也不得消费它。""" + agent = MoviePilotAgent( + session_id="session-secret", + user_id="1", + channel=MessageChannel.Telegram.value, + source="telegram-main", + username="admin", + original_chat_id="chat-1", + ) + tool = QuerySystemSettingsTool(session_id="session-secret", user_id="1") + + async def scenario() -> tuple[str, str, str]: + first = await agent._register_secret_confirmation( + tool, + {"setting_key": "TMDB_API_KEY", "show_secrets": True}, + ) + second = await agent._register_secret_confirmation( + tool, + {"setting_key": "API_TOKEN", "show_secrets": True}, + ) + agent.original_chat_id = "chat-2" + result = await agent.process("确认") + return first, second, result + + with ( + patch.object(agent, "_is_system_admin_context", new=AsyncMock(return_value=True)), + patch.object(agent, "_execute_agent", new=AsyncMock(return_value="普通回复")), + patch.object(QuerySystemSettingsTool, "_load_setting_value") as load_value, + ): + first, second, result = asyncio.run(scenario()) + + assert "TMDB_API_KEY" in first + assert "已有待确认" in second + assert result == "普通回复" + load_value.assert_not_called() + assert agent.has_pending_secret_confirmation() is True diff --git a/tests/test_agent_system_settings_tools.py b/tests/test_agent_system_settings_tools.py index 8e4b2f578..70504632b 100644 --- a/tests/test_agent_system_settings_tools.py +++ b/tests/test_agent_system_settings_tools.py @@ -86,6 +86,55 @@ class TestAgentSystemSettingsTools(unittest.TestCase): self.assertFalse(item["redacted"]) self.assertEqual("site-api-key", item["value"][0]["apikey"]) + def test_agent_tool_refuses_unconfirmed_secret_read(self): + """Agent 宿主要求确认时，工具自身也不得直接返回密钥。""" + tool = QuerySystemSettingsTool(session_id="session-1", user_id="10001") + tool.set_agent_context( + { + "is_admin": True, + "require_secret_confirmation": True, + "secret_read_confirmed": False, + } + ) + + with patch.object( + QuerySystemSettingsTool, + "_load_setting_value", + return_value="must-not-load", + ) as load_value: + result = asyncio.run( + tool.run(setting_key="TMDB_API_KEY", show_secrets=True) + ) + + payload = json.loads(result) + self.assertFalse(payload["success"]) + self.assertIn("确认", payload["message"]) + load_value.assert_not_called() + + def test_agent_tool_ignores_forged_confirmation_context(self): + """模型可影响的共享上下文不得伪造宿主确认。""" + tool = QuerySystemSettingsTool(session_id="session-1", user_id="10001") + tool.set_agent_context( + { + "is_admin": True, + "require_secret_confirmation": True, + "secret_read_confirmed": True, + } + ) + + with patch.object( + QuerySystemSettingsTool, + "_load_setting_value", + return_value="must-not-load", + ) as load_value: + result = asyncio.run( + tool.run(setting_key="TMDB_API_KEY", show_secrets=True) + ) + + payload = json.loads(result) + self.assertFalse(payload["success"]) + load_value.assert_not_called() + def test_query_system_settings_group_defaults_to_summary_for_multiple_items(self): tool = QuerySystemSettingsTool(session_id="session-1", user_id="10001") diff --git a/tests/test_web_agent_stream.py b/tests/test_web_agent_stream.py index fca4fac6d..623a38127 100644 --- a/tests/test_web_agent_stream.py +++ b/tests/test_web_agent_stream.py @@ -1,4 +1,5 @@ import asyncio +import json import time from queue import Queue from threading import Event as ThreadEvent @@ -31,6 +32,7 @@ from app.api.endpoints.agent import ( ) from app.core.event import Event from app.db.agentchat_oper import AgentChatOper +from app.db.models.agentchat import AgentChat from app.helper.agent import build_web_agent_message_update_event from app.helper.interaction import AgentInteractionOption, agent_interaction_manager, skills_interaction_manager from app.chain.message import MessageChain @@ -668,6 +670,10 @@ def test_web_agent_stream_binds_session_to_agent_manager(): """更新当前 SSE 输出回调。""" self.output_callback = output_callback + def set_protected_output_callback(self, protected_output_callback): + """更新当前 SSE 受保护输出回调。""" + self.protected_output_callback = protected_output_callback + def set_notification_callback(self, notification_callback): """更新当前 SSE 通知回调。""" self.notification_callback = notification_callback @@ -715,6 +721,294 @@ def test_web_agent_stream_binds_session_to_agent_manager(): worker.cancel() +def test_web_agent_stream_emits_secret_result_only_as_protected_event(): + """敏感结果只能进入命名 protected SSE，不能进入普通快照。""" + secret_marker = "WEB_SECRET_MARKER **literal** " + payload = schemas.AgentWebChatRequest( + text="确认", + session_id="browser-secret", + echo_user=True, + ) + request = SimpleNamespace( + headers={"X-MoviePilot-Agent-Interaction": "1"}, + is_disconnected=AsyncMock(return_value=False), + ) + user = SimpleNamespace(id=1, name="admin", is_superuser=True) + + class FakeProtectedAgent: + """直接触发受保护输出的 WebAgent 测试替身。""" + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + self._pending_secret_confirmation = SimpleNamespace( + channel=MessageChannel.WebAgent.value, + source="web-agent", + original_chat_id="", + ) + + def has_pending_secret_confirmation(self): + """模拟当前会话存在有效的敏感读取确认。""" + return self._pending_secret_confirmation is not None + + def set_output_callback(self, output_callback): + self.output_callback = output_callback + + def set_notification_callback(self, notification_callback): + self.notification_callback = notification_callback + + def set_protected_output_callback(self, protected_output_callback): + self.protected_output_callback = protected_output_callback + + async def process(self, _message, **_kwargs): + self.protected_output_callback(secret_marker) + return "敏感设置确认已处理。" + + async def cleanup(self): + return None + + session_id = _build_web_agent_session_id(user, payload.session_id) + existing_messages = [ + {"role": "user", "content": "此前的问题", "status": "done"}, + {"role": "assistant", "content": "此前的回答", "status": "done"}, + ] + existing_chat = AgentChatOper().save_display_messages( + session_id=session_id, + user_id="1", + username="admin", + channel=MessageChannel.WebAgent.value, + source="web-agent", + messages=existing_messages, + client_session_id=payload.session_id, + ) + agent_manager.active_agents[session_id] = FakeProtectedAgent( + session_id=session_id, + user_id="1", + channel=MessageChannel.WebAgent.value, + source="web-agent", + username="admin", + ) + agent_manager._session_queues.pop(session_id, None) + worker = agent_manager._session_workers.pop(session_id, None) + if worker: + worker.cancel() + + async def scenario(): + response = await web_agent_stream(payload, request, user) + return "".join(await _collect_streaming_response(response)) + + try: + with patch("app.api.endpoints.agent.settings.AI_AGENT_ENABLE", True), patch( + "app.api.endpoints.agent._WebAgentMoviePilotAgent", + FakeProtectedAgent, + ), patch( + "app.api.endpoints.agent._save_web_agent_display_snapshot", + ) as save_snapshot: + body = asyncio.run(scenario()) + + assert "event: interaction-protected\n" in body + assert secret_marker in body + save_snapshot.assert_not_called() + preserved_chat = AgentChatOper().get(session_id=session_id, user_id="1") + assert preserved_chat.display_messages == existing_messages + assert preserved_chat.message_count == 2 + assert preserved_chat.preview == "此前的回答" + finally: + agent = agent_manager.active_agents.pop(session_id, None) + if agent: + asyncio.run(agent.cleanup()) + agent_manager._session_queues.pop(session_id, None) + worker = agent_manager._session_workers.pop(session_id, None) + if worker: + worker.cancel() + AgentChat.delete(rid=existing_chat.id) + + +def test_web_agent_cancel_keeps_existing_display_history(): + """取消敏感读取不得覆盖当前会话已有的普通展示历史。""" + payload = schemas.AgentWebChatRequest( + text="取消", + session_id="browser-secret-cancel", + echo_user=True, + ) + request = SimpleNamespace( + headers={"X-MoviePilot-Agent-Interaction": "1"}, + is_disconnected=AsyncMock(return_value=False), + ) + user = SimpleNamespace(id=1, name="admin", is_superuser=True) + session_id = _build_web_agent_session_id(user, payload.session_id) + existing_messages = [ + {"role": "user", "content": "保留的问题", "status": "done"}, + {"role": "assistant", "content": "保留的回答", "status": "done"}, + ] + existing_chat = AgentChatOper().save_display_messages( + session_id=session_id, + user_id="1", + username="admin", + channel=MessageChannel.WebAgent.value, + source="web-agent", + messages=existing_messages, + client_session_id=payload.session_id, + ) + + async def scenario(): + response = await web_agent_stream(payload, request, user) + return "".join(await _collect_streaming_response(response)) + + try: + with patch("app.api.endpoints.agent.settings.AI_AGENT_ENABLE", True), patch.object( + agent_manager, + "matches_secret_confirmation", + return_value=True, + ), patch.object( + agent_manager, + "process_message", + new=AsyncMock(return_value="已取消敏感设置读取。"), + ), patch( + "app.api.endpoints.agent._save_web_agent_display_snapshot", + ) as save_snapshot: + body = asyncio.run(scenario()) + + assert '"type": "done"' in body + save_snapshot.assert_not_called() + preserved_chat = AgentChatOper().get(session_id=session_id, user_id="1") + assert preserved_chat.display_messages == existing_messages + assert preserved_chat.message_count == 2 + assert preserved_chat.preview == "保留的回答" + finally: + AgentChat.delete(rid=existing_chat.id) + + +def test_web_agent_stream_rejects_confirmation_without_protected_capability(): + """旧客户端未声明 protected 能力时不得把确认交给 Agent。""" + payload = schemas.AgentWebChatRequest( + text="确认", + session_id="browser-secret-legacy", + echo_user=False, + ) + request = SimpleNamespace( + headers={}, + is_disconnected=AsyncMock(return_value=False), + ) + user = SimpleNamespace(id=1, name="admin", is_superuser=True) + + async def scenario(): + response = await web_agent_stream(payload, request, user) + return "".join(await _collect_streaming_response(response)) + + with patch("app.api.endpoints.agent.settings.AI_AGENT_ENABLE", True), patch.object( + agent_manager, + "matches_secret_confirmation", + return_value=True, + ), patch.object(agent_manager, "process_message", new=AsyncMock()) as process: + body = asyncio.run(scenario()) + + assert "不支持安全交付" in body + process.assert_not_awaited() + + +def test_web_agent_stream_keeps_confirmation_without_pending_on_normal_path(): + """无待确认操作时，纯文本确认仍是普通 Agent 消息。""" + payload = schemas.AgentWebChatRequest( + text="确认", + session_id="browser-ordinary-confirmation", + echo_user=True, + ) + request = SimpleNamespace(headers={}, is_disconnected=AsyncMock(return_value=False)) + user = SimpleNamespace(id=1, name="admin", is_superuser=True) + + async def scenario(): + response = await web_agent_stream(payload, request, user) + body = "".join(await _collect_streaming_response(response)) + return response, body + + with patch("app.api.endpoints.agent.settings.AI_AGENT_ENABLE", True), patch.object( + agent_manager, + "process_message", + new=AsyncMock(return_value="普通回复"), + ) as process: + response, body = asyncio.run(scenario()) + + assert "不支持安全交付" not in body + assert response.headers.get("X-MoviePilot-Agent-Control") is None + process.assert_awaited_once() + + +def test_web_agent_stream_drops_secret_result_after_disconnect(): + """确认请求断线后不改造通用队列，并拒绝向关闭的连接投递密钥。""" + payload = schemas.AgentWebChatRequest( + text="确认", + session_id="browser-secret-disconnect", + echo_user=True, + ) + + agent_started = asyncio.Event() + release_agent = asyncio.Event() + agent_completed = asyncio.Event() + async def disconnect_after_agent_starts(): + """等待确认进入处理流程后再模拟浏览器断线。""" + await agent_started.wait() + return True + + request = SimpleNamespace( + headers={"X-MoviePilot-Agent-Interaction": "1"}, + is_disconnected=AsyncMock(side_effect=disconnect_after_agent_starts), + ) + user = SimpleNamespace(id=1, name="admin", is_superuser=True) + session_id = _build_web_agent_session_id(user, payload.session_id) + existing_messages = [ + {"role": "user", "content": "断线前的问题", "status": "done"}, + {"role": "assistant", "content": "断线前的回答", "status": "done"}, + ] + existing_chat = AgentChatOper().save_display_messages( + session_id=session_id, + user_id="1", + username="admin", + channel=MessageChannel.WebAgent.value, + source="web-agent", + messages=existing_messages, + client_session_id=payload.session_id, + ) + + async def finish_after_disconnect(**kwargs): + """断线后继续完成只读任务，并尝试向已关闭发布器投递。""" + agent_started.set() + await release_agent.wait() + kwargs["protected_output_callback"]("DISCONNECTED_SECRET_MARKER") + agent_completed.set() + + async def scenario(): + response = await web_agent_stream(payload, request, user) + body = "".join(await _collect_streaming_response(response)) + release_agent.set() + await asyncio.wait_for(agent_completed.wait(), timeout=1) + await asyncio.sleep(0) + return body + + try: + with patch("app.api.endpoints.agent.settings.AI_AGENT_ENABLE", True), patch.object( + agent_manager, + "matches_secret_confirmation", + return_value=True, + ), patch.object( + agent_manager, + "process_message", + new=AsyncMock(side_effect=finish_after_disconnect), + ) as process, patch( + "app.api.endpoints.agent._save_web_agent_display_snapshot", + ) as save_snapshot: + body = asyncio.run(scenario()) + + assert '"type": "start"' in body + assert "cancel_on_waiter_cancel" not in process.await_args.kwargs + save_snapshot.assert_not_called() + preserved_chat = AgentChatOper().get(session_id=session_id, user_id="1") + assert preserved_chat.display_messages == existing_messages + assert preserved_chat.message_count == 2 + assert preserved_chat.preview == "断线前的回答" + finally: + AgentChat.delete(rid=existing_chat.id) + + def test_web_agent_stream_emits_heartbeat_during_idle_tool_wait(): """长时间没有 Agent 事件时应发送 SSE heartbeat 保持连接。""" payload = schemas.AgentWebChatRequest(text="分析系统状态", session_id="browser-heartbeat")