codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/agent/codex_threads.py

@@ -0,0 +1,114 @@
+"""이 사용자가 연 codex 스레드 목록 (이슈 #121 — `/resume`).
+
+codex 의 `thread/list` 는 **서버의 모든 스레드**를 준다. centri 처럼 여러 사용자가 한
+app-server 를 쓰면 남의 스레드가 섞이고, `originators`·`sourceKinds` 로도 사용자를
+가를 수 없다(스키마 확인). 그래서 목록의 정본은 codex 가 아니라 이 파일이다 — 이 pod
+(=이 사용자)의 세션이 `thread/start`·`thread/fork` 로 만든 id 만 적는다. `/resume` 은
+여기 있는 id 만 받는다.
+
+쓰기는 매번 다시 읽어 합치고 `os.replace` 로 바꾼다. 같은 사용자의 탭 두 개가 거의
+동시에 스레드를 열어도 한쪽 기록이 통째로 사라지지 않게 하려는 것이고, 락은 걸지 않는다
+(같은 순간의 두 쓰기 중 하나를 잃는 것은 목록 한 줄이라 감수한다).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+#: 테스트와 운영이 다른 곳을 쓰게 하는 환경 변수. 없으면 `$CODEBUFF_HOME` (워크트리
+#: 격리, 이슈 #116), 그것도 없으면 `~/.codebuff` 아래 `codex_threads.json`.
+PATH_ENV = "CODEBUFF_CODEX_THREADS"
+
+#: 오래된 것부터 버린다. `/resume` 목록은 이 중 앞 `LIST_LIMIT` 개만 codex 에 묻는다.
+MAX_THREADS = 50
+LIST_LIMIT = 20
+
+
+def registry_path() -> str:
+    home = os.environ.get("CODEBUFF_HOME") or os.path.expanduser("~/.codebuff")
+    return os.environ.get(PATH_ENV) or os.path.join(home, "codex_threads.json")
+
+
+def load() -> List[Dict[str, Any]]:
+    """새것부터. 파일이 없거나 깨졌으면 빈 목록이다."""
+    try:
+        with open(registry_path(), encoding="utf-8") as f:
+            data = json.load(f)
+    except (OSError, ValueError):
+        return []
+    threads = data.get("threads") if isinstance(data, dict) else None
+    return [t for t in threads or [] if isinstance(t, dict) and t.get("id")]
+
+
+def _save(threads: List[Dict[str, Any]]) -> None:
+    path = registry_path()
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    fd, tmp = tempfile.mkstemp(dir=os.path.dirname(path), prefix=".codex_threads.")
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as f:
+            json.dump(
+                {"version": 1, "threads": threads[:MAX_THREADS]},
+                f,
+                ensure_ascii=False,
+                indent=2,
+            )
+        os.replace(tmp, path)
+    except BaseException:
+        try:
+            os.unlink(tmp)
+        except OSError:
+            pass
+        raise
+
+
+def record(
+    thread: Dict[str, Any], *, kind: str, workspace: str, parent: str = ""
+) -> None:
+    """새로 연 스레드를 맨 앞에 적는다. 이미 있으면 그 줄을 앞으로 옮긴다."""
+    tid = str(thread.get("id") or "")
+    if not tid:
+        return
+    threads = [t for t in load() if t.get("id") != tid]
+    entry: Dict[str, Any] = {
+        "id": tid,
+        "kind": kind,
+        "workspace": workspace,
+        "createdAt": int(time.time()),
+    }
+    if parent:
+        entry["parent"] = parent
+    _save([entry] + threads)
+
+
+def update(thread_id: str, **fields: Any) -> None:
+    threads = load()
+    for t in threads:
+        if t.get("id") == thread_id:
+            t.update(fields)
+    _save(threads)
+
+
+def forget(thread_id: str) -> None:
+    _save([t for t in load() if t.get("id") != thread_id])
+
+
+def find(thread_id: str) -> Optional[Dict[str, Any]]:
+    return next((t for t in load() if t.get("id") == thread_id), None)
+
+
+def recorder(workspace: str) -> Callable[[str, Dict[str, Any]], None]:
+    """`CodexSession(on_thread=...)` 에 넘기는 기록기."""
+
+    def _on_thread(kind: str, thread: Dict[str, Any]) -> None:
+        record(
+            thread,
+            kind=kind,
+            workspace=workspace,
+            parent=str(thread.get("forkedFromId") or ""),
+        )
+
+    return _on_thread

jupyter_ext/agent/content_injection.py

@@ -0,0 +1,377 @@
+"""
+ContentInjectionMiddleware for multi-agent code/SQL auto-injection.
+
+Parses [CODE]/[SQL] blocks from subagent ToolMessage responses and injects
+them into target tool calls (jupyter_cell, file_write, jupyter_markdown).
+
+Ported from chat-mode-enhancement branch, adapted for DeepAgents ToolMessage
+parsing instead of State-based approach (DeepAgents has no state_schema param).
+"""
+
+import logging
+import re
+from typing import Optional
+
+from ..log_paths import get_debug_log_path
+
+logger = logging.getLogger(__name__)
+
+# Direct file-based debug logging (same as agent_factory._dbg)
+
+_DEBUG_LOG = get_debug_log_path()
+
+
+def _dbg(msg: str) -> None:
+    """Write debug message directly to agent_debug.log."""
+    try:
+        import datetime
+
+        ts = datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
+        with open(_DEBUG_LOG, "a") as f:
+            f.write(f"[{ts}] {msg}\n")
+            f.flush()
+    except OSError:
+        pass
+
+
+def _extract_code_block(response: str, lang: str = "python") -> Optional[str]:
+    """Extract code/SQL from response using [CODE]/[SQL] markers or fenced blocks.
+
+    Tries in order:
+    1. [CODE] or [SQL] section marker with fenced block
+    2. [CODE] or [SQL] marker with content (no fenced block)
+    3. Fenced code block with language
+    4. Generic fenced code block (fallback)
+    """
+    if not response:
+        return None
+
+    marker = "[CODE]" if lang == "python" else "[SQL]"
+
+    # Pattern 1: [CODE]/[SQL] marker + fenced block
+    marker_pattern = re.escape(marker) + r"\s*\n\s*```(?:\w+)?\s*\n(.*?)```"
+    match = re.search(marker_pattern, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        if code:
+            return code
+
+    # Pattern 2: [CODE]/[SQL] marker with content after it (no fenced block)
+    marker_pattern2 = re.escape(marker) + r"\s*\n(.*?)(?=\n\[|\Z)"
+    match = re.search(marker_pattern2, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        if code.startswith(f"```{lang}"):
+            code = code[len(f"```{lang}") :].strip()
+        if code.startswith("```"):
+            code = code[3:].strip()
+        if code.endswith("```"):
+            code = code[:-3].strip()
+        if code:
+            return code
+
+    # Pattern 3: Fenced code block with language
+    fenced_pattern = rf"```{lang}\s*\n(.*?)```"
+    match = re.search(fenced_pattern, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        if code:
+            return code
+
+    # Pattern 4: Generic fenced code block (fallback)
+    generic_pattern = r"```\s*\n(.*?)```"
+    match = re.search(generic_pattern, response, re.DOTALL)
+    if match:
+        code = match.group(1).strip()
+        if code and len(code) > 10:
+            return code
+
+    return None
+
+
+def _extract_description(response: str) -> Optional[str]:
+    """Extract [DESCRIPTION] section from subagent response."""
+    if not response:
+        return None
+
+    patterns = [
+        r"\[DESCRIPTION\]\s*\n(.*?)(?=\n\s*\[(?:CODE|SQL|ANALYSIS|RECOMMENDATION)\])",
+        r"\[DESCRIPTION\]\s*(.*?)(?=\s*\[(?:CODE|SQL|ANALYSIS|RECOMMENDATION)\])",
+        r"\[DESCRIPTION\]\s*\n(.*?)(?=\n\s*```)",
+        r"\[DESCRIPTION\]\s*\n(.+?)(?=\n\n|\Z)",
+    ]
+
+    for pattern in patterns:
+        match = re.search(pattern, response, re.DOTALL | re.IGNORECASE)
+        if match:
+            description = match.group(1).strip()
+            if description and len(description) > 5:
+                return description
+
+    return None
+
+
+try:
+    from langchain.agents.middleware.types import AgentMiddleware
+except ImportError:
+    # Fallback: define a minimal base class
+    class AgentMiddleware:  # type: ignore[no-redef]
+        tools: list = []
+
+
+class ContentInjectionMiddleware(AgentMiddleware):
+    """Inject subagent-generated code/SQL into target tool calls.
+
+    DeepAgents' task tool returns subagent responses as ToolMessage(text).
+    This middleware:
+    1. In before_model(): scans recent ToolMessages for [CODE]/[SQL] blocks
+    2. In after_model(): injects content into non-HITL tool calls (e.g. jupyter_markdown)
+    3. In awrap_tool_call(): fallback injection during tool execution
+
+    Note: For HITL-interrupted tools (jupyter_cell, file_write), content injection
+    happens in ds_tool_description() callback (agent_factory.py) so the frontend
+    sees injected args in the approval request.
+    """
+
+    def __init__(self):
+        self.tools: list = []
+        self._last_content: Optional[str] = None
+        self._last_type: Optional[str] = None  # "python" | "sql" | "full_analysis"
+        self._last_description: Optional[str] = None
+        # Track consumed ToolMessage IDs to prevent re-extraction
+        self._consumed_msg_ids: set = set()
+
+    def _consume(self) -> None:
+        """Mark current content as consumed and clear state."""
+        _mid = getattr(self, "_pending_msg_id", None)
+        _dbg(
+            f"CI._consume: pending_id={_mid} content_len={len(self._last_content) if self._last_content else 0}"
+        )
+        if _mid:
+            self._consumed_msg_ids.add(_mid)
+        self._last_content = None
+        self._last_type = None
+        self._last_description = None
+        self._pending_msg_id = None
+
+    def before_model(self, state, runtime):
+        """Scan recent ToolMessages for [CODE]/[SQL] blocks before LLM call."""
+        messages = state.get("messages", [])
+        # 현재 상태 로깅
+        _dbg(
+            f"CI.before_model: msg_count={len(messages)} _last_content="
+            f"{'%dchars' % len(self._last_content) if self._last_content else 'None'} "
+            f"_last_type={self._last_type} consumed_ids={self._consumed_msg_ids} "
+            f"pending_id={getattr(self, '_pending_msg_id', None)}"
+        )
+        # 최근 메시지 내용 디버그 로깅 (AI + Tool 모두)
+        msgs_found = 0
+        for _m in reversed(messages):
+            _mtype = getattr(_m, "type", "unknown")
+            if _mtype == "tool":
+                _c = getattr(_m, "content", "")
+                _mid = getattr(_m, "tool_call_id", None)
+                _consumed = _mid and _mid in self._consumed_msg_ids
+                _dbg(
+                    f"CI.before_model: [{len(messages) - msgs_found - 1}] ToolMessage "
+                    f"id={_mid} consumed={_consumed} "
+                    f"has_CODE={'[CODE]' in _c if isinstance(_c, str) else 'N/A'} "
+                    f"has_SQL={'[SQL]' in _c if isinstance(_c, str) else 'N/A'} "
+                    f"len={len(_c) if isinstance(_c, str) else 0} "
+                    f"preview={repr(_c[:200]) if isinstance(_c, str) else repr(type(_c))}"
+                )
+            elif _mtype == "ai":
+                _tc = getattr(_m, "tool_calls", [])
+                _tc_names = [tc.get("name") for tc in _tc] if _tc else []
+                _dbg(
+                    f"CI.before_model: [{len(messages) - msgs_found - 1}] AIMessage "
+                    f"tool_calls={_tc_names}"
+                )
+            msgs_found += 1
+            if msgs_found >= 6:
+                break
+
+        for msg in reversed(messages):
+            if not (hasattr(msg, "type") and msg.type == "tool"):
+                continue
+            content = getattr(msg, "content", "")
+            if not isinstance(content, str):
+                continue
+
+            # Skip already-consumed ToolMessages
+            msg_id = getattr(msg, "tool_call_id", None) or getattr(msg, "id", None)
+            if msg_id and msg_id in self._consumed_msg_ids:
+                continue
+
+            if "[CODE]" in content:
+                code = _extract_code_block(content, lang="python")
+                if code:
+                    self._last_content = code
+                    self._last_type = "python"
+                    self._last_description = _extract_description(content)
+                    self._pending_msg_id = msg_id
+                    _dbg(
+                        f"CI.EXTRACT: python code {len(code)}chars from ToolMessage id={msg_id}"
+                    )
+                    break
+            elif "[ANALYSIS]" in content and "[RECOMMENDATION]" in content:
+                # Airflow expert analysis response — keep entire content
+                self._last_content = content
+                self._last_type = "full_analysis"
+                self._last_description = None
+                self._pending_msg_id = msg_id
+                _dbg(
+                    f"CI.EXTRACT: full_analysis (airflow) {len(content)}chars from ToolMessage id={msg_id}"
+                )
+                break
+            elif "[SQL]" in content and (
+                "[ANALYSIS]" in content
+                or "[MERMAID]" in content
+                or "[PERFORMANCE]" in content
+                or "[SYNTAX]" in content
+                or "[PLAN]" in content
+                or "## " in content
+            ):
+                # Full analysis response (bq_agent Format 4 etc.) — keep entire content
+                self._last_content = content
+                self._last_type = "full_analysis"
+                self._last_description = None  # description is already in content
+                self._pending_msg_id = msg_id
+                _dbg(
+                    f"CI.EXTRACT: full_analysis {len(content)}chars from ToolMessage id={msg_id}"
+                )
+                break
+            elif "[SQL]" in content:
+                sql = _extract_code_block(content, lang="sql")
+                if sql:
+                    self._last_content = sql
+                    self._last_type = "sql"
+                    self._last_description = _extract_description(content)
+                    self._pending_msg_id = msg_id
+                    _dbg(
+                        f"CI.EXTRACT: SQL {len(sql)}chars from ToolMessage id={msg_id}"
+                    )
+                    break
+        return None  # no state modification
+
+    def after_model(self, state, runtime):
+        """Inject stored content into LLM output's tool_calls args (in-place)."""
+        messages = state.get("messages", [])
+        last_msg = messages[-1] if messages else None
+        _tc = getattr(last_msg, "tool_calls", []) if last_msg else []
+        _tc_info = [
+            {
+                "name": tc.get("name"),
+                "has_code": bool(tc.get("args", {}).get("code")),
+                "has_content": bool(tc.get("args", {}).get("content")),
+            }
+            for tc in _tc
+        ]
+        _dbg(
+            f"CI.after_model: _last_content="
+            f"{'%dchars' % len(self._last_content) if self._last_content else 'None'} "
+            f"_last_type={self._last_type} tool_calls={_tc_info}"
+        )
+        if not self._last_content:
+            return None
+
+        messages = state.get("messages", [])
+        if not messages:
+            return None
+
+        last_msg = messages[-1]
+        tool_calls = getattr(last_msg, "tool_calls", [])
+        if not tool_calls:
+            return None
+
+        for tc in tool_calls:
+            if self._last_type == "full_analysis":
+                if tc["name"] == "jupyter_markdown" and not tc["args"].get("content"):
+                    tc["args"]["content"] = self._last_content
+                    _dbg(
+                        f"CI.INJECT: full_analysis {len(self._last_content)}chars → jupyter_markdown"
+                    )
+                    self._consume()
+                    break
+            elif self._last_type == "python":
+                if tc["name"] == "jupyter_cell" and not tc["args"].get("code"):
+                    tc["args"]["code"] = self._last_content
+                    if self._last_description:
+                        tc["args"]["description"] = self._last_description
+                    _dbg(
+                        f"CI.INJECT: python {len(self._last_content)}chars → jupyter_cell"
+                    )
+                    self._consume()
+                    break
+                elif tc["name"] == "file_write" and not tc["args"].get("content"):
+                    tc["args"]["content"] = self._last_content
+                    _dbg(
+                        f"CI.INJECT: python {len(self._last_content)}chars → file_write"
+                    )
+                    self._consume()
+                    break
+            elif self._last_type == "sql":
+                if tc["name"] == "jupyter_markdown" and not tc["args"].get("content"):
+                    sql_md = f"```sql\n{self._last_content}\n```"
+                    if self._last_description:
+                        sql_md = f"{self._last_description}\n\n{sql_md}"
+                    tc["args"]["content"] = sql_md
+                    _dbg(
+                        f"CI.INJECT: SQL {len(self._last_content)}chars → jupyter_markdown"
+                    )
+                    self._consume()
+                    break
+
+        return None  # tool_calls are modified in-place
+
+    async def awrap_tool_call(self, request, handler):
+        """Fallback: inject stored content during tool execution.
+
+        This handles cases where after_model didn't inject (e.g. content was
+        already injected by ds_tool_description, or HITL skipped after_model).
+        Also consumes state when the target tool already has content
+        (injected by _inject_content_from_state in ds_tool_description).
+        """
+        tool_name = request.tool_call["name"]
+        args = request.tool_call.get("args", {})
+        _dbg(
+            f"CI.awrap_tool_call: tool={tool_name} _last_content="
+            f"{'%dchars' % len(self._last_content) if self._last_content else 'None'} "
+            f"has_code={bool(args.get('code'))} has_content={bool(args.get('content'))}"
+        )
+
+        if self._last_content:
+            new_args = None
+
+            if self._last_type == "full_analysis":
+                if tool_name == "jupyter_markdown":
+                    if not args.get("content"):
+                        new_args = {**args, "content": self._last_content}
+                    self._consume()
+            elif self._last_type == "python":
+                if tool_name == "jupyter_cell":
+                    if not args.get("code"):
+                        new_args = {**args, "code": self._last_content}
+                        if self._last_description and not args.get("description"):
+                            new_args["description"] = self._last_description
+                    # Consume regardless - code was either injected here or by _inject_content_from_state
+                    self._consume()
+                elif tool_name == "file_write":
+                    if not args.get("content"):
+                        new_args = {**args, "content": self._last_content}
+                    self._consume()
+            elif self._last_type == "sql":
+                if tool_name == "jupyter_markdown":
+                    if not args.get("content"):
+                        sql_md = f"```sql\n{self._last_content}\n```"
+                        if self._last_description:
+                            sql_md = f"{self._last_description}\n\n{sql_md}"
+                        new_args = {**args, "content": sql_md}
+                    self._consume()
+
+            if new_args is not None:
+                request = request.override(
+                    tool_call={**request.tool_call, "args": new_args}
+                )
+
+        return await handler(request)