feat: 实现用户权限控制、会话管理和审计日志功能
- 添加用户权限检查功能,支持配置允许使用的用户列表 - 实现会话管理功能,包括会话创建、关闭、列表和切换 - 新增审计日志模块,记录所有交互信息 - 改进WebSocket连接,增加自动重连机制 - 添加健康检查端点,包含Claude服务可用性测试 - 实现会话持久化功能,重启后恢复会话状态 - 增加命令行功能支持,包括/new、/list、/close等命令 - 优化消息处理流程,支持直接传递模式
This commit is contained in:
+34
-2
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user