codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/agent/background_procs.py

@@ -0,0 +1,372 @@
+"""pod 에서 끝나지 않는 명령을 띄우고 거둔다 (이슈 #130).
+
+`execute_command` 는 명령이 끝날 때까지 기다린다. 개발 서버(`python -m http.server`,
+`npm run dev`, `./mvnw spring-boot:run`)는 끝나지 않으므로 무출력 한도에 걸려 죽거나
+턴을 붙잡는다. 여기서는 그런 명령을 **백그라운드로** 띄우고, 목록·로그·종료를 따로
+제공한다.
+
+- 출력은 파이프가 아니라 **로그 파일**로 받는다. 파이프는 아무도 안 읽으면 버퍼가
+  차서 서버가 멈춘다.
+- 셸마다 프로세스 그룹을 따로 둔다(`start_new_session`). 종료는 그룹 단위라 셸이 띄운
+  자식(java·node)까지 같이 끊긴다.
+- 생존과 LISTEN 포트는 `psutil` 로 본다. 포트는 패널의 "열기" 링크가 쓴다.
+- 레지스트리는 jupyter 서버 프로세스 하나에 하나이고 **대화와 무관한 전역**이다. pod 은
+  사용자 1명 = 서버 1개라 전역이 곧 "내 것" 이다. 대화(thread_id)는 중지·새로고침마다
+  버려져서(프론트가 저장하지 않는다) 소유 단위로 쓰면 못 찾는 프로세스가 생긴다.
+  `owner` 는 표시용으로만 남긴다.
+- 죽이는 때는 셋이다: `kill_process` 도구·패널 중지 버튼, 서버 종료(`atexit`),
+  서버가 `kill -9` 로 죽어 `atexit` 이 안 돈 경우 다음 기동의 `reap_orphans`.
+  턴 취소·모델 변경·유휴 회수·대화 리셋에서는 죽이지 않는다 — 사용자가 브라우저에서
+  쓰는 서버가 그 사이에 사라지면 기능이 무의미하다.
+- **생존은 셸이 아니라 그룹으로 본다.** 셸(그룹 리더)이 먼저 끝나도 `cmd &`·`nohup` 이
+  남긴 자식이 그룹에 살아 포트를 물고 있을 수 있다. 커널은 그룹에 멤버가 있는 동안 그
+  번호를 새 프로세스에 주지 않으므로, 리더가 없는데 그룹이 살아 있으면 그 그룹은 우리
+  것이다. 반대로 리더를 회수했는데 같은 pid 의 프로세스가 보이면 번호가 재사용된 남의
+  것이라 신호를 보내지 않는다. `setsid`·데몬화로 그룹을 벗어난 프로세스는 못 잡는다.
+- 레지스트리 파일은 서버가 아니라 `CODEBUFF_HOME` 단위다. 같은 홈을 쓰는 서버를 둘 띄우면
+  나중 서버의 `reap_orphans` 가 앞 서버의 프로세스를 거둔다 — pod 은 서버가 하나라 해당
+  없고, 개발 기계의 워크트리는 홈이 갈린다.
+"""
+
+from __future__ import annotations
+
+import atexit
+import itertools
+import json
+import logging
+import os
+import shutil
+import signal
+import subprocess
+import tempfile
+import threading
+import time
+from typing import Any, Dict, List, Optional
+
+import psutil
+
+logger = logging.getLogger(__name__)
+
+# 종료 신호를 보낸 뒤 SIGKILL 까지 기다리는 시간(초)
+_KILL_GRACE_S = 3.0
+# 로그를 돌려줄 때의 줄 수 기본값과 상한
+_TAIL_DEFAULT = 200
+_TAIL_MAX = 2000
+# 로그 끝부분만 읽는다 — 서버 로그는 하루면 수백 MB 가 된다
+_TAIL_READ_BYTES = 256 * 1024
+# 끝난 프로세스를 목록에 남기는 개수. 넘으면 오래된 것부터 로그와 함께 지운다
+_MAX_FINISHED = 20
+
+
+class _Proc:
+    __slots__ = (
+        "id",
+        "popen",
+        "command",
+        "cwd",
+        "owner",
+        "log",
+        "started_at",
+        "create_time",
+    )
+
+    def __init__(
+        self, id_: str, popen: Any, command: str, cwd: str, owner: str, log: str
+    ):
+        self.id = id_
+        self.popen = popen
+        self.command = command
+        self.cwd = cwd
+        self.owner = owner
+        self.log = log
+        self.started_at = time.time()
+        # pid 재사용을 가르는 값 — 재기동 뒤 `reap_orphans` 가 남의 프로세스를 죽이지 않게
+        try:
+            self.create_time: Optional[float] = psutil.Process(popen.pid).create_time()
+        except (psutil.NoSuchProcess, psutil.AccessDenied):
+            self.create_time = None
+
+
+_procs: Dict[str, _Proc] = {}
+_lock = threading.Lock()
+_ids = itertools.count(1)
+_log_dir: Optional[str] = None
+
+
+def _registry_path() -> str:
+    home = os.environ.get("CODEBUFF_HOME") or os.path.expanduser("~/.codebuff")
+    return os.path.join(home, "background_procs.json")
+
+
+def _group_alive(pgid: int) -> bool:
+    """그룹에 프로세스가 하나라도 남아 있는가 (신호 0 은 보내지 않고 존재만 묻는다)."""
+    try:
+        os.killpg(pgid, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    return True
+
+
+def _owns_group(proc: _Proc) -> bool:
+    """그 프로세스 그룹이 아직 우리 것이고 살아 있는가."""
+    if proc.popen.poll() is None:
+        return True
+    pid = proc.popen.pid
+    try:
+        psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        # 리더는 회수됐고 번호도 비어 있다 — 그룹이 살아 있으면 남은 자식이다
+        return _group_alive(pid)
+    except psutil.AccessDenied:
+        return False
+    # 리더를 이미 회수했는데 같은 번호가 보인다 — 재사용된 남의 프로세스다
+    return False
+
+
+def _save() -> None:
+    """살아 있는 프로세스를 파일에 적는다. 서버가 `kill -9` 로 죽으면 `atexit` 이
+    안 돌고, 자식은 세션을 따로 두어 살아남는다 — 다음 기동이 이 파일로 거둔다."""
+    rows = [
+        {
+            "id": p.id,
+            "pid": p.popen.pid,
+            "create_time": p.create_time,
+            "command": p.command,
+        }
+        for p in _procs.values()
+        if p.create_time is not None and _owns_group(p)
+    ]
+    path = _registry_path()
+    try:
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        tmp = path + ".tmp"
+        with open(tmp, "w", encoding="utf-8") as f:
+            json.dump(rows, f, ensure_ascii=False)
+        os.replace(tmp, path)
+    except OSError as e:
+        # 파일이 없으면 서버가 `kill -9` 로 죽었을 때 다음 기동이 못 거둔다
+        logger.warning("[Background] 레지스트리를 못 썼습니다 (%s): %s", path, e)
+
+
+def reap_orphans() -> int:
+    """지난 서버가 남긴 백그라운드 프로세스를 거둔다. 확장 로드 때 한 번 부른다.
+
+    pid 만 보면 재사용된 남의 프로세스를 죽일 수 있어 생성 시각까지 맞는 것만 죽인다.
+    리더가 없으면 그룹이 살아 있을 때만 죽인다 — 멤버가 있는 동안 그 번호는 재사용되지
+    않는다.
+    """
+    try:
+        with open(_registry_path(), encoding="utf-8") as f:
+            rows = json.load(f)
+    except (OSError, ValueError):
+        return 0
+    reaped = 0
+    for row in rows if isinstance(rows, list) else []:
+        try:
+            pid = int(row["pid"])
+            created = float(row["create_time"])
+        except (KeyError, TypeError, ValueError):
+            continue
+        try:
+            ours = abs(psutil.Process(pid).create_time() - created) < 1.0
+        except psutil.NoSuchProcess:
+            ours = _group_alive(pid)
+        except psutil.Error:
+            continue
+        if ours and _signal_group(pid, signal.SIGKILL):
+            reaped += 1
+    with _lock:
+        _save()
+    return reaped
+
+
+def _logs() -> str:
+    global _log_dir
+    if _log_dir is None or not os.path.isdir(_log_dir):
+        _log_dir = tempfile.mkdtemp(prefix="codebuff-bg-")
+    return _log_dir
+
+
+def start(
+    command: str,
+    cwd: str,
+    env: Optional[Dict[str, str]] = None,
+    owner: str = "",
+) -> Dict[str, Any]:
+    """명령을 백그라운드로 띄우고 곧바로 돌아온다."""
+    if not command:
+        return {"success": False, "error": "명령이 비어 있습니다."}
+    id_ = f"bg{next(_ids)}"
+    log = os.path.join(_logs(), f"{id_}.log")
+    with open(log, "wb") as out:
+        popen = subprocess.Popen(
+            command,
+            shell=True,
+            cwd=cwd if os.path.isdir(cwd) else None,
+            stdin=subprocess.DEVNULL,
+            stdout=out,
+            stderr=subprocess.STDOUT,
+            start_new_session=True,
+            env=env,
+        )
+    proc = _Proc(id_, popen, command, cwd, owner, log)
+    with _lock:
+        _procs[id_] = proc
+        _trim_finished()
+        _save()
+    return {"success": True, "process_id": id_, **snapshot(proc)}
+
+
+def _family(pid: int) -> List[Any]:
+    """리더와 자손. 리더가 끝났으면 같은 그룹에 남은 프로세스."""
+    try:
+        root = psutil.Process(pid)
+        return [root] + root.children(recursive=True)
+    except psutil.AccessDenied:
+        return []
+    except psutil.NoSuchProcess:
+        pass
+    members = []
+    for p in psutil.process_iter():
+        try:
+            if os.getpgid(p.pid) == pid:
+                members.append(p)
+        except (OSError, psutil.Error):
+            continue
+    return members
+
+
+def _ports(pid: int) -> List[int]:
+    """그 그룹이 LISTEN 중인 TCP 포트."""
+    ports = set()
+    for p in _family(pid):
+        try:
+            # psutil 6 에서 `connections` 가 `net_connections` 로 바뀌었다
+            conns = (getattr(p, "net_connections", None) or p.connections)(kind="tcp")
+        except (psutil.NoSuchProcess, psutil.AccessDenied):
+            continue
+        for c in conns:
+            if c.status == psutil.CONN_LISTEN and c.laddr:
+                ports.add(c.laddr.port)
+    return sorted(ports)
+
+
+def snapshot(proc: _Proc) -> Dict[str, Any]:
+    code = proc.popen.poll()
+    running = _owns_group(proc)
+    return {
+        "id": proc.id,
+        "pid": proc.popen.pid,
+        "command": proc.command,
+        "cwd": proc.cwd,
+        "owner": proc.owner,
+        "state": "running" if running else "exited",
+        "exit_code": None if running else code,
+        "ports": _ports(proc.popen.pid) if running else [],
+        "started_at": proc.started_at,
+        "elapsed_s": round(time.time() - proc.started_at, 1),
+    }
+
+
+def list_processes() -> Dict[str, Any]:
+    with _lock:
+        procs = list(_procs.values())
+    return {"success": True, "processes": [snapshot(p) for p in procs]}
+
+
+def read_output(process_id: str, tail: Any = None) -> Dict[str, Any]:
+    proc = _procs.get(process_id)
+    if proc is None:
+        return {"success": False, "error": f"없는 프로세스: {process_id}"}
+    try:
+        lines = int(tail or _TAIL_DEFAULT)
+    except (TypeError, ValueError):
+        lines = _TAIL_DEFAULT
+    lines = max(1, min(lines, _TAIL_MAX))
+    try:
+        with open(proc.log, "rb") as f:
+            f.seek(0, os.SEEK_END)
+            size = f.tell()
+            f.seek(max(0, size - _TAIL_READ_BYTES))
+            text = f.read().decode("utf-8", errors="replace")
+    except OSError as e:
+        return {"success": False, "error": f"로그를 못 읽었습니다: {e}"}
+    kept = text.splitlines()[-lines:]
+    return {"success": True, **snapshot(proc), "output": "\n".join(kept)}
+
+
+def _signal_group(pid: int, sig: int) -> bool:
+    try:
+        os.killpg(pid, sig)
+    except (ProcessLookupError, PermissionError):
+        return False
+    return True
+
+
+def _wait_group(proc: _Proc, timeout: float) -> bool:
+    """그룹이 빌 때까지 기다린다. 리더는 회수해야 좀비로 그룹에 남지 않는다."""
+    deadline = time.time() + timeout
+    while True:
+        try:
+            proc.popen.wait(timeout=0.1)
+        except subprocess.TimeoutExpired:
+            pass
+        if not _group_alive(proc.popen.pid):
+            return True
+        if time.time() >= deadline:
+            return False
+
+
+def _stop(proc: _Proc) -> None:
+    """그룹 전체에 SIGTERM, 유예 뒤 남았으면 SIGKILL.
+
+    셸(그룹 리더)이 먼저 끝나도 그룹에 자식이 남을 수 있어 리더가 아니라 그룹이 빌
+    때까지 기다린다. 이미 끝났거나 번호가 재사용됐으면 아무것도 보내지 않는다.
+    """
+    if not _owns_group(proc):
+        return
+    pid = proc.popen.pid
+    _signal_group(pid, signal.SIGTERM)
+    if not _wait_group(proc, _KILL_GRACE_S):
+        _signal_group(pid, signal.SIGKILL)
+        _wait_group(proc, _KILL_GRACE_S)
+
+
+def kill(process_id: str) -> Dict[str, Any]:
+    proc = _procs.get(process_id)
+    if proc is None:
+        return {"success": False, "error": f"없는 프로세스: {process_id}"}
+    _stop(proc)
+    with _lock:
+        _save()
+    return {"success": True, **snapshot(proc)}
+
+
+def kill_all() -> None:
+    with _lock:
+        procs = list(_procs.values())
+    if not procs:
+        # 띄운 적이 없으면 파일도 건드리지 않는다 (`atexit` 은 모든 서버에서 돈다)
+        return
+    for p in procs:
+        _stop(p)
+    with _lock:
+        _save()
+    if _log_dir is not None:
+        shutil.rmtree(_log_dir, ignore_errors=True)
+
+
+def _trim_finished() -> None:
+    finished = [p for p in _procs.values() if not _owns_group(p)]
+    for p in finished[: max(0, len(finished) - _MAX_FINISHED)]:
+        _procs.pop(p.id, None)
+        try:
+            os.unlink(p.log)
+        except OSError:
+            pass
+
+
+# jupyter 서버가 내려갈 때 띄운 서버가 남지 않게 한다
+atexit.register(kill_all)

jupyter_ext/agent/code_history_middleware.py

@@ -0,0 +1,211 @@
+"""CodeHistoryTracker + CodeHistoryMiddleware.
+
+코드 실행 이력을 추적하고, 서브에이전트(python_expert)에 이전 코드 컨텍스트를 주입.
+
+데이터 흐름:
+    HITL edit decision (execution_result)
+        -> http_handlers.py: track_tool_execution()
+        -> CodeHistoryTracker: 이력 저장
+        -> Main Agent -> task(python_expert)
+        -> CodeHistoryMiddleware.awrap_tool_call: description에 코드 히스토리 append
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from langchain.agents.middleware.types import AgentMiddleware
+
+logger = logging.getLogger(__name__)
+
+# 토큰 추정: 1 토큰 ~= 4 chars
+_CHARS_PER_TOKEN = 4
+_MAX_CONTEXT_TOKENS = 50000
+_MAX_CONTEXT_CHARS = _MAX_CONTEXT_TOKENS * _CHARS_PER_TOKEN
+_RECENT_FULL_COUNT = 3  # 최근 N개는 전체 코드+출력 유지
+
+
+@dataclass
+class CodeHistoryEntry:
+    """단일 코드 실행 이력."""
+
+    tool_name: str  # jupyter_cell, file_write
+    code: Optional[str] = None
+    output: Optional[str] = None
+    file_path: Optional[str] = None
+
+    def to_context_string(self) -> str:
+        """전체 코드+출력 포맷."""
+        parts = [f"[{self.tool_name}]"]
+        if self.file_path:
+            parts.append(f"File: {self.file_path}")
+        if self.code:
+            parts.append(f"```python\n{self.code}\n```")
+        if self.output:
+            parts.append(f"Output:\n{self.output[:2000]}")
+        return "\n".join(parts)
+
+    def to_summary_string(self) -> str:
+        """한 줄 요약."""
+        if self.tool_name == "file_write":
+            return f"[file_write] {self.file_path or 'unknown'}"
+        code_preview = (self.code or "")[:80].replace("\n", " ")
+        output_preview = (self.output or "")[:60].replace("\n", " ")
+        return f"[{self.tool_name}] {code_preview} -> {output_preview}"
+
+
+class CodeHistoryTracker:
+    """스레드별 코드 실행 이력 관리."""
+
+    def __init__(self) -> None:
+        self._history: list[CodeHistoryEntry] = []
+
+    def add_jupyter_cell(self, code: str, output: Optional[str] = None) -> None:
+        self._history.append(
+            CodeHistoryEntry(tool_name="jupyter_cell", code=code, output=output)
+        )
+
+    def add_write_file(self, file_path: str, content: Optional[str] = None) -> None:
+        self._history.append(
+            CodeHistoryEntry(tool_name="file_write", code=content, file_path=file_path)
+        )
+
+    @property
+    def entries(self) -> list[CodeHistoryEntry]:
+        return list(self._history)
+
+    def get_context_for_subagent(
+        self,
+        existing_context: str = "",
+        max_tokens: int = _MAX_CONTEXT_TOKENS,
+    ) -> str:
+        """서브에이전트에 주입할 코드 히스토리 컨텍스트 생성.
+
+        최근 _RECENT_FULL_COUNT개는 전체 코드+출력,
+        그 이전은 한 줄 요약으로 토큰 관리.
+        """
+        if not self._history:
+            return existing_context
+
+        max_chars = max_tokens * _CHARS_PER_TOKEN
+        existing_chars = len(existing_context)
+        budget = max_chars - existing_chars
+        if budget <= 200:
+            return existing_context
+
+        history_parts: list[str] = []
+        total = len(self._history)
+
+        # 오래된 항목: 요약
+        if total > _RECENT_FULL_COUNT:
+            for entry in self._history[: total - _RECENT_FULL_COUNT]:
+                history_parts.append(entry.to_summary_string())
+
+        # 최근 항목: 전체
+        recent_start = max(0, total - _RECENT_FULL_COUNT)
+        for entry in self._history[recent_start:]:
+            history_parts.append(entry.to_context_string())
+
+        history_text = "\n\n".join(history_parts)
+
+        # 예산 초과 시 잘라내기
+        if len(history_text) > budget:
+            history_text = history_text[-(budget - 50) :]
+            history_text = "...(truncated)\n" + history_text
+
+        section = f"\n\n## Previous Code Execution History\n{history_text}\n"
+
+        if existing_context:
+            return existing_context + section
+        return section.strip()
+
+
+# ---------------------------------------------------------------------------
+# Global tracker registry (per thread_id)
+# ---------------------------------------------------------------------------
+
+_code_history_trackers: dict[str, CodeHistoryTracker] = {}
+
+
+def get_code_history_tracker(thread_id: str) -> CodeHistoryTracker:
+    """thread_id별 CodeHistoryTracker 반환 (없으면 새로 생성)."""
+    if thread_id not in _code_history_trackers:
+        _code_history_trackers[thread_id] = CodeHistoryTracker()
+    return _code_history_trackers[thread_id]
+
+
+def track_tool_execution(
+    tool_name: str, args: dict, thread_id: str
+) -> None:
+    """HITL edit decision에서 호출: 코드 실행 이력 기록."""
+    tracker = get_code_history_tracker(thread_id)
+
+    if tool_name == "jupyter_cell":
+        code = args.get("code", "")
+        exec_result = args.get("execution_result", {})
+        output = None
+        if isinstance(exec_result, dict):
+            output = exec_result.get("output", exec_result.get("text", ""))
+        tracker.add_jupyter_cell(code, output)
+        logger.debug(
+            "CodeHistory: tracked jupyter_cell (code_len=%d, has_output=%s)",
+            len(code),
+            bool(output),
+        )
+    elif tool_name == "file_write":
+        file_path = args.get("path", "")
+        content = args.get("content", "")
+        tracker.add_write_file(file_path, content)
+        logger.debug("CodeHistory: tracked file_write (path=%s)", file_path)
+
+
+# ---------------------------------------------------------------------------
+# CodeHistoryMiddleware
+# ---------------------------------------------------------------------------
+
+
+class CodeHistoryMiddleware(AgentMiddleware):
+    """task 도구 호출 인터셉트하여 python_expert에 코드 히스토리 주입."""
+
+    async def awrap_tool_call(self, request, handler):
+        tool_call = request.tool_call
+        tool_name = tool_call.get("name", "")
+
+        if tool_name == "task":
+            args = tool_call.get("args", {})
+            subagent_type = args.get("subagent_type", args.get("agent_name", ""))
+
+            if subagent_type in ("python_expert", "python_developer"):
+                # thread_id 추출 시도
+                thread_id = self._extract_thread_id(request)
+                if thread_id:
+                    tracker = get_code_history_tracker(thread_id)
+                    if tracker.entries:
+                        description = args.get("description", "")
+                        enhanced = tracker.get_context_for_subagent(description)
+                        new_args = {**args, "description": enhanced}
+                        new_tool_call = {**tool_call, "args": new_args}
+                        request = request.override(tool_call=new_tool_call)
+                        logger.debug(
+                            "CodeHistory: injected %d entries into %s task",
+                            len(tracker.entries),
+                            subagent_type,
+                        )
+
+        return await handler(request)
+
+    def _extract_thread_id(self, request: Any) -> str:
+        """request에서 thread_id 추출."""
+        # runtime config에서 추출 시도
+        if hasattr(request, "config"):
+            config = request.config
+            if isinstance(config, dict):
+                return config.get("configurable", {}).get("thread_id", "")
+        # state에서 추출 시도
+        if hasattr(request, "state"):
+            state = request.state
+            if isinstance(state, dict):
+                return state.get("thread_id", "")
+        return ""