360 lines
11 KiB
Python
360 lines
11 KiB
Python
"""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)
|