mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-08-14 10:14:36 +08:00
feat(agent): 敏感设置读取增加宿主二次确认 (#6283)
* feat(agent): require confirmation for secret settings * fix(agent): preserve history during secret confirmation
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 {}
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
395
tests/test_agent_secret_confirmation.py
Normal file
395
tests/test_agent_secret_confirmation.py
Normal 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
|
||||
@@ -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")
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user