feat: add SDK session implementation with approval flow and audit logging
- 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
This commit is contained in:
Yuyao Huang
@@ -58,6 +58,61 @@ def log_interaction(
|
||||
logger.exception("Failed to log audit entry for session %s", conv_id)
|
||||
|
||||
|
||||
def log_tool_use(
|
||||
session_id: str,
|
||||
tool_name: str,
|
||||
tool_input: dict,
|
||||
tool_response: Optional[object] = None,
|
||||
) -> None:
|
||||
"""Log a tool call to the audit JSONL file."""
|
||||
try:
|
||||
_ensure_audit_dir()
|
||||
log_file = AUDIT_DIR / f"{session_id}.jsonl"
|
||||
|
||||
entry = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"type": "tool_use",
|
||||
"session_id": session_id,
|
||||
"tool_name": tool_name,
|
||||
"tool_input": str(tool_input)[:500],
|
||||
}
|
||||
if tool_response is not None:
|
||||
entry["tool_response"] = str(tool_response)[:500]
|
||||
|
||||
with open(log_file, "a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
||||
|
||||
except Exception:
|
||||
logger.exception("Failed to log tool use for session %s", session_id)
|
||||
|
||||
|
||||
def log_permission_decision(
|
||||
conv_id: str,
|
||||
tool_name: str,
|
||||
tool_input: dict,
|
||||
approved: bool,
|
||||
) -> None:
|
||||
"""Log a permission approval/denial decision."""
|
||||
try:
|
||||
_ensure_audit_dir()
|
||||
log_file = AUDIT_DIR / f"{conv_id}.jsonl"
|
||||
|
||||
entry = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"type": "permission_decision",
|
||||
"conv_id": conv_id,
|
||||
"tool_name": tool_name,
|
||||
"tool_input": str(tool_input)[:300],
|
||||
"approved": approved,
|
||||
}
|
||||
|
||||
with open(log_file, "a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
|
||||
|
||||
except Exception:
|
||||
logger.exception("Failed to log permission decision for session %s", conv_id)
|
||||
|
||||
|
||||
def get_audit_log(conv_id: str, limit: int = 100) -> list[dict]:
|
||||
"""Read the audit log for a session."""
|
||||
log_file = AUDIT_DIR / f"{conv_id}.jsonl"
|
||||
|
||||
+105
-93
@@ -5,18 +5,20 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
from typing import Optional
|
||||
|
||||
from agent.cc_runner import run_claude, DEFAULT_PERMISSION_MODE, VALID_PERMISSION_MODES
|
||||
from agent.audit import log_interaction
|
||||
from agent.sdk_session import (
|
||||
SDKSession,
|
||||
SessionProgress,
|
||||
DEFAULT_PERMISSION_MODE,
|
||||
VALID_PERMISSION_MODES,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_IDLE_TIMEOUT = 30 * 60
|
||||
DEFAULT_CC_TIMEOUT = 300.0
|
||||
PERSISTENCE_FILE = Path(__file__).parent.parent / "data" / "sessions.json"
|
||||
|
||||
|
||||
@@ -25,21 +27,27 @@ class Session:
|
||||
conv_id: str
|
||||
cwd: str
|
||||
owner_id: str = ""
|
||||
cc_session_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
last_activity: float = 0.0
|
||||
started: bool = False
|
||||
idle_timeout: int = DEFAULT_IDLE_TIMEOUT
|
||||
cc_timeout: float = DEFAULT_CC_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:
|
||||
return asdict(self)
|
||||
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)
|
||||
|
||||
|
||||
@@ -60,6 +68,9 @@ class SessionManager:
|
||||
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()
|
||||
@@ -70,8 +81,8 @@ class SessionManager:
|
||||
working_dir: str,
|
||||
owner_id: str = "",
|
||||
idle_timeout: int = DEFAULT_IDLE_TIMEOUT,
|
||||
cc_timeout: float = DEFAULT_CC_TIMEOUT,
|
||||
permission_mode: str = DEFAULT_PERMISSION_MODE,
|
||||
chat_id: str | None = None,
|
||||
) -> Session:
|
||||
async with self._lock:
|
||||
session = Session(
|
||||
@@ -79,103 +90,79 @@ class SessionManager:
|
||||
cwd=working_dir,
|
||||
owner_id=owner_id,
|
||||
idle_timeout=idle_timeout,
|
||||
cc_timeout=cc_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, cc=%.0fs, perm=%s)",
|
||||
"Created session %s (owner=...%s) in %s (idle=%ds, perm=%s)",
|
||||
conv_id, owner_id[-8:] if owner_id else "-", working_dir,
|
||||
idle_timeout, cc_timeout, permission_mode,
|
||||
idle_timeout, permission_mode,
|
||||
)
|
||||
return session
|
||||
|
||||
async def send(self, conv_id: str, message: str, user_id: Optional[str] = None, direct: bool = False) -> str:
|
||||
async with self._lock:
|
||||
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")
|
||||
session.touch()
|
||||
cwd = session.cwd
|
||||
cc_session_id = session.cc_session_id
|
||||
cc_timeout = session.cc_timeout
|
||||
permission_mode = session.permission_mode
|
||||
first_message = not session.started
|
||||
if first_message:
|
||||
session.started = True
|
||||
self._save()
|
||||
# --- Secretary model: async send (returns immediately) ---
|
||||
|
||||
if not direct and cc_timeout > 60:
|
||||
from agent.task_runner import task_runner
|
||||
from orchestrator.tools import get_current_chat, set_current_chat, set_current_user
|
||||
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)
|
||||
|
||||
chat_id = get_current_chat()
|
||||
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)
|
||||
|
||||
async def run_task():
|
||||
output = await run_claude(
|
||||
message,
|
||||
cwd=cwd,
|
||||
cc_session_id=cc_session_id,
|
||||
resume=not first_message,
|
||||
timeout=cc_timeout,
|
||||
permission_mode=permission_mode,
|
||||
)
|
||||
log_interaction(
|
||||
conv_id=conv_id,
|
||||
prompt=message,
|
||||
response=output,
|
||||
cwd=cwd,
|
||||
user_id=user_id,
|
||||
)
|
||||
return output
|
||||
# --- Kept for backward compatibility (used by bot/commands.py _cmd_new) ---
|
||||
|
||||
async def on_task_complete(task) -> None:
|
||||
if not chat_id or not user_id or not task.result:
|
||||
return
|
||||
set_current_user(user_id)
|
||||
set_current_chat(chat_id)
|
||||
from orchestrator.agent import agent
|
||||
follow_up = (
|
||||
f"CC task completed. Output:\n{task.result}\n\n"
|
||||
f"Original request was: {message}\n\n"
|
||||
"If the user asked you to send a file, use send_file now. "
|
||||
"Otherwise just acknowledge completion."
|
||||
)
|
||||
reply = await agent.run(user_id, follow_up)
|
||||
if reply:
|
||||
from bot.feishu import send_text
|
||||
await send_text(chat_id, "chat_id", reply)
|
||||
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)
|
||||
|
||||
task_id = await task_runner.submit(
|
||||
run_task(),
|
||||
description=f"CC session {conv_id}: {message[:50]}",
|
||||
notify_chat_id=chat_id,
|
||||
user_id=user_id,
|
||||
on_complete=on_task_complete,
|
||||
)
|
||||
return f"⏳ Task #{task_id} started (timeout: {int(cc_timeout)}s). I'll notify you when it's done."
|
||||
# --- Progress, interrupt, approve ---
|
||||
|
||||
output = await run_claude(
|
||||
message,
|
||||
cwd=cwd,
|
||||
cc_session_id=cc_session_id,
|
||||
resume=not first_message,
|
||||
timeout=cc_timeout,
|
||||
permission_mode=permission_mode,
|
||||
)
|
||||
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()
|
||||
|
||||
log_interaction(
|
||||
conv_id=conv_id,
|
||||
prompt=message,
|
||||
response=output,
|
||||
cwd=cwd,
|
||||
user_id=user_id,
|
||||
)
|
||||
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
|
||||
|
||||
return output
|
||||
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:
|
||||
@@ -184,6 +171,8 @@ class SessionManager:
|
||||
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)
|
||||
@@ -198,10 +187,8 @@ class SessionManager:
|
||||
"conv_id": s.conv_id,
|
||||
"cwd": s.cwd,
|
||||
"owner_id": s.owner_id[-8:] if s.owner_id else None,
|
||||
"cc_session_id": s.cc_session_id,
|
||||
"started": s.started,
|
||||
"busy": s.sdk_session._busy if s.sdk_session else False,
|
||||
"idle_timeout": s.idle_timeout,
|
||||
"cc_timeout": s.cc_timeout,
|
||||
"permission_mode": s.permission_mode,
|
||||
}
|
||||
for s in sessions
|
||||
@@ -217,9 +204,31 @@ class SessionManager:
|
||||
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()}
|
||||
@@ -255,6 +264,9 @@ class SessionManager:
|
||||
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:
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
"""SDK hooks for audit logging and dangerous command blocking."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from claude_agent_sdk import HookContext, HookInput, HookJSONOutput, HookMatcher
|
||||
|
||||
BLOCKED_PATTERNS = [
|
||||
r"\brm\s+-rf\s+/",
|
||||
r"\brm\s+-rf\s+~",
|
||||
r"\bformat\s+",
|
||||
r"\bmkfs\b",
|
||||
r"\bshutdown\b",
|
||||
r"\breboot\b",
|
||||
r"\bdd\s+if=",
|
||||
r":\(\)\{:\|:&\};:",
|
||||
r"\bchmod\s+777\s+/",
|
||||
r"\bchown\s+.*\s+/",
|
||||
r"\bsudo\s+rm\b",
|
||||
r"\bsudo\s+chmod\b",
|
||||
r"\bsudo\s+chown\b",
|
||||
r"\bsudo\s+dd\b",
|
||||
r"\bkill\s+-9\s+1\b",
|
||||
]
|
||||
|
||||
|
||||
async def audit_hook(
|
||||
input_data: HookInput, tool_use_id: str | None, context: HookContext
|
||||
) -> HookJSONOutput:
|
||||
"""PostToolUse hook — log tool calls to audit JSONL."""
|
||||
from agent.audit import log_tool_use
|
||||
|
||||
log_tool_use(
|
||||
session_id=input_data.get("session_id", ""),
|
||||
tool_name=input_data.get("tool_name", ""),
|
||||
tool_input=input_data.get("tool_input", {}),
|
||||
tool_response=input_data.get("tool_response"),
|
||||
)
|
||||
return {}
|
||||
|
||||
|
||||
async def deny_dangerous_hook(
|
||||
input_data: HookInput, tool_use_id: str | None, context: HookContext
|
||||
) -> HookJSONOutput:
|
||||
"""PreToolUse hook — block dangerous Bash commands."""
|
||||
if input_data.get("tool_name") != "Bash":
|
||||
return {}
|
||||
|
||||
command = input_data.get("tool_input", {}).get("command", "")
|
||||
for pattern in BLOCKED_PATTERNS:
|
||||
if re.search(pattern, command, re.IGNORECASE):
|
||||
return {
|
||||
"hookSpecificOutput": {
|
||||
"hookEventName": "PreToolUse",
|
||||
"permissionDecision": "deny",
|
||||
"permissionDecisionReason": f"Blocked by policy: matches {pattern}",
|
||||
}
|
||||
}
|
||||
return {}
|
||||
|
||||
|
||||
def build_hooks(conv_id: str) -> dict[str, list[HookMatcher]]:
|
||||
"""Build hooks configuration for a session."""
|
||||
return {
|
||||
"PostToolUse": [
|
||||
HookMatcher(matcher="Bash|Edit|Write|MultiEdit", hooks=[audit_hook]),
|
||||
],
|
||||
"PreToolUse": [
|
||||
HookMatcher(matcher="Bash", hooks=[deny_dangerous_hook]),
|
||||
],
|
||||
}
|
||||
@@ -0,0 +1,355 @@
|
||||
"""SDK-based Claude Code session — secretary model.
|
||||
|
||||
Messages are buffered in memory, not pushed to Feishu in real-time.
|
||||
Only key events (completion, error, approval) trigger notifications.
|
||||
The secretary AI queries get_progress() to answer user questions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from claude_agent_sdk import (
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ClaudeSDKClient,
|
||||
PermissionMode,
|
||||
PermissionResult,
|
||||
PermissionResultAllow,
|
||||
PermissionResultDeny,
|
||||
ResultMessage,
|
||||
SystemMessage,
|
||||
TextBlock,
|
||||
ToolPermissionContext,
|
||||
ToolUseBlock,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
VALID_PERMISSION_MODES = ["default", "acceptEdits", "plan", "bypassPermissions", "dontAsk"]
|
||||
DEFAULT_PERMISSION_MODE = "default"
|
||||
APPROVAL_TIMEOUT = 120 # seconds
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionProgress:
|
||||
"""Session progress snapshot for the secretary AI to inspect."""
|
||||
|
||||
busy: bool = False
|
||||
current_prompt: str = ""
|
||||
started_at: float = 0.0
|
||||
elapsed_seconds: float = 0.0
|
||||
text_messages: list[str] = field(default_factory=list)
|
||||
tool_calls: list[str] = field(default_factory=list)
|
||||
last_result: str = ""
|
||||
error: str = ""
|
||||
pending_approval: str = "" # non-empty → waiting for approval, value is tool description
|
||||
|
||||
|
||||
class SDKSession:
|
||||
"""One session = one long-lived ClaudeSDKClient + background message buffer loop.
|
||||
|
||||
Secretary model design:
|
||||
- _message_loop buffers all messages to memory, does NOT push to Feishu
|
||||
- Only pushes on key events: completion (ResultMessage), error, approval needed
|
||||
- get_progress() returns a snapshot for the secretary AI to inspect
|
||||
"""
|
||||
|
||||
MAX_BUFFER_TEXTS = 20
|
||||
MAX_BUFFER_TOOLS = 50
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
conv_id: str,
|
||||
cwd: str,
|
||||
owner_id: str,
|
||||
permission_mode: str = DEFAULT_PERMISSION_MODE,
|
||||
chat_id: str | None = None,
|
||||
):
|
||||
self.conv_id = conv_id
|
||||
self.cwd = cwd
|
||||
self.owner_id = owner_id
|
||||
self.permission_mode = permission_mode
|
||||
self.chat_id = chat_id
|
||||
|
||||
self.client: ClaudeSDKClient | None = None
|
||||
self.session_id: str | None = None
|
||||
|
||||
# Message buffers
|
||||
self._text_buffer: list[str] = []
|
||||
self._tool_buffer: list[str] = []
|
||||
self._last_result: str = ""
|
||||
self._error: str = ""
|
||||
self._current_prompt: str = ""
|
||||
self._started_at: float = 0.0
|
||||
|
||||
# Task state
|
||||
self._message_loop_task: asyncio.Task | None = None
|
||||
self._busy = False
|
||||
self._busy_event = asyncio.Event()
|
||||
self._busy_event.set() # initially idle
|
||||
|
||||
# Approval mechanism
|
||||
self._pending_approval: asyncio.Future | None = None
|
||||
self._pending_approval_desc: str = ""
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Create and connect the ClaudeSDKClient, start the message loop."""
|
||||
from agent.sdk_hooks import build_hooks
|
||||
|
||||
env = self._build_env()
|
||||
hooks = build_hooks(self.conv_id)
|
||||
|
||||
options = ClaudeAgentOptions(
|
||||
cwd=self.cwd,
|
||||
permission_mode=self.permission_mode,
|
||||
allowed_tools=[
|
||||
"Read", "Glob", "Grep", "Bash", "Edit", "Write",
|
||||
"MultiEdit", "WebFetch", "WebSearch",
|
||||
],
|
||||
can_use_tool=self._permission_callback,
|
||||
hooks=hooks,
|
||||
env=env,
|
||||
)
|
||||
self.client = ClaudeSDKClient(options)
|
||||
await self.client.connect()
|
||||
|
||||
self._message_loop_task = asyncio.create_task(
|
||||
self._message_loop(), name=f"sdk-loop-{self.conv_id}"
|
||||
)
|
||||
logger.info("SDKSession %s started in %s", self.conv_id, self.cwd)
|
||||
|
||||
async def send(self, prompt: str, chat_id: str | None = None) -> str:
|
||||
"""Send a message. Returns immediately; execution happens in background."""
|
||||
if not self.client:
|
||||
await self.start()
|
||||
|
||||
if chat_id:
|
||||
self.chat_id = chat_id
|
||||
|
||||
# If busy, interrupt the current task first
|
||||
if self._busy:
|
||||
await self.interrupt()
|
||||
try:
|
||||
await asyncio.wait_for(self._busy_event.wait(), timeout=10)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
|
||||
self._busy = True
|
||||
self._busy_event.clear()
|
||||
self._current_prompt = prompt
|
||||
self._started_at = time.time()
|
||||
self._last_result = ""
|
||||
self._error = ""
|
||||
self._text_buffer.clear()
|
||||
self._tool_buffer.clear()
|
||||
|
||||
await self.client.query(prompt)
|
||||
return "⏳ 已开始执行"
|
||||
|
||||
async def send_and_wait(self, prompt: str, chat_id: str | None = None) -> str:
|
||||
"""Send and wait for completion. For LLM agent tool calls."""
|
||||
await self.send(prompt, chat_id)
|
||||
await self._busy_event.wait()
|
||||
return self._last_result or self._error or "(no output)"
|
||||
|
||||
def get_progress(self) -> SessionProgress:
|
||||
"""Return a progress snapshot. Primary query interface for the secretary AI."""
|
||||
return SessionProgress(
|
||||
busy=self._busy,
|
||||
current_prompt=self._current_prompt,
|
||||
started_at=self._started_at,
|
||||
elapsed_seconds=time.time() - self._started_at if self._busy else 0,
|
||||
text_messages=list(self._text_buffer[-5:]),
|
||||
tool_calls=list(self._tool_buffer[-10:]),
|
||||
last_result=self._last_result[:1000],
|
||||
error=self._error,
|
||||
pending_approval=self._pending_approval_desc,
|
||||
)
|
||||
|
||||
async def interrupt(self) -> None:
|
||||
"""Interrupt the currently running task."""
|
||||
if self.client and self._busy:
|
||||
await self.client.interrupt()
|
||||
logger.info("SDKSession %s interrupted", self.conv_id)
|
||||
|
||||
async def set_permission_mode(self, mode: PermissionMode) -> None:
|
||||
"""Dynamically change the permission mode."""
|
||||
if self.client:
|
||||
await self.client.set_permission_mode(mode)
|
||||
self.permission_mode = mode
|
||||
logger.info("SDKSession %s permission_mode → %s", self.conv_id, mode)
|
||||
|
||||
async def approve(self, approved: bool) -> None:
|
||||
"""Resolve a pending tool approval."""
|
||||
if self._pending_approval and not self._pending_approval.done():
|
||||
self._pending_approval.set_result(approved)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Disconnect and clean up."""
|
||||
if self._message_loop_task and not self._message_loop_task.done():
|
||||
self._message_loop_task.cancel()
|
||||
try:
|
||||
await self._message_loop_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
if self.client:
|
||||
await self.client.disconnect()
|
||||
self.client = None
|
||||
logger.info("SDKSession %s closed", self.conv_id)
|
||||
|
||||
# --- Internal ---
|
||||
|
||||
async def _message_loop(self) -> None:
|
||||
"""Background message consumption loop. Buffers messages, notifies on key events."""
|
||||
from agent.audit import log_interaction
|
||||
|
||||
try:
|
||||
async for msg in self.client.receive_messages():
|
||||
if isinstance(msg, SystemMessage) and msg.subtype == "init":
|
||||
self.session_id = msg.data.get("session_id")
|
||||
|
||||
elif isinstance(msg, AssistantMessage):
|
||||
for block in msg.content:
|
||||
if isinstance(block, TextBlock):
|
||||
self._text_buffer.append(block.text)
|
||||
if len(self._text_buffer) > self.MAX_BUFFER_TEXTS:
|
||||
self._text_buffer.pop(0)
|
||||
elif isinstance(block, ToolUseBlock):
|
||||
summary = f"{block.name}({self._summarize_input(block.input)})"
|
||||
self._tool_buffer.append(summary)
|
||||
if len(self._tool_buffer) > self.MAX_BUFFER_TOOLS:
|
||||
self._tool_buffer.pop(0)
|
||||
|
||||
elif isinstance(msg, ResultMessage):
|
||||
self._last_result = msg.result or ""
|
||||
self._busy = False
|
||||
self._busy_event.set()
|
||||
|
||||
# Key event: task completed → notify Feishu
|
||||
if self.chat_id:
|
||||
await self._notify_completion()
|
||||
|
||||
log_interaction(
|
||||
conv_id=self.conv_id,
|
||||
prompt=self._current_prompt,
|
||||
response=self._last_result[:2000],
|
||||
cwd=self.cwd,
|
||||
user_id=self.owner_id,
|
||||
)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Message loop cancelled for %s", self.conv_id)
|
||||
except Exception as exc:
|
||||
logger.exception("Message loop error for %s", self.conv_id)
|
||||
self._error = str(exc)
|
||||
self._busy = False
|
||||
self._busy_event.set()
|
||||
if self.chat_id:
|
||||
await self._notify_error(str(exc))
|
||||
|
||||
async def _notify_completion(self) -> None:
|
||||
from bot.feishu import send_markdown
|
||||
|
||||
result_preview = self._last_result[:800]
|
||||
if len(self._last_result) > 800:
|
||||
result_preview += "\n...[truncated]"
|
||||
elapsed = int(time.time() - self._started_at)
|
||||
tools_used = len(self._tool_buffer)
|
||||
msg = f"✅ **任务完成** ({elapsed}s, {tools_used} tool calls)\n\n{result_preview}"
|
||||
try:
|
||||
await send_markdown(self.chat_id, "chat_id", msg)
|
||||
except Exception:
|
||||
logger.exception("Failed to notify completion")
|
||||
|
||||
async def _notify_error(self, error: str) -> None:
|
||||
from bot.feishu import send_markdown
|
||||
|
||||
try:
|
||||
await send_markdown(
|
||||
self.chat_id, "chat_id",
|
||||
f"❌ **任务出错**\n\n```\n{error[:500]}\n```",
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to notify error")
|
||||
|
||||
async def _permission_callback(
|
||||
self, tool_name: str, input_data: dict, context: ToolPermissionContext
|
||||
) -> PermissionResult:
|
||||
"""can_use_tool — send approval card to Feishu, wait for card callback."""
|
||||
# Auto-allow read-only tools
|
||||
if tool_name in ("Read", "Glob", "Grep", "WebSearch", "WebFetch"):
|
||||
return PermissionResultAllow()
|
||||
|
||||
if not self.chat_id:
|
||||
return PermissionResultAllow()
|
||||
|
||||
# Send approval card
|
||||
from bot.feishu import send_card, build_approval_card
|
||||
|
||||
summary = self._format_tool_summary(tool_name, input_data)
|
||||
self._pending_approval_desc = f"{tool_name}: {summary}"
|
||||
|
||||
card = build_approval_card(
|
||||
conv_id=self.conv_id,
|
||||
tool_name=tool_name,
|
||||
summary=summary,
|
||||
timeout=APPROVAL_TIMEOUT,
|
||||
)
|
||||
await send_card(self.chat_id, "chat_id", card)
|
||||
|
||||
# Wait for card callback or text reply y/n
|
||||
loop = asyncio.get_running_loop()
|
||||
self._pending_approval = loop.create_future()
|
||||
try:
|
||||
approved = await asyncio.wait_for(
|
||||
self._pending_approval, timeout=APPROVAL_TIMEOUT
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
approved = False
|
||||
from bot.feishu import send_markdown
|
||||
|
||||
await send_markdown(self.chat_id, "chat_id", "⏰ 审批超时，已自动拒绝。")
|
||||
finally:
|
||||
self._pending_approval_desc = ""
|
||||
|
||||
from agent.audit import log_permission_decision
|
||||
|
||||
log_permission_decision(
|
||||
conv_id=self.conv_id,
|
||||
tool_name=tool_name,
|
||||
tool_input=input_data,
|
||||
approved=approved,
|
||||
)
|
||||
if approved:
|
||||
return PermissionResultAllow()
|
||||
return PermissionResultDeny(message="用户拒绝了此操作")
|
||||
|
||||
def _format_tool_summary(self, tool_name: str, input_data: dict) -> str:
|
||||
if tool_name == "Bash":
|
||||
return f"`{input_data.get('command', '')[:200]}`"
|
||||
if tool_name in ("Edit", "Write", "MultiEdit"):
|
||||
return f"file: `{input_data.get('file_path', input_data.get('path', ''))}`"
|
||||
return str(input_data)[:200]
|
||||
|
||||
@staticmethod
|
||||
def _summarize_input(input_data: dict) -> str:
|
||||
if "command" in input_data:
|
||||
return input_data["command"][:80]
|
||||
if "file_path" in input_data:
|
||||
return input_data["file_path"]
|
||||
return str(input_data)[:60]
|
||||
|
||||
def _build_env(self) -> dict[str, str]:
|
||||
import os
|
||||
|
||||
env = {}
|
||||
for key in ("ANTHROPIC_BASE_URL", "ANTHROPIC_AUTH_TOKEN", "CLAUDE_CODE_OAUTH_TOKEN"):
|
||||
val = os.environ.get(key, "")
|
||||
if val:
|
||||
env[key] = val
|
||||
return env
|
||||
Reference in New Issue
Block a user