"""TerminalSession - 终端会话状态管理 维护 cwd、env、history，支持跨命令保持状态。 通过在命令前注入 cd 和 export 语句实现跨命令状态持久化。 """ from __future__ import annotations import asyncio import logging import os import re import shlex import time from collections import deque from dataclasses import dataclass, field from agentkit.tools.output_parser import OutputParser, ParsedOutput logger = logging.getLogger(__name__) _ENV_KEY_PATTERN = re.compile(r'^[A-Za-z_][A-Za-z0-9_]*$') @dataclass class CommandRecord: """命令执行记录 Attributes: command: 执行的命令 exit_code: 退出码 output: 标准输出+错误输出 cwd: 执行时的工作目录 timestamp: 执行时间戳 duration_ms: 执行耗时（毫秒） """ command: str exit_code: int output: str cwd: str timestamp: float duration_ms: int class TerminalSession: """终端会话 - 跨命令保持 cwd/env/history 状态 通过在命令前注入 `cd {cwd} && ` 和 `export K=V && ` 实现跨命令状态持久化。 每次命令执行后自动更新 cwd 和 env 状态。 Usage: session = TerminalSession(session_id="build-01") result = await session.execute("cd /tmp") result = await session.execute("pwd") # 输出 /tmp """ def __init__( self, session_id: str, cwd: str | None = None, env: dict[str, str] | None = None, max_history: int = 1000, ): self.session_id = session_id self._cwd = cwd or os.getcwd() self._env: dict[str, str] = dict(env or os.environ) self._env_delta: set[str] = set() # Track only env vars explicitly set in this session self._history: deque[CommandRecord] = deque(maxlen=max_history) self._output_parser = OutputParser() self._created_at = time.time() @property def cwd(self) -> str: """当前工作目录""" return self._cwd @property def env(self) -> dict[str, str]: """当前环境变量（副本）""" return dict(self._env) @property def history(self) -> list[CommandRecord]: """命令执行历史（副本）""" return list(self._history) @property def created_at(self) -> float: """会话创建时间戳""" return self._created_at def get_cwd(self) -> str: """获取当前工作目录""" return self._cwd def set_cwd(self, cwd: str) -> None: """手动设置当前工作目录""" self._cwd = cwd def get_env(self) -> dict[str, str]: """获取当前环境变量（副本）""" return dict(self._env) def set_env(self, key: str, value: str) -> None: """设置单个环境变量""" self._env[key] = value def update_env(self, env: dict[str, str]) -> None: """批量更新环境变量""" self._env.update(env) def get_history(self) -> list[CommandRecord]: """获取命令执行历史（副本）""" return list(self._history) async def execute( self, command: str, timeout: float | None = None, ) -> ParsedOutput: """在会话上下文中执行命令 自动在命令前注入 cd 和 export 语句以保持会话状态。 执行后自动更新 cwd 和 env。 Args: command: 要执行的命令 timeout: 超时时间（秒），None 表示不超时 Returns: ParsedOutput 结构化解析结果 """ full_command = self._build_command(command) start = time.monotonic() try: proc = await asyncio.create_subprocess_shell( full_command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, env=self._env, ) try: stdout, _ = await asyncio.wait_for( proc.communicate(), timeout=timeout, ) except asyncio.TimeoutError: try: proc.kill() except ProcessLookupError: pass # Process already exited except OSError: pass await proc.wait() output = f"命令执行超时（{timeout}s）" exit_code = -1 else: output = stdout.decode("utf-8", errors="replace") if stdout else "" exit_code = proc.returncode if proc.returncode is not None else 0 except Exception as e: output = str(e) exit_code = -1 duration_ms = int((time.monotonic() - start) * 1000) # 更新会话状态 self._update_state_after_execution(command, output, exit_code) # 记录历史 record = CommandRecord( command=command, exit_code=exit_code, output=output, cwd=self._cwd, timestamp=time.time(), duration_ms=duration_ms, ) self._add_history(record) # 解析输出 parsed = self._output_parser.parse(output, exit_code) logger.debug( "Session %s: command=%r exit_code=%d duration=%dms", self.session_id, command, exit_code, duration_ms, ) return parsed def _build_command(self, command: str) -> str: """构建带会话状态的完整命令 在原始命令前注入 cd 和 export 语句。 """ parts: list[str] = [] # 注入 cd if self._cwd: parts.append(f"cd {shlex.quote(self._cwd)}") # 注入环境变量（仅注入 session 中显式设置的增量，避免泄露完整 os.environ） if self._env_delta: for key in self._env_delta: value = self._env.get(key, "") if _ENV_KEY_PATTERN.match(key): parts.append(f"export {shlex.quote(key)}={shlex.quote(value)}") parts.append(command) return " && ".join(parts) def _update_state_after_execution( self, command: str, output: str, exit_code: int ) -> None: """命令执行后更新会话状态 解析 cd 和 export 命令更新 cwd 和 env。 """ if exit_code != 0: return # 解析 cd 命令更新 cwd self._parse_cd_commands(command, output) # 解析 export 命令更新 env self._parse_export_commands(command) def _parse_cd_commands(self, command: str, output: str) -> None: """从命令中解析 cd 并更新 cwd 支持: - cd /path - cd dir (相对路径，需要通过 pwd 获取实际路径) - cd - (切换到上一个目录) """ import re # 匹配 cd 命令（可能出现在 && 链中） cd_pattern = re.compile(r"(?:^|\s|&&\s*)cd\s+(.+?)(?:\s*&&|\s*$)") matches = cd_pattern.findall(command) for target in matches: target = target.strip().strip("'\"") if not target: continue if target == "-": # cd - 切换到 OLDPWD old_pwd = self._env.get("OLDPWD") if old_pwd: self._cwd = old_pwd elif os.path.isabs(target): self._cwd = target else: # 相对路径：拼接后规范化 new_cwd = os.path.normpath(os.path.join(self._cwd, target)) self._cwd = new_cwd def _parse_export_commands(self, command: str) -> None: """从命令中解析 export 并更新 env 支持: - export KEY=VALUE - export KEY="VALUE WITH SPACES" """ import re export_pattern = re.compile( r"(?:^|\s|&&\s*)export\s+(\w+)=(.+?)(?:\s*&&|\s*$)" ) matches = export_pattern.findall(command) for key, value in matches: value = value.strip().strip("'\"") self._env[key] = value self._env_delta.add(key) def _add_history(self, record: CommandRecord) -> None: """添加命令记录到历史，deque maxlen 自动淘汰最旧记录""" self._history.append(record) def close(self) -> None: """关闭会话，清理资源""" logger.info( "Session %s closed: %d commands executed", self.session_id, len(self._history), ) class TerminalSessionManager: """终端会话管理器 - 按 ID 管理多个 TerminalSession Usage: manager = TerminalSessionManager() session = manager.get_or_create("build-01") session = manager.get("build-01") manager.remove("build-01") """ def __init__(self, max_sessions: int = 100): self._sessions: dict[str, TerminalSession] = {} self._max_sessions = max_sessions def get_or_create( self, session_id: str, cwd: str | None = None, env: dict[str, str] | None = None, ) -> TerminalSession: """获取或创建会话 Args: session_id: 会话 ID cwd: 初始工作目录（仅创建时使用） env: 初始环境变量（仅创建时使用） Returns: TerminalSession 实例 """ if session_id not in self._sessions: if len(self._sessions) >= self._max_sessions: # 移除最旧的会话 oldest_id = min( self._sessions, key=lambda k: self._sessions[k].created_at ) self.remove(oldest_id) self._sessions[session_id] = TerminalSession( session_id=session_id, cwd=cwd, env=env, ) logger.info("Session created: %s", session_id) return self._sessions[session_id] def get(self, session_id: str) -> TerminalSession | None: """获取会话，不存在返回 None""" return self._sessions.get(session_id) def remove(self, session_id: str) -> None: """移除并关闭会话""" session = self._sessions.pop(session_id, None) if session: session.close() def list_sessions(self) -> list[str]: """列出所有会话 ID""" return list(self._sessions.keys()) def has_session(self, session_id: str) -> bool: """检查会话是否存在""" return session_id in self._sessions def close_all(self) -> None: """关闭所有会话""" for session_id in list(self._sessions.keys()): self.remove(session_id)