feat: 实现用户权限控制、会话管理和审计日志功能

- 添加用户权限检查功能,支持配置允许使用的用户列表
- 实现会话管理功能,包括会话创建、关闭、列表和切换
- 新增审计日志模块,记录所有交互信息
- 改进WebSocket连接,增加自动重连机制
- 添加健康检查端点,包含Claude服务可用性测试
- 实现会话持久化功能,重启后恢复会话状态
- 增加命令行功能支持,包括/new、/list、/close等命令
- 优化消息处理流程,支持直接传递模式
This commit is contained in:
Yuyao Huang (Sam)
2026-03-28 08:39:32 +08:00
parent 29c0f2e403
commit 6307deb701
11 changed files with 921 additions and 222 deletions
+34 -2
View File
@@ -5,8 +5,10 @@ Uses LangChain 1.x tool-calling pattern: bind_tools + manual agentic loop.
from __future__ import annotations
import asyncio
import json
import logging
import re
from collections import defaultdict
from typing import Dict, List, Optional
@@ -19,8 +21,9 @@ from langchain_core.messages import (
)
from langchain_openai import ChatOpenAI
from agent.manager import manager
from config import OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL, WORKING_DIR
from orchestrator.tools import TOOLS
from orchestrator.tools import TOOLS, set_current_user
logger = logging.getLogger(__name__)
@@ -52,6 +55,12 @@ Guidelines:
MAX_ITERATIONS = 10
_TOOL_MAP = {t.name: t for t in TOOLS}
COMMAND_PATTERN = re.compile(r"^/(new|list|close|switch|retry|help)", re.IGNORECASE)
def _looks_like_command(text: str) -> bool:
return bool(COMMAND_PATTERN.match(text.strip()))
class OrchestrationAgent:
"""Per-user agent with conversation history and active session tracking."""
@@ -69,6 +78,8 @@ class OrchestrationAgent:
self._history: Dict[str, List[BaseMessage]] = defaultdict(list)
# user_id -> most recently active conv_id
self._active_conv: Dict[str, Optional[str]] = defaultdict(lambda: None)
# user_id -> asyncio.Lock (prevents concurrent processing per user)
self._user_locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
def _build_system_prompt(self, user_id: str) -> str:
conv_id = self._active_conv[user_id]
@@ -81,13 +92,35 @@ class OrchestrationAgent:
active_session_line=active_line,
)
def get_active_conv(self, user_id: str) -> Optional[str]:
return self._active_conv.get(user_id)
async def run(self, user_id: str, text: str) -> str:
"""Process a user message and return the agent's reply."""
async with self._user_locks[user_id]:
return await self._run_locked(user_id, text)
async def _run_locked(self, user_id: str, text: str) -> str:
"""Internal implementation, must be called with user lock held."""
set_current_user(user_id)
active_conv = self._active_conv[user_id]
short_uid = user_id[-8:]
logger.info(">>> user=...%s conv=%s msg=%r", short_uid, active_conv, text[:80])
logger.debug(" history_len=%d", len(self._history[user_id]))
# Passthrough mode: if active session and not a command, bypass LLM
if active_conv and not _looks_like_command(text):
try:
reply = await manager.send(active_conv, text, user_id=user_id)
logger.info("<<< [passthrough] reply: %r", reply[:120])
return reply
except KeyError:
logger.warning("Session %s no longer exists, clearing active_conv", active_conv)
self._active_conv[user_id] = None
except Exception as exc:
logger.exception("Passthrough error for user=%s", user_id)
return f"[Error] {exc}"
messages: List[BaseMessage] = (
[SystemMessage(content=self._build_system_prompt(user_id))]
+ self._history[user_id]
@@ -161,5 +194,4 @@ class OrchestrationAgent:
return reply
# Module-level singleton
agent = OrchestrationAgent()
+46 -11
View File
@@ -4,6 +4,7 @@ from __future__ import annotations
import json
import uuid
from contextvars import ContextVar
from pathlib import Path
from typing import Optional, Type
@@ -13,6 +14,16 @@ from pydantic import BaseModel, Field
from agent.manager import manager
from config import WORKING_DIR
_current_user_id: ContextVar[Optional[str]] = ContextVar("current_user_id", default=None)
def set_current_user(user_id: Optional[str]) -> None:
_current_user_id.set(user_id)
def get_current_user() -> Optional[str]:
return _current_user_id.get()
def _resolve_dir(working_dir: str) -> Path:
"""
@@ -21,14 +32,21 @@ def _resolve_dir(working_dir: str) -> Path:
Rules:
- Absolute paths are used as-is (but must stay within WORKING_DIR for safety).
- Relative paths / bare names are joined onto WORKING_DIR.
- Path traversal attempts (..) are blocked.
- The resolved directory is created if it doesn't exist.
"""
working_dir = working_dir.strip()
if ".." in working_dir.split("/") or ".." in working_dir.split("\\"):
raise ValueError(
"Path traversal not allowed. Use a subfolder name or path inside the working directory."
)
p = Path(working_dir)
if not p.is_absolute():
p = WORKING_DIR / p
p = p.resolve()
# Safety: must be inside WORKING_DIR
try:
p.relative_to(WORKING_DIR)
except ValueError:
@@ -55,6 +73,8 @@ class CreateConversationInput(BaseModel):
),
)
initial_message: Optional[str] = Field(None, description="Optional first message to send after spawning")
idle_timeout: Optional[int] = Field(None, description="Idle timeout in seconds (default 1800)")
cc_timeout: Optional[float] = Field(None, description="Claude Code execution timeout in seconds (default 300)")
class SendToConversationInput(BaseModel):
@@ -79,17 +99,24 @@ class CreateConversationTool(BaseTool):
)
args_schema: Type[BaseModel] = CreateConversationInput
def _run(self, working_dir: str, initial_message: Optional[str] = None) -> str:
def _run(self, working_dir: str, initial_message: Optional[str] = None, idle_timeout: Optional[int] = None, cc_timeout: Optional[float] = None) -> str:
raise NotImplementedError("Use async version")
async def _arun(self, working_dir: str, initial_message: Optional[str] = None) -> str:
async def _arun(self, working_dir: str, initial_message: Optional[str] = None, idle_timeout: Optional[int] = None, cc_timeout: Optional[float] = None) -> str:
try:
resolved = _resolve_dir(working_dir)
except ValueError as exc:
return json.dumps({"error": str(exc)})
user_id = get_current_user()
conv_id = str(uuid.uuid4())[:8]
await manager.create(conv_id, str(resolved))
await manager.create(
conv_id,
str(resolved),
owner_id=user_id or "",
idle_timeout=idle_timeout or 1800,
cc_timeout=cc_timeout or 300.0,
)
result: dict = {
"conv_id": conv_id,
@@ -97,7 +124,7 @@ class CreateConversationTool(BaseTool):
}
if initial_message:
output = await manager.send(conv_id, initial_message)
output = await manager.send(conv_id, initial_message, user_id=user_id)
result["response"] = output
else:
result["status"] = "Session created. Send a message to start working."
@@ -117,11 +144,14 @@ class SendToConversationTool(BaseTool):
raise NotImplementedError("Use async version")
async def _arun(self, conv_id: str, message: str) -> str:
user_id = get_current_user()
try:
output = await manager.send(conv_id, message)
output = await manager.send(conv_id, message, user_id=user_id)
return json.dumps({"conv_id": conv_id, "response": output}, ensure_ascii=False)
except KeyError:
return json.dumps({"error": f"No active session for conv_id={conv_id!r}"})
except PermissionError as e:
return json.dumps({"error": str(e)})
class ListConversationsTool(BaseTool):
@@ -132,7 +162,8 @@ class ListConversationsTool(BaseTool):
raise NotImplementedError("Use async version")
async def _arun(self) -> str:
sessions = manager.list_sessions()
user_id = get_current_user()
sessions = manager.list_sessions(user_id=user_id)
if not sessions:
return "No active sessions."
return json.dumps(sessions, ensure_ascii=False, indent=2)
@@ -147,10 +178,14 @@ class CloseConversationTool(BaseTool):
raise NotImplementedError("Use async version")
async def _arun(self, conv_id: str) -> str:
closed = await manager.close(conv_id)
if closed:
return f"Session {conv_id!r} closed."
return f"Session {conv_id!r} not found."
user_id = get_current_user()
try:
closed = await manager.close(conv_id, user_id=user_id)
if closed:
return f"Session {conv_id!r} closed."
return f"Session {conv_id!r} not found."
except PermissionError as e:
return str(e)
# Module-level tool list for easy import