eac90941ef
- Implement SDK session with secretary model for tool approval flow - Add audit logging for tool usage and permission decisions - Support Feishu card interactions for approval requests - Add new commands for task interruption and progress checking - Remove old test files and update documentation
277 lines
10 KiB
Python
277 lines
10 KiB
Python
"""Session registry: tracks active project sessions by conversation ID."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass, field, asdict
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from agent.sdk_session import (
|
|
SDKSession,
|
|
SessionProgress,
|
|
DEFAULT_PERMISSION_MODE,
|
|
VALID_PERMISSION_MODES,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_IDLE_TIMEOUT = 30 * 60
|
|
PERSISTENCE_FILE = Path(__file__).parent.parent / "data" / "sessions.json"
|
|
|
|
|
|
@dataclass
|
|
class Session:
|
|
conv_id: str
|
|
cwd: str
|
|
owner_id: str = ""
|
|
last_activity: float = 0.0
|
|
idle_timeout: int = DEFAULT_IDLE_TIMEOUT
|
|
permission_mode: str = field(default_factory=lambda: DEFAULT_PERMISSION_MODE)
|
|
chat_id: str | None = None
|
|
# Runtime only — not serialized
|
|
sdk_session: SDKSession | None = field(default=None, repr=False)
|
|
|
|
def touch(self) -> None:
|
|
self.last_activity = asyncio.get_event_loop().time()
|
|
|
|
def to_dict(self) -> dict:
|
|
d = asdict(self)
|
|
d.pop("sdk_session", None)
|
|
return d
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict) -> "Session":
|
|
data.pop("sdk_session", None)
|
|
# Migration: remove old cc_runner fields if present in persisted data
|
|
for old_key in ("cc_session_id", "started", "cc_timeout"):
|
|
data.pop(old_key, None)
|
|
return cls(**data)
|
|
|
|
|
|
class SessionManager:
|
|
"""Registry of active Claude Code project sessions with persistence and user isolation."""
|
|
|
|
def __init__(self) -> None:
|
|
self._sessions: dict[str, Session] = {}
|
|
self._lock = asyncio.Lock()
|
|
self._reaper_task: Optional[asyncio.Task] = None
|
|
|
|
async def start(self) -> None:
|
|
self._load()
|
|
loop = asyncio.get_event_loop()
|
|
self._reaper_task = loop.create_task(self._reaper_loop())
|
|
|
|
async def stop(self) -> None:
|
|
if self._reaper_task:
|
|
self._reaper_task.cancel()
|
|
async with self._lock:
|
|
for session in self._sessions.values():
|
|
if session.sdk_session:
|
|
await session.sdk_session.close()
|
|
self._sessions.clear()
|
|
if PERSISTENCE_FILE.exists():
|
|
PERSISTENCE_FILE.unlink()
|
|
|
|
async def create(
|
|
self,
|
|
conv_id: str,
|
|
working_dir: str,
|
|
owner_id: str = "",
|
|
idle_timeout: int = DEFAULT_IDLE_TIMEOUT,
|
|
permission_mode: str = DEFAULT_PERMISSION_MODE,
|
|
chat_id: str | None = None,
|
|
) -> Session:
|
|
async with self._lock:
|
|
session = Session(
|
|
conv_id=conv_id,
|
|
cwd=working_dir,
|
|
owner_id=owner_id,
|
|
idle_timeout=idle_timeout,
|
|
permission_mode=permission_mode,
|
|
chat_id=chat_id,
|
|
)
|
|
self._sessions[conv_id] = session
|
|
self._save()
|
|
logger.info(
|
|
"Created session %s (owner=...%s) in %s (idle=%ds, perm=%s)",
|
|
conv_id, owner_id[-8:] if owner_id else "-", working_dir,
|
|
idle_timeout, permission_mode,
|
|
)
|
|
return session
|
|
|
|
# --- Secretary model: async send (returns immediately) ---
|
|
|
|
async def send_message(
|
|
self, conv_id: str, message: str,
|
|
user_id: Optional[str] = None, chat_id: Optional[str] = None,
|
|
) -> str:
|
|
"""Send a message to the session. Returns immediately; result pushed to Feishu on completion."""
|
|
session = self._get_session(conv_id, user_id)
|
|
session.touch()
|
|
self._ensure_sdk_session(session, chat_id)
|
|
return await session.sdk_session.send(message, chat_id)
|
|
|
|
async def send_and_wait(
|
|
self, conv_id: str, message: str,
|
|
user_id: Optional[str] = None, chat_id: Optional[str] = None,
|
|
) -> str:
|
|
"""Send and wait for completion. For LLM agent tool calls that need the result."""
|
|
session = self._get_session(conv_id, user_id)
|
|
session.touch()
|
|
self._ensure_sdk_session(session, chat_id)
|
|
return await session.sdk_session.send_and_wait(message, chat_id)
|
|
|
|
# --- Kept for backward compatibility (used by bot/commands.py _cmd_new) ---
|
|
|
|
async def send(
|
|
self, conv_id: str, message: str,
|
|
user_id: Optional[str] = None, direct: bool = False,
|
|
) -> str:
|
|
"""Backward-compatible send. Maps to send_message (async, secretary model)."""
|
|
from orchestrator.tools import get_current_chat
|
|
chat_id = get_current_chat()
|
|
return await self.send_message(conv_id, message, user_id=user_id, chat_id=chat_id)
|
|
|
|
# --- Progress, interrupt, approve ---
|
|
|
|
def get_progress(self, conv_id: str, user_id: Optional[str] = None) -> SessionProgress | None:
|
|
"""Query session progress. Primary interface for the secretary AI."""
|
|
session = self._sessions.get(conv_id)
|
|
if not session:
|
|
return None
|
|
if session.owner_id and user_id and session.owner_id != user_id:
|
|
return None
|
|
if session.sdk_session:
|
|
return session.sdk_session.get_progress()
|
|
return SessionProgress()
|
|
|
|
async def interrupt(self, conv_id: str, user_id: Optional[str] = None) -> bool:
|
|
"""Interrupt the currently running task in a session."""
|
|
session = self._get_session(conv_id, user_id)
|
|
if session.sdk_session:
|
|
await session.sdk_session.interrupt()
|
|
return True
|
|
return False
|
|
|
|
async def approve(self, conv_id: str, approved: bool) -> None:
|
|
"""Resolve a pending tool approval for a session."""
|
|
session = self._sessions.get(conv_id)
|
|
if session and session.sdk_session:
|
|
await session.sdk_session.approve(approved)
|
|
|
|
# --- Close, list, permission ---
|
|
|
|
async def close(self, conv_id: str, user_id: Optional[str] = None) -> bool:
|
|
async with self._lock:
|
|
session = self._sessions.get(conv_id)
|
|
if session is None:
|
|
return False
|
|
if session.owner_id and user_id and session.owner_id != user_id:
|
|
raise PermissionError(f"Session {conv_id} belongs to another user")
|
|
if session.sdk_session:
|
|
await session.sdk_session.close()
|
|
del self._sessions[conv_id]
|
|
self._save()
|
|
logger.info("Closed session %s", conv_id)
|
|
return True
|
|
|
|
def list_sessions(self, user_id: Optional[str] = None) -> list[dict]:
|
|
sessions = self._sessions.values()
|
|
if user_id:
|
|
sessions = [s for s in sessions if not s.owner_id or s.owner_id == user_id]
|
|
return [
|
|
{
|
|
"conv_id": s.conv_id,
|
|
"cwd": s.cwd,
|
|
"owner_id": s.owner_id[-8:] if s.owner_id else None,
|
|
"busy": s.sdk_session._busy if s.sdk_session else False,
|
|
"idle_timeout": s.idle_timeout,
|
|
"permission_mode": s.permission_mode,
|
|
}
|
|
for s in sessions
|
|
]
|
|
|
|
def set_permission_mode(self, conv_id: str, mode: str, user_id: Optional[str] = None) -> None:
|
|
"""Change the permission mode for an existing session."""
|
|
session = self._sessions.get(conv_id)
|
|
if session is None:
|
|
raise KeyError(f"No session for conv_id={conv_id!r}")
|
|
if session.owner_id and user_id and session.owner_id != user_id:
|
|
raise PermissionError(f"Session {conv_id} belongs to another user")
|
|
if mode not in VALID_PERMISSION_MODES:
|
|
raise ValueError(f"Invalid permission mode {mode!r}. Valid: {VALID_PERMISSION_MODES}")
|
|
session.permission_mode = mode
|
|
if session.sdk_session:
|
|
asyncio.create_task(session.sdk_session.set_permission_mode(mode))
|
|
self._save()
|
|
logger.info("Set permission_mode=%s for session %s", mode, conv_id)
|
|
|
|
# --- Internal ---
|
|
|
|
def _get_session(self, conv_id: str, user_id: Optional[str] = None) -> Session:
|
|
session = self._sessions.get(conv_id)
|
|
if session is None:
|
|
raise KeyError(f"No session for conv_id={conv_id!r}")
|
|
if session.owner_id and user_id and session.owner_id != user_id:
|
|
raise PermissionError(f"Session {conv_id} belongs to another user")
|
|
return session
|
|
|
|
def _ensure_sdk_session(self, session: Session, chat_id: str | None = None) -> None:
|
|
if session.sdk_session is None:
|
|
session.sdk_session = SDKSession(
|
|
conv_id=session.conv_id,
|
|
cwd=session.cwd,
|
|
owner_id=session.owner_id,
|
|
permission_mode=session.permission_mode,
|
|
chat_id=chat_id or session.chat_id,
|
|
)
|
|
|
|
def _save(self) -> None:
|
|
try:
|
|
data = {cid: s.to_dict() for cid, s in self._sessions.items()}
|
|
PERSISTENCE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(PERSISTENCE_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2)
|
|
logger.debug("Saved %d sessions to %s", len(data), PERSISTENCE_FILE)
|
|
except Exception:
|
|
logger.exception("Failed to save sessions")
|
|
|
|
def _load(self) -> None:
|
|
if not PERSISTENCE_FILE.exists():
|
|
return
|
|
try:
|
|
with open(PERSISTENCE_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
for cid, sdata in data.items():
|
|
self._sessions[cid] = Session.from_dict(sdata)
|
|
logger.info("Loaded %d sessions from %s", len(self._sessions), PERSISTENCE_FILE)
|
|
except Exception:
|
|
logger.exception("Failed to load sessions")
|
|
|
|
async def _reaper_loop(self) -> None:
|
|
while True:
|
|
await asyncio.sleep(60)
|
|
await self._reap_idle()
|
|
|
|
async def _reap_idle(self) -> None:
|
|
now = asyncio.get_event_loop().time()
|
|
async with self._lock:
|
|
to_close = []
|
|
for cid, s in self._sessions.items():
|
|
if s.last_activity > 0 and (now - s.last_activity) > s.idle_timeout:
|
|
to_close.append(cid)
|
|
for cid in to_close:
|
|
session = self._sessions[cid]
|
|
if session.sdk_session:
|
|
await session.sdk_session.close()
|
|
del self._sessions[cid]
|
|
logger.info("Reaped idle session %s", cid)
|
|
if to_close:
|
|
self._save()
|
|
|
|
|
|
manager = SessionManager()
|