feat(agent): 敏感设置读取增加宿主二次确认 (#6283)

* feat(agent): require confirmation for secret settings

* fix(agent): preserve history during secret confirmation
This commit is contained in:
InfinityPacer
2026-08-13 08:10:32 +08:00
committed by GitHub
parent 949464d064
commit a6e8ba8b57
9 changed files with 1216 additions and 14 deletions

View File

@@ -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)

View File

@@ -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,

View File

@@ -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:

View File

@@ -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 {}
),
},
)

View File

@@ -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,

View File

@@ -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()

View File

@@ -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

View File

@@ -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")

View File

@@ -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** <img src=x>"
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")