feat: 新增会话通信工具与子会话嵌套管理能力
1. 新增会话通信工具集,支持列出会话、查看会话历史、发送会话消息和查询智能体进度 2. 实现子智能体嵌套深度限制与父子会话关系管理 3. 新增数据库字段与索引支持会话层级存储 4. 添加会话工具单元测试与LITE模式适配
This commit is contained in:
parent
436c09fb90
commit
32b213593e
@ -17,14 +17,17 @@ from langgraph.types import Command
|
||||
from yuxi.agents.context import build_agent_input_context
|
||||
from yuxi.repositories.agent_repository import SUB_AGENT_BACKEND_ID, AgentRepository
|
||||
from yuxi.repositories.agent_run_repository import AgentRunRepository
|
||||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||||
from yuxi.repositories.user_repository import UserRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.storage.postgres.models_business import Agent
|
||||
from yuxi.utils import logger
|
||||
from yuxi.utils.datetime_utils import utc_isoformat
|
||||
from yuxi.utils.subagent_thread_utils import make_child_thread_id
|
||||
|
||||
_CHILD_STATE_INHERIT_KEYS: frozenset[str] = frozenset()
|
||||
_TERMINAL_RUN_STATUSES = {"completed", "failed", "cancelled", "interrupted"}
|
||||
_MAX_SPAWN_DEPTH = 3
|
||||
|
||||
TASK_SYSTEM_PROMPT = """## `task`(子智能体任务工具)
|
||||
|
||||
@ -372,6 +375,19 @@ class YuxiSubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
|
||||
if not backend or agent.backend_id != SUB_AGENT_BACKEND_ID:
|
||||
return f"无法调用子智能体 {subagent_type}:后端配置无效"
|
||||
|
||||
# 嵌套深度限制:查询当前会话的 spawn_depth,超过上限则拒绝
|
||||
parent_spawn_depth: int | None = None
|
||||
try:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
conv_repo = ConversationRepository(db)
|
||||
current_conv = await conv_repo.get_conversation_by_thread_id(parent_thread_id)
|
||||
parent_spawn_depth = current_conv.spawn_depth if current_conv else None
|
||||
current_depth = parent_spawn_depth if parent_spawn_depth is not None else 0
|
||||
if current_depth >= _MAX_SPAWN_DEPTH:
|
||||
return f"子智能体嵌套深度已达上限({_MAX_SPAWN_DEPTH} 层),无法继续派生子智能体"
|
||||
except Exception:
|
||||
logger.warning("Failed to check spawn depth, allowing subagent invocation")
|
||||
|
||||
child_thread_id, continuing = _new_child_thread_id(
|
||||
thread_id,
|
||||
parent_thread_id=parent_thread_id,
|
||||
@ -463,6 +479,20 @@ class YuxiSubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
|
||||
)
|
||||
return _failed_tool_response(exc, runtime.tool_call_id, _with_run_payload(subagent_run, failed_run))
|
||||
completed_run = await self._set_subagent_run_status(run.id, "completed")
|
||||
|
||||
# 更新子会话的父子关系(复用之前查询的 parent_spawn_depth,避免重复开 session 查询父会话)
|
||||
try:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
conv_repo = ConversationRepository(db)
|
||||
child_conv = await conv_repo.get_conversation_by_thread_id(child_thread_id)
|
||||
if child_conv:
|
||||
depth = (parent_spawn_depth + 1) if parent_spawn_depth is not None else 1
|
||||
child_conv.parent_thread_id = parent_thread_id
|
||||
child_conv.spawn_depth = depth
|
||||
await db.commit()
|
||||
except Exception:
|
||||
logger.warning(f"Failed to update parent relationship for child thread {child_thread_id}")
|
||||
|
||||
return _completed_tool_response(
|
||||
result, runtime.tool_call_id, _with_run_payload(subagent_run, completed_run)
|
||||
)
|
||||
|
||||
@ -2,8 +2,22 @@
|
||||
from .install_skill import install_skill
|
||||
from .tools import ask_user_question, present_artifacts
|
||||
|
||||
# 会话通信工具仅在非 LITE 模式下注册
|
||||
from .session_tools import _LITE_MODE as _session_lite_mode
|
||||
|
||||
if not _session_lite_mode:
|
||||
from .session_tools import get_agent_progress, get_session_history, list_sessions, send_to_session
|
||||
|
||||
__all__ = [
|
||||
"ask_user_question",
|
||||
"install_skill",
|
||||
"present_artifacts",
|
||||
]
|
||||
|
||||
if not _session_lite_mode:
|
||||
__all__.extend([
|
||||
"get_agent_progress",
|
||||
"get_session_history",
|
||||
"list_sessions",
|
||||
"send_to_session",
|
||||
])
|
||||
|
||||
318
backend/package/yuxi/agents/toolkits/buildin/session_tools.py
Normal file
318
backend/package/yuxi/agents/toolkits/buildin/session_tools.py
Normal file
@ -0,0 +1,318 @@
|
||||
"""会话通信工具集 — Agent 间会话通信能力"""
|
||||
|
||||
import os
|
||||
|
||||
from langgraph.prebuilt.tool_node import ToolRuntime
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from yuxi.agents.toolkits.registry import tool
|
||||
from yuxi.utils import logger
|
||||
|
||||
_LITE_MODE = os.environ.get("LITE_MODE", "").lower() in ("true", "1")
|
||||
|
||||
# ── 常量 ──────────────────────────────────────────────────
|
||||
|
||||
_HISTORY_CONTENT_MAX_CHARS = 4000
|
||||
_HISTORY_BYTES_MAX = 80 * 1024
|
||||
|
||||
|
||||
# ── 辅助函数 ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def _get_runtime_uid(runtime: ToolRuntime) -> str:
|
||||
"""从 ToolRuntime.context 获取当前用户 uid。"""
|
||||
uid = getattr(runtime.context, "uid", None)
|
||||
if not uid:
|
||||
raise ValueError("当前运行时缺少 uid")
|
||||
return str(uid)
|
||||
|
||||
|
||||
def _get_runtime_thread_id(runtime: ToolRuntime) -> str:
|
||||
"""从 ToolRuntime.context 获取当前线程 ID。
|
||||
|
||||
注意:子智能体的 file_thread_id 继承自父会话,不能用于标识当前会话。
|
||||
会话通信场景应使用 thread_id(当前智能体自身的线程)。
|
||||
"""
|
||||
thread_id = getattr(runtime.context, "thread_id", None)
|
||||
if not thread_id:
|
||||
raise ValueError("当前运行时缺少 thread_id")
|
||||
return str(thread_id)
|
||||
|
||||
|
||||
# ── 权限校验 ──────────────────────────────────────────────
|
||||
|
||||
|
||||
async def _check_session_access(uid: str, thread_id: str) -> None:
|
||||
"""
|
||||
检查用户是否有权访问指定会话(tree 级可见性)。
|
||||
1. 会话的 uid 与当前用户一致 → 允许
|
||||
2. 沿 parent_thread_id 链上溯,任一祖先属于当前用户 → 允许
|
||||
3. 否则 → 拒绝
|
||||
"""
|
||||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
repo = ConversationRepository(db)
|
||||
visited: set[str] = set()
|
||||
current_tid = thread_id
|
||||
|
||||
while current_tid and current_tid not in visited:
|
||||
visited.add(current_tid)
|
||||
conv = await repo.get_conversation_by_thread_id(current_tid)
|
||||
if not conv:
|
||||
raise ValueError(f"会话不存在: {thread_id}")
|
||||
if conv.uid == uid:
|
||||
return
|
||||
parent_tid = conv.parent_thread_id
|
||||
if not parent_tid:
|
||||
break
|
||||
current_tid = parent_tid
|
||||
|
||||
raise PermissionError(f"无权访问会话: {thread_id}")
|
||||
|
||||
|
||||
# ── LITE 模式下不注册会话通信工具 ─────────────────────────
|
||||
|
||||
if _LITE_MODE:
|
||||
logger.info("LITE_MODE enabled, session communication tools not registered")
|
||||
else:
|
||||
|
||||
# ── list_sessions ─────────────────────────────────────────
|
||||
|
||||
class ListSessionsInput(BaseModel):
|
||||
status: str | None = Field(default=None, description="按状态过滤: active / archived")
|
||||
agent_id: str | None = Field(default=None, description="按智能体 ID 过滤")
|
||||
limit: int = Field(default=20, ge=1, le=50, description="返回数量上限")
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["会话"],
|
||||
display_name="列出会话",
|
||||
args_schema=ListSessionsInput,
|
||||
)
|
||||
async def list_sessions(
|
||||
status: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
limit: int = 20,
|
||||
runtime: ToolRuntime = None,
|
||||
) -> str:
|
||||
"""
|
||||
列出当前用户可见的对话会话,包括自己创建的和子智能体的会话。
|
||||
返回每个会话的 thread_id、agent_id、状态、标题、最近更新时间。
|
||||
"""
|
||||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
uid = _get_runtime_uid(runtime)
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
repo = ConversationRepository(db)
|
||||
conversations = await repo.list_conversations(
|
||||
uid=uid,
|
||||
agent_id=agent_id,
|
||||
status=status or "active",
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
if not conversations:
|
||||
return "当前无可见的会话"
|
||||
|
||||
lines = []
|
||||
for conv in conversations:
|
||||
parent_marker = " [子智能体]" if conv.parent_thread_id else ""
|
||||
lines.append(
|
||||
f"- thread_id: {conv.thread_id} | agent: {conv.agent_id} | "
|
||||
f"状态: {conv.status} | 标题: {conv.title or '无'}{parent_marker} | "
|
||||
f"更新: {conv.updated_at}"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
# ── get_session_history ───────────────────────────────────
|
||||
|
||||
class GetSessionHistoryInput(BaseModel):
|
||||
thread_id: str = Field(description="目标会话线程 ID")
|
||||
limit: int = Field(default=10, ge=1, le=50, description="返回消息数量")
|
||||
include_tools: bool = Field(default=False, description="是否包含工具调用消息")
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["会话"],
|
||||
display_name="查看会话历史",
|
||||
args_schema=GetSessionHistoryInput,
|
||||
)
|
||||
async def get_session_history(
|
||||
thread_id: str,
|
||||
limit: int = 10,
|
||||
include_tools: bool = False,
|
||||
runtime: ToolRuntime = None,
|
||||
) -> str:
|
||||
"""
|
||||
获取指定会话的最近消息。
|
||||
仅可查看自己或子智能体的会话历史。
|
||||
超长内容会被截断,工具调用消息默认不包含。
|
||||
"""
|
||||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
uid = _get_runtime_uid(runtime)
|
||||
await _check_session_access(uid, thread_id)
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
repo = ConversationRepository(db)
|
||||
# get_messages_by_thread_id 按 created_at.asc() 排序 + limit 返回最旧消息,
|
||||
# 但我们需要最近消息,所以先降序查再反转。
|
||||
conversation = await repo.get_conversation_by_thread_id(thread_id)
|
||||
if not conversation:
|
||||
return f"会话 {thread_id} 暂无消息"
|
||||
from sqlalchemy import select
|
||||
|
||||
from yuxi.storage.postgres.models_business import Message
|
||||
|
||||
result = await db.execute(
|
||||
select(Message)
|
||||
.where(Message.conversation_id == conversation.id)
|
||||
.order_by(Message.created_at.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
messages = list(reversed(result.scalars().unique().all()))
|
||||
|
||||
if not messages:
|
||||
return f"会话 {thread_id} 暂无消息"
|
||||
|
||||
total_bytes = 0
|
||||
lines = []
|
||||
truncated = False
|
||||
|
||||
for msg in messages:
|
||||
# 过滤工具消息
|
||||
if not include_tools and msg.message_type in ("tool_call", "tool_result"):
|
||||
continue
|
||||
|
||||
content = msg.content or ""
|
||||
# 截断超长内容
|
||||
if len(content) > _HISTORY_CONTENT_MAX_CHARS:
|
||||
content = content[:_HISTORY_CONTENT_MAX_CHARS] + "...[已截断]"
|
||||
truncated = True
|
||||
|
||||
line = f"[{msg.role}] {content}"
|
||||
line_bytes = len(line.encode("utf-8"))
|
||||
if total_bytes + line_bytes > _HISTORY_BYTES_MAX:
|
||||
truncated = True
|
||||
break
|
||||
total_bytes += line_bytes
|
||||
lines.append(line)
|
||||
|
||||
suffix = "\n(部分内容已截断)" if truncated else ""
|
||||
return "\n".join(lines) + suffix
|
||||
|
||||
# ── send_to_session ───────────────────────────────────────
|
||||
|
||||
class SendToSessionInput(BaseModel):
|
||||
thread_id: str = Field(description="目标会话线程 ID")
|
||||
message: str = Field(description="要发送的消息内容")
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["会话"],
|
||||
display_name="发送会话消息",
|
||||
args_schema=SendToSessionInput,
|
||||
)
|
||||
async def send_to_session(
|
||||
thread_id: str,
|
||||
message: str,
|
||||
runtime: ToolRuntime = None,
|
||||
) -> str:
|
||||
"""
|
||||
向指定会话发送消息。
|
||||
消息会以 user 角色注入目标会话的历史记录。
|
||||
如果目标会话正在运行,消息将作为下一轮输入;否则仅写入历史。
|
||||
不能向自己所在的会话发送消息(避免循环)。
|
||||
"""
|
||||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
uid = _get_runtime_uid(runtime)
|
||||
current_thread_id = _get_runtime_thread_id(runtime)
|
||||
|
||||
# 防止向自身会话发送消息
|
||||
if thread_id == current_thread_id:
|
||||
return "不能向当前所在的会话发送消息"
|
||||
|
||||
await _check_session_access(uid, thread_id)
|
||||
|
||||
# 标注消息来源
|
||||
annotated_message = f"[来自会话 {current_thread_id}] {message}"
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
repo = ConversationRepository(db)
|
||||
await repo.add_message_by_thread_id(
|
||||
thread_id=thread_id,
|
||||
role="user",
|
||||
content=annotated_message,
|
||||
message_type="text",
|
||||
extra_metadata={
|
||||
"source_thread_id": current_thread_id,
|
||||
"source_type": "inter_session",
|
||||
},
|
||||
)
|
||||
|
||||
return f"消息已发送到会话 {thread_id}"
|
||||
|
||||
# ── get_agent_progress ────────────────────────────────────
|
||||
|
||||
class GetAgentProgressInput(BaseModel):
|
||||
thread_id: str = Field(description="会话线程 ID")
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["会话"],
|
||||
display_name="查看智能体进度",
|
||||
args_schema=GetAgentProgressInput,
|
||||
)
|
||||
async def get_agent_progress(
|
||||
thread_id: str,
|
||||
runtime: ToolRuntime = None,
|
||||
) -> str:
|
||||
"""
|
||||
查看指定智能体的运行进度。
|
||||
返回运行状态、开始时间、运行类型等信息。
|
||||
"""
|
||||
from sqlalchemy import select
|
||||
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.storage.postgres.models_business import AgentRun
|
||||
|
||||
uid = _get_runtime_uid(runtime)
|
||||
await _check_session_access(uid, thread_id)
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
result = await db.execute(
|
||||
select(AgentRun)
|
||||
.where(
|
||||
AgentRun.thread_id == thread_id,
|
||||
AgentRun.uid == uid,
|
||||
)
|
||||
.order_by(AgentRun.created_at.desc())
|
||||
.limit(1)
|
||||
)
|
||||
run = result.scalar_one_or_none()
|
||||
|
||||
if not run:
|
||||
return f"会话 {thread_id} 暂无运行记录"
|
||||
|
||||
lines = [
|
||||
f"运行 ID: {run.id}",
|
||||
f"状态: {run.status}",
|
||||
f"类型: {run.run_type}",
|
||||
f"智能体: {run.agent_id}",
|
||||
f"创建时间: {run.created_at}",
|
||||
]
|
||||
if run.started_at:
|
||||
lines.append(f"开始时间: {run.started_at}")
|
||||
if run.finished_at:
|
||||
lines.append(f"完成时间: {run.finished_at}")
|
||||
if run.error_message:
|
||||
lines.append(f"错误: {run.error_message}")
|
||||
|
||||
return "\n".join(lines)
|
||||
@ -73,6 +73,20 @@ class ConversationRepository:
|
||||
result = await self.db.execute(select(Conversation).where(Conversation.thread_id == thread_id))
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def get_conversations_by_parent_thread_id(
|
||||
self, parent_thread_id: str, limit: int | None = None
|
||||
) -> list[Conversation]:
|
||||
"""查询指定父会话下的子会话列表。"""
|
||||
query = (
|
||||
select(Conversation)
|
||||
.where(Conversation.parent_thread_id == parent_thread_id)
|
||||
.order_by(Conversation.created_at.desc())
|
||||
)
|
||||
if limit is not None:
|
||||
query = query.limit(limit)
|
||||
result = await self.db.execute(query)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def _get_conversation_by_id(self, conversation_id: int) -> Conversation | None:
|
||||
result = await self.db.execute(select(Conversation).where(Conversation.id == conversation_id))
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
@ -465,6 +465,9 @@ class PostgresManager(metaclass=SingletonMeta):
|
||||
"CREATE INDEX IF NOT EXISTS ix_conversations_is_pinned ON conversations(is_pinned)",
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS ix_model_providers_provider_id ON model_providers(provider_id)",
|
||||
"CREATE INDEX IF NOT EXISTS ix_model_providers_is_enabled ON model_providers(is_enabled)",
|
||||
"ALTER TABLE IF EXISTS conversations ADD COLUMN IF NOT EXISTS parent_thread_id VARCHAR(64)",
|
||||
"ALTER TABLE IF EXISTS conversations ADD COLUMN IF NOT EXISTS spawn_depth INTEGER NOT NULL DEFAULT 0",
|
||||
"CREATE INDEX IF NOT EXISTS idx_conversations_parent ON conversations(parent_thread_id)",
|
||||
]
|
||||
async with self.async_engine.begin() as conn:
|
||||
for stmt in stmts:
|
||||
|
||||
@ -263,6 +263,8 @@ class Conversation(Base):
|
||||
created_at = Column(DateTime, default=utc_now_naive, comment="Creation time")
|
||||
updated_at = Column(DateTime, default=utc_now_naive, onupdate=utc_now_naive, comment="Update time")
|
||||
extra_metadata = Column(JSON, nullable=True, comment="Additional metadata")
|
||||
parent_thread_id = Column(String(64), nullable=True, index=True, comment="Parent thread ID for sub-agent sessions")
|
||||
spawn_depth = Column(Integer, default=0, nullable=False, comment="Nesting depth: 0=main, 1=first-level sub-agent, etc.")
|
||||
|
||||
# Relationships
|
||||
messages = relationship("Message", back_populates="conversation", cascade="all, delete-orphan")
|
||||
@ -283,6 +285,8 @@ class Conversation(Base):
|
||||
"created_at": format_utc_datetime(self.created_at),
|
||||
"updated_at": format_utc_datetime(self.updated_at),
|
||||
"metadata": metadata,
|
||||
"parent_thread_id": self.parent_thread_id,
|
||||
"spawn_depth": self.spawn_depth,
|
||||
}
|
||||
|
||||
|
||||
|
||||
470
backend/test/unit/agents/toolkits/buildin/test_session_tools.py
Normal file
470
backend/test/unit/agents/toolkits/buildin/test_session_tools.py
Normal file
@ -0,0 +1,470 @@
|
||||
"""会话通信工具集单元测试"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from yuxi.agents.toolkits.buildin import session_tools
|
||||
|
||||
|
||||
# ── 辅助函数测试 ──────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_runtime(uid: str = "user-1", thread_id: str = "thread-1", file_thread_id: str | None = None):
|
||||
"""构造模拟的 ToolRuntime 对象。"""
|
||||
context = MagicMock()
|
||||
context.uid = uid
|
||||
context.thread_id = thread_id
|
||||
context.file_thread_id = file_thread_id or thread_id
|
||||
runtime = MagicMock()
|
||||
runtime.context = context
|
||||
return runtime
|
||||
|
||||
|
||||
class TestGetRuntimeUid:
|
||||
def test_returns_uid(self):
|
||||
runtime = _make_runtime(uid="user-123")
|
||||
assert session_tools._get_runtime_uid(runtime) == "user-123"
|
||||
|
||||
def test_raises_when_uid_missing(self):
|
||||
runtime = _make_runtime(uid="")
|
||||
with pytest.raises(ValueError, match="uid"):
|
||||
session_tools._get_runtime_uid(runtime)
|
||||
|
||||
def test_raises_when_uid_none(self):
|
||||
runtime = _make_runtime(uid=None)
|
||||
with pytest.raises(ValueError, match="uid"):
|
||||
session_tools._get_runtime_uid(runtime)
|
||||
|
||||
|
||||
class TestGetRuntimeThreadId:
|
||||
def test_returns_thread_id(self):
|
||||
runtime = _make_runtime(thread_id="thread-1")
|
||||
assert session_tools._get_runtime_thread_id(runtime) == "thread-1"
|
||||
|
||||
def test_ignores_file_thread_id(self):
|
||||
"""file_thread_id 继承自父会话,不应作为当前会话标识"""
|
||||
runtime = _make_runtime(thread_id="thread-1", file_thread_id="file-thread-1")
|
||||
assert session_tools._get_runtime_thread_id(runtime) == "thread-1"
|
||||
|
||||
def test_raises_when_missing(self):
|
||||
runtime = _make_runtime(thread_id="", file_thread_id=None)
|
||||
with pytest.raises(ValueError, match="thread_id"):
|
||||
session_tools._get_runtime_thread_id(runtime)
|
||||
|
||||
|
||||
# ── 权限校验测试 ──────────────────────────────────────────
|
||||
|
||||
|
||||
class TestCheckSessionAccess:
|
||||
@pytest.mark.asyncio
|
||||
async def test_allows_own_conversation(self):
|
||||
"""用户访问自己的会话 → 允许"""
|
||||
conv = MagicMock()
|
||||
conv.uid = "user-1"
|
||||
conv.parent_thread_id = None
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=conv)
|
||||
MockPg.get_async_session_context = MagicMock()
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
# Should not raise
|
||||
await session_tools._check_session_access("user-1", "thread-1")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_allows_via_parent_chain(self):
|
||||
"""沿 parent 链上溯找到自己的会话 → 允许"""
|
||||
child_conv = MagicMock()
|
||||
child_conv.uid = "other-user"
|
||||
child_conv.parent_thread_id = "parent-thread"
|
||||
|
||||
parent_conv = MagicMock()
|
||||
parent_conv.uid = "user-1"
|
||||
parent_conv.parent_thread_id = None
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(side_effect=[child_conv, parent_conv])
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
await session_tools._check_session_access("user-1", "child-thread")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_unauthorized_conversation(self):
|
||||
"""非自己的会话且不在子会话树内 → 拒绝"""
|
||||
conv = MagicMock()
|
||||
conv.uid = "other-user"
|
||||
conv.parent_thread_id = None
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=conv)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(PermissionError, match="无权访问"):
|
||||
await session_tools._check_session_access("user-1", "other-thread")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_nonexistent_conversation(self):
|
||||
"""会话不存在 → 抛出 ValueError"""
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=None)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(ValueError, match="会话不存在"):
|
||||
await session_tools._check_session_access("user-1", "nonexistent")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prevents_infinite_loop_in_parent_chain(self):
|
||||
"""parent 链出现循环时不会无限循环"""
|
||||
conv_a = MagicMock()
|
||||
conv_a.uid = "other-user"
|
||||
conv_a.parent_thread_id = "thread-b"
|
||||
|
||||
conv_b = MagicMock()
|
||||
conv_b.uid = "other-user"
|
||||
conv_b.parent_thread_id = "thread-a" # 循环回 A
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(side_effect=[conv_a, conv_b])
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(PermissionError, match="无权访问"):
|
||||
await session_tools._check_session_access("user-1", "thread-a")
|
||||
|
||||
|
||||
# ── list_sessions 测试 ────────────────────────────────────
|
||||
|
||||
|
||||
class TestListSessions:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_formatted_list(self):
|
||||
conv = MagicMock()
|
||||
conv.thread_id = "thread-1"
|
||||
conv.agent_id = "agent-1"
|
||||
conv.status = "active"
|
||||
conv.title = "测试会话"
|
||||
conv.parent_thread_id = None
|
||||
conv.updated_at = "2026-06-13"
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.list_conversations = AsyncMock(return_value=[conv])
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.list_sessions.coroutine(
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "thread-1" in result
|
||||
assert "agent-1" in result
|
||||
assert "测试会话" in result
|
||||
assert "[子智能体]" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_marks_subagent_sessions(self):
|
||||
conv = MagicMock()
|
||||
conv.thread_id = "sub-thread-1"
|
||||
conv.agent_id = "sub-agent"
|
||||
conv.status = "active"
|
||||
conv.title = "子智能体会话"
|
||||
conv.parent_thread_id = "parent-thread"
|
||||
conv.updated_at = "2026-06-13"
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.list_conversations = AsyncMock(return_value=[conv])
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.list_sessions.coroutine(
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "[子智能体]" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_message(self):
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.list_conversations = AsyncMock(return_value=[])
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.list_sessions.coroutine(
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "无可见" in result
|
||||
|
||||
|
||||
# ── get_session_history 测试 ──────────────────────────────
|
||||
|
||||
|
||||
class TestGetSessionHistory:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_messages(self):
|
||||
msg = MagicMock()
|
||||
msg.role = "user"
|
||||
msg.content = "你好"
|
||||
msg.message_type = "text"
|
||||
|
||||
conv = MagicMock()
|
||||
conv.id = 1
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalars.return_value.unique.return_value.all.return_value = [msg]
|
||||
mock_db.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=conv)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=mock_db)
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.get_session_history.coroutine(
|
||||
thread_id="thread-1",
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "[user] 你好" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_filters_tool_messages_by_default(self):
|
||||
msg_tool = MagicMock()
|
||||
msg_tool.role = "assistant"
|
||||
msg_tool.content = "tool result"
|
||||
msg_tool.message_type = "tool_result"
|
||||
|
||||
msg_text = MagicMock()
|
||||
msg_text.role = "user"
|
||||
msg_text.content = "hello"
|
||||
msg_text.message_type = "text"
|
||||
|
||||
conv = MagicMock()
|
||||
conv.id = 1
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalars.return_value.unique.return_value.all.return_value = [msg_tool, msg_text]
|
||||
mock_db.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=conv)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=mock_db)
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.get_session_history.coroutine(
|
||||
thread_id="thread-1",
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "tool result" not in result
|
||||
assert "hello" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_includes_tool_messages_when_requested(self):
|
||||
msg_tool = MagicMock()
|
||||
msg_tool.role = "assistant"
|
||||
msg_tool.content = "tool result"
|
||||
msg_tool.message_type = "tool_result"
|
||||
|
||||
conv = MagicMock()
|
||||
conv.id = 1
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalars.return_value.unique.return_value.all.return_value = [msg_tool]
|
||||
mock_db.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=conv)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=mock_db)
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.get_session_history.coroutine(
|
||||
thread_id="thread-1",
|
||||
include_tools=True,
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "tool result" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_truncates_long_content(self):
|
||||
msg = MagicMock()
|
||||
msg.role = "assistant"
|
||||
msg.content = "x" * 5000
|
||||
msg.message_type = "text"
|
||||
|
||||
conv = MagicMock()
|
||||
conv.id = 1
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalars.return_value.unique.return_value.all.return_value = [msg]
|
||||
mock_db.execute = AsyncMock(return_value=mock_result)
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=conv)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=mock_db)
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.get_session_history.coroutine(
|
||||
thread_id="thread-1",
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "已截断" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_message_when_no_conversation(self):
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.get_conversation_by_thread_id = AsyncMock(return_value=None)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.get_session_history.coroutine(
|
||||
thread_id="thread-1",
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "暂无消息" in result
|
||||
|
||||
|
||||
# ── send_to_session 测试 ──────────────────────────────────
|
||||
|
||||
|
||||
class TestSendToSession:
|
||||
@pytest.mark.asyncio
|
||||
async def test_sends_message_with_annotation(self):
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.ConversationRepository") as MockRepo, \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_repo = MockRepo.return_value
|
||||
mock_repo.add_message_by_thread_id = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=MagicMock())
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.send_to_session.coroutine(
|
||||
thread_id="other-thread",
|
||||
message="hello",
|
||||
runtime=_make_runtime(thread_id="my-thread"),
|
||||
)
|
||||
assert "已发送" in result
|
||||
# 验证消息带有来源标注
|
||||
call_args = mock_repo.add_message_by_thread_id.call_args
|
||||
assert "[来自会话 my-thread]" in call_args.kwargs["content"]
|
||||
assert call_args.kwargs["extra_metadata"]["source_type"] == "inter_session"
|
||||
assert call_args.kwargs["extra_metadata"]["source_thread_id"] == "my-thread"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_sending_to_self(self):
|
||||
"""禁止向自身会话发送消息"""
|
||||
result = await session_tools.send_to_session.coroutine(
|
||||
thread_id="my-thread",
|
||||
message="hello",
|
||||
runtime=_make_runtime(thread_id="my-thread"),
|
||||
)
|
||||
assert "不能向当前所在的会话发送消息" in result
|
||||
|
||||
|
||||
# ── get_agent_progress 测试 ───────────────────────────────
|
||||
|
||||
|
||||
class TestGetAgentProgress:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_run_info(self):
|
||||
run = MagicMock()
|
||||
run.id = "run-123"
|
||||
run.status = "running"
|
||||
run.run_type = "chat"
|
||||
run.agent_id = "agent-1"
|
||||
run.created_at = "2026-06-13"
|
||||
run.started_at = "2026-06-13T10:00"
|
||||
run.finished_at = None
|
||||
run.error_message = None
|
||||
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = run
|
||||
mock_db.execute = AsyncMock(return_value=mock_result)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=mock_db)
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.get_agent_progress.coroutine(
|
||||
thread_id="thread-1",
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "run-123" in result
|
||||
assert "running" in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_no_run_message(self):
|
||||
with patch("yuxi.agents.toolkits.buildin.session_tools._check_session_access", new_callable=AsyncMock), \
|
||||
patch("yuxi.agents.toolkits.buildin.session_tools.pg_manager") as MockPg:
|
||||
mock_db = MagicMock()
|
||||
mock_result = MagicMock()
|
||||
mock_result.scalar_one_or_none.return_value = None
|
||||
mock_db.execute = AsyncMock(return_value=mock_result)
|
||||
MockPg.get_async_session_context.return_value.__aenter__ = AsyncMock(return_value=mock_db)
|
||||
MockPg.get_async_session_context.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
result = await session_tools.get_agent_progress.coroutine(
|
||||
thread_id="thread-1",
|
||||
runtime=_make_runtime(),
|
||||
)
|
||||
assert "暂无运行记录" in result
|
||||
|
||||
|
||||
# ── LITE 模式测试 ─────────────────────────────────────────
|
||||
|
||||
|
||||
class TestLiteMode:
|
||||
def test_lite_mode_flag(self):
|
||||
"""验证 LITE_MODE 常量正确读取环境变量"""
|
||||
# 当前测试环境未设置 LITE_MODE,应为 False
|
||||
assert session_tools._LITE_MODE is False
|
||||
|
||||
def test_lite_mode_tools_not_in_global_registry(self):
|
||||
"""非 LITE 模式下,工具应已注册到全局列表"""
|
||||
from yuxi.agents.toolkits.registry import _all_tool_instances, _extra_registry
|
||||
|
||||
tool_names = {t.name for t in _all_tool_instances}
|
||||
assert "list_sessions" in tool_names
|
||||
assert "get_session_history" in tool_names
|
||||
assert "send_to_session" in tool_names
|
||||
assert "get_agent_progress" in tool_names
|
||||
|
||||
assert "list_sessions" in _extra_registry
|
||||
assert "get_session_history" in _extra_registry
|
||||
assert "send_to_session" in _extra_registry
|
||||
assert "get_agent_progress" in _extra_registry
|
||||
Loading…
Reference in New Issue
Block a user