codebuff-agent 2.1.74__py3-none-any.whl

codebuff_core/managers/session_manager.py

@@ -0,0 +1,255 @@
+"""
+Session Manager for persistent conversation storage.
+
+Provides file-based persistence for chat sessions, allowing conversation
+history to survive server restarts.
+"""
+
+import json
+import logging
+import uuid
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from codebuff_core.legacy_paths import resolve_jupyter_config
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ChatMessage:
+    """Individual chat message in a session."""
+
+    role: str  # 'user' | 'assistant'
+    content: str
+    timestamp: float = field(default_factory=lambda: datetime.now().timestamp())
+
+
+@dataclass
+class Session:
+    """Chat session with message history."""
+
+    id: str
+    messages: List[ChatMessage] = field(default_factory=list)
+    created_at: float = field(default_factory=lambda: datetime.now().timestamp())
+    updated_at: float = field(default_factory=lambda: datetime.now().timestamp())
+    metadata: Optional[Dict] = None
+
+
+class SessionManager:
+    """
+    Manages chat session persistence.
+
+    Singleton pattern ensures consistent state across handlers.
+    Sessions are stored as JSON in ~/.jupyter/codebuff_sessions.json.
+    """
+
+    _instance: Optional["SessionManager"] = None
+
+    def __new__(cls) -> "SessionManager":
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if hasattr(self, "_initialized") and self._initialized:
+            return
+        self._initialized = True
+        self._sessions: Dict[str, Session] = {}
+        self._storage_path = resolve_jupyter_config(
+            "codebuff_sessions.json", "hdsp_agent_sessions.json"
+        )
+        self._load_sessions()
+
+    @classmethod
+    def reset_instance(cls):
+        """Reset singleton instance (for testing)"""
+        cls._instance = None
+
+    def _load_sessions(self) -> None:
+        """Load sessions from persistent storage."""
+        try:
+            if self._storage_path.exists():
+                data = json.loads(self._storage_path.read_text(encoding="utf-8"))
+                for sid, sdata in data.items():
+                    self._sessions[sid] = Session(
+                        id=sdata["id"],
+                        messages=[ChatMessage(**m) for m in sdata.get("messages", [])],
+                        created_at=sdata.get("created_at", datetime.now().timestamp()),
+                        updated_at=sdata.get("updated_at", datetime.now().timestamp()),
+                        metadata=sdata.get("metadata"),
+                    )
+                logger.info(
+                    f"Loaded {len(self._sessions)} sessions from {self._storage_path}"
+                )
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse sessions file: {e}. Starting fresh.")
+            self._sessions = {}
+        except Exception as e:
+            logger.error(f"Failed to load sessions: {e}")
+            self._sessions = {}
+
+    def _save_sessions(self) -> None:
+        """Persist sessions to storage."""
+        try:
+            data = {}
+            for sid, session in self._sessions.items():
+                data[sid] = {
+                    "id": session.id,
+                    "messages": [asdict(m) for m in session.messages],
+                    "created_at": session.created_at,
+                    "updated_at": session.updated_at,
+                    "metadata": session.metadata,
+                }
+            self._storage_path.parent.mkdir(parents=True, exist_ok=True)
+            self._storage_path.write_text(
+                json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8"
+            )
+        except Exception as e:
+            logger.error(f"Failed to save sessions: {e}")
+
+    def get_session(self, session_id: str) -> Optional[Session]:
+        """Get session by ID."""
+        return self._sessions.get(session_id)
+
+    def create_session(self, session_id: Optional[str] = None) -> Session:
+        """Create a new session with optional custom ID."""
+        sid = session_id or str(uuid.uuid4())
+        now = datetime.now().timestamp()
+        session = Session(id=sid, messages=[], created_at=now, updated_at=now)
+        self._sessions[sid] = session
+        self._save_sessions()
+        logger.debug(f"Created new session: {sid}")
+        return session
+
+    def get_or_create_session(self, session_id: Optional[str] = None) -> Session:
+        """Get existing session or create new one."""
+        if session_id and session_id in self._sessions:
+            return self._sessions[session_id]
+        return self.create_session(session_id)
+
+    def add_message(self, session_id: str, role: str, content: str) -> ChatMessage:
+        """Add a message to a session. Creates session if needed."""
+        session = self._sessions.get(session_id)
+        if not session:
+            session = self.create_session(session_id)
+
+        msg = ChatMessage(role=role, content=content)
+        session.messages.append(msg)
+        session.updated_at = msg.timestamp
+        self._save_sessions()
+        return msg
+
+    def store_messages(
+        self, session_id: str, user_message: str, assistant_response: str
+    ) -> None:
+        """Store user and assistant message pair. Creates session if needed."""
+        session = self._sessions.get(session_id)
+        if not session:
+            session = self.create_session(session_id)
+
+        now = datetime.now().timestamp()
+        session.messages.append(
+            ChatMessage(role="user", content=user_message, timestamp=now)
+        )
+        session.messages.append(
+            ChatMessage(role="assistant", content=assistant_response, timestamp=now)
+        )
+        session.updated_at = now
+        self._save_sessions()
+
+    def get_recent_messages(
+        self, session_id: str, limit: int = 10
+    ) -> List[ChatMessage]:
+        """Get most recent messages from a session."""
+        session = self._sessions.get(session_id)
+        if not session:
+            return []
+        return session.messages[-limit:]
+
+    def build_context(
+        self,
+        session_id: str,
+        max_messages: int = 5,
+        compress: bool = False,
+        target_tokens: Optional[int] = None,
+    ) -> Optional[str]:
+        """Build conversation context string from recent history.
+
+        Args:
+            session_id: Session ID to get context from
+            max_messages: Maximum messages to include (before compression)
+            compress: Enable context compression for token efficiency
+            target_tokens: Target token count for compression (default: auto)
+
+        Returns:
+            Formatted context string or None if session empty
+        """
+        session = self._sessions.get(session_id)
+        if not session or not session.messages:
+            return None
+
+        messages = session.messages[-max_messages:]
+
+        if compress:
+            # Lazy import to avoid circular dependency
+            try:
+                from codebuff_core.core.context_condenser import get_context_condenser
+
+                condenser = get_context_condenser()
+                msg_dicts = [{"role": m.role, "content": m.content} for m in messages]
+                compressed, _stats = condenser.condense(msg_dicts, target_tokens)
+                messages_to_format = compressed
+            except ImportError:
+                # Fallback if context_condenser not available
+                messages_to_format = [
+                    {"role": m.role, "content": m.content} for m in messages
+                ]
+        else:
+            messages_to_format = [
+                {"role": m.role, "content": m.content} for m in messages
+            ]
+
+        return "\n".join(
+            [
+                f"{'User' if msg['role'] == 'user' else 'Assistant'}: {msg['content']}"
+                for msg in messages_to_format
+            ]
+        )
+
+    def list_sessions(self, limit: int = 50) -> List[Session]:
+        """List sessions sorted by most recently updated."""
+        sessions = sorted(
+            self._sessions.values(), key=lambda s: s.updated_at, reverse=True
+        )
+        return sessions[:limit]
+
+    def delete_session(self, session_id: str) -> bool:
+        """Delete a session by ID."""
+        if session_id in self._sessions:
+            del self._sessions[session_id]
+            self._save_sessions()
+            logger.debug(f"Deleted session: {session_id}")
+            return True
+        return False
+
+    def clear_all_sessions(self) -> int:
+        """Delete all sessions. Returns count of deleted sessions."""
+        count = len(self._sessions)
+        self._sessions = {}
+        self._save_sessions()
+        logger.info(f"Cleared {count} sessions")
+        return count
+
+
+# Singleton accessor
+_session_manager: Optional[SessionManager] = None
+
+
+def get_session_manager() -> SessionManager:
+    """Get the singleton SessionManager instance."""
+    global _session_manager
+    if _session_manager is None:
+        _session_manager = SessionManager()
+    return _session_manager

codebuff_core/middleware/__init__.py

@@ -0,0 +1,3 @@
+from codebuff_core.middleware.error_recovery import create_error_recovery_middleware
+
+__all__ = ["create_error_recovery_middleware"]

codebuff_core/middleware/error_recovery.py

@@ -0,0 +1,394 @@
+"""
+Error Recovery Middleware - 런타임 에러 발생 시 자동 Re-plan 프롬프트 주입
+
+HITL 도구(jupyter_cell_tool 등)에서 에러가 반환되면:
+1. ErrorClassifier로 에러를 분류 (기존 코드 재사용)
+2. ReplanDecision에 따른 re-plan 프롬프트를 HumanMessage로 주입
+3. LLM이 write_todos 업데이트 후 수정된 코드를 실행하도록 유도
+"""
+
+import json
+import logging
+import re
+from typing import Optional, Tuple
+
+from langchain_core.messages import HumanMessage
+
+from codebuff_core.core.error_classifier import (
+    ErrorAnalysis,
+    ReplanDecision,
+    get_error_classifier,
+)
+from codebuff_core.middleware.logging_utils import _with_middleware_logging
+
+logger = logging.getLogger(__name__)
+
+# HITL tools that execute code and can produce runtime errors
+HITL_ERROR_TOOLS = {
+    "jupyter_cell_tool",
+    "jupyter_cell",
+    "execute_command_tool",
+    "execute_command",
+}
+
+# Marker prefix for error recovery messages (used for retry counting)
+ERROR_RECOVERY_MARKER = "[ERROR_RECOVERY]"
+
+# Error patterns in tool output
+ERROR_PATTERNS = [
+    r"(\w+Error):\s*(.+)",
+    r"(\w+Exception):\s*(.+)",
+    r"Traceback \(most recent call last\)",
+]
+
+
+def create_error_recovery_middleware(wrap_model_call, max_retries: int = 3):
+    """Create middleware that injects re-plan prompts when HITL tool errors occur.
+
+    Uses ErrorClassifier (codebuff_core.core.error_classifier) to classify errors
+    and generates decision-specific prompts to guide the LLM toward self-correction.
+
+    Args:
+        wrap_model_call: LangChain's wrap_model_call decorator
+        max_retries: Maximum consecutive error recovery attempts (default: 3)
+
+    Returns:
+        Middleware function
+    """
+    classifier = get_error_classifier()
+
+    @wrap_model_call
+    @_with_middleware_logging("error_recovery")
+    def error_recovery_replan(request, handler):
+        messages = request.messages
+        if not messages:
+            return handler(request)
+
+        last_msg = messages[-1]
+
+        # 1. Is this a ToolMessage from a HITL tool with an error?
+        if not _is_hitl_tool_error(last_msg):
+            return handler(request)
+
+        # 2. Check consecutive error recovery count
+        recovery_count = _count_error_recoveries(messages)
+        if recovery_count >= max_retries:
+            logger.warning(
+                "Error recovery: max retries reached (%d/%d). Passing through.",
+                recovery_count,
+                max_retries,
+            )
+            return handler(request)
+
+        # 3. Extract error parts and classify
+        error_type, error_msg, traceback_str = _extract_error_parts(last_msg.content)
+        if not error_type:
+            return handler(request)
+
+        analysis = classifier.classify(error_type, error_msg, traceback_str)
+
+        logger.info(
+            "Error recovery: %s -> %s (attempt %d/%d)",
+            error_type,
+            analysis.decision.value,
+            recovery_count + 1,
+            max_retries,
+        )
+
+        # 4. Build re-plan prompt
+        todos = request.state.get("todos", [])
+        prompt = _build_replan_prompt(analysis, todos, error_type, error_msg)
+
+        # 5. Inject HumanMessage with re-plan instructions
+        new_messages = list(messages) + [HumanMessage(content=prompt)]
+        request = request.override(messages=new_messages)
+
+        return handler(request)
+
+    return error_recovery_replan
+
+
+def _is_hitl_tool_error(msg) -> bool:
+    """Check if a message is a ToolMessage from a HITL tool containing an error.
+
+    Supports both plain text errors and structured JSON with execution_result.
+    """
+    if getattr(msg, "type", "") != "tool":
+        return False
+
+    tool_name = getattr(msg, "name", "") or ""
+    if tool_name not in HITL_ERROR_TOOLS:
+        return False
+
+    content = getattr(msg, "content", "") or ""
+    if not isinstance(content, str):
+        return False
+
+    # Check structured execution_result first (from HITL tool execution)
+    try:
+        data = json.loads(content)
+        if isinstance(data, dict):
+            exec_result = data.get("execution_result")
+            if isinstance(exec_result, dict):
+                return exec_result.get("success") is False or bool(
+                    exec_result.get("error")
+                )
+    except (json.JSONDecodeError, TypeError, ValueError):
+        pass
+
+    return _has_error_indicators(content)
+
+
+def _has_error_indicators(content: str) -> bool:
+    """Check if content contains Python error indicators."""
+    for pattern in ERROR_PATTERNS:
+        if re.search(pattern, content):
+            return True
+    return False
+
+
+def _count_error_recoveries(messages) -> int:
+    """Count consecutive error recovery attempts since last successful HITL result.
+
+    Scans backwards through messages. Counts ERROR_RECOVERY_MARKER occurrences.
+    Resets count when a successful (non-error) HITL tool result is found.
+    """
+    count = 0
+    for msg in reversed(messages):
+        content = getattr(msg, "content", "") or ""
+        msg_type = getattr(msg, "type", "")
+
+        if isinstance(content, str) and ERROR_RECOVERY_MARKER in content:
+            count += 1
+        elif msg_type == "tool":
+            tool_name = getattr(msg, "name", "") or ""
+            if tool_name in HITL_ERROR_TOOLS:
+                if isinstance(content, str) and _has_error_indicators(content):
+                    # Error result - continue counting
+                    continue
+                else:
+                    # Successful result - reset and stop
+                    break
+            # Non-HITL tool result - skip
+            continue
+        elif msg_type in ("ai", "human"):
+            # Skip AI/human messages that aren't recovery markers
+            if msg_type == "human" and isinstance(content, str):
+                if ERROR_RECOVERY_MARKER not in content:
+                    break
+            continue
+        else:
+            break
+
+    return count
+
+
+def _extract_error_parts(content: str) -> Tuple[Optional[str], str, str]:
+    """Extract error type, message, and traceback from tool output.
+
+    Supports two formats:
+    1. Structured JSON with execution_result (from HITL tool execution)
+    2. Plain text with Python error patterns (traceback, ErrorType: message)
+
+    Returns:
+        (error_type, error_message, traceback_string)
+    """
+    if not content or not isinstance(content, str):
+        return None, "", ""
+
+    # Strategy 1: Structured extraction from JSON tool output with execution_result
+    try:
+        data = json.loads(content)
+        if isinstance(data, dict):
+            exec_result = data.get("execution_result")
+            if isinstance(exec_result, dict):
+                error_type = exec_result.get("error_type")
+                if error_type:
+                    error_msg = exec_result.get("error", "")
+                    # Remove error_type prefix from error message
+                    prefix = f"{error_type}: "
+                    if error_msg.startswith(prefix):
+                        error_msg = error_msg[len(prefix) :]
+                    traceback_str = exec_result.get("traceback", "") or ""
+                    return error_type, error_msg, traceback_str
+    except (json.JSONDecodeError, TypeError, ValueError):
+        pass
+
+    # Strategy 2: Regex-based extraction from plain text
+    lines = content.strip().split("\n")
+
+    error_type = None
+    error_msg = ""
+
+    # Search from the end for error type pattern (Python convention)
+    for line in reversed(lines):
+        line = line.strip()
+        # Remove ANSI escape codes
+        line = re.sub(r"\x1b\[[0-9;]*m", "", line)
+
+        # Match "ErrorType: error message" anywhere in the line
+        # (re.search handles both standalone lines and embedded in JSON/dict repr)
+        match = re.search(r"(\w+(?:Error|Exception)):\s*(.*)", line)
+        if match:
+            error_type = match.group(1)
+            error_msg = match.group(2)
+            break
+
+        # Match standalone error type at start of line
+        match = re.match(r"^(\w+(?:Error|Exception))\s*$", line)
+        if match:
+            error_type = match.group(1)
+            break
+
+    # Extract traceback
+    traceback_str = ""
+    traceback_match = re.search(
+        r"(Traceback \(most recent call last\)[\s\S]*)", content
+    )
+    if traceback_match:
+        traceback_str = traceback_match.group(1)
+
+    return error_type, error_msg, traceback_str
+
+
+def _build_replan_prompt(
+    analysis: ErrorAnalysis,
+    todos: list,
+    error_type: str,
+    error_msg: str,
+) -> str:
+    """Build re-plan prompt based on ErrorClassifier analysis."""
+    decision = analysis.decision
+    guidance = _get_error_guidance(error_type, error_msg)
+
+    # Pending todos context
+    pending_todos = [t for t in todos if t.get("status") in ("pending", "in_progress")]
+    todos_context = ""
+    if pending_todos:
+        todo_items = "\n".join(
+            f"  - [{t.get('status', '?')}] {t.get('content', '')[:60]}"
+            for t in pending_todos[:5]
+        )
+        todos_context = f"\n\nCurrent pending tasks:\n{todo_items}"
+
+    error_preview = error_msg[:200] if error_msg else "unknown error"
+
+    if decision == ReplanDecision.REFINE:
+        prompt = (
+            f"{ERROR_RECOVERY_MARKER} The previous code execution failed "
+            f"with {error_type}: {error_preview}\n\n"
+            f"Root cause: {analysis.root_cause}\n\n"
+            f"Instructions:\n"
+            f"1. Analyze the error carefully\n"
+            f"2. {guidance}\n"
+            f"3. Call write_todos to update your plan with the fix approach\n"
+            f"4. Call jupyter_cell_tool with corrected code\n"
+            f"5. Do NOT call final_answer_tool until the error is resolved"
+            f"{todos_context}"
+        )
+    elif decision == ReplanDecision.INSERT_STEPS:
+        pkg = analysis.missing_package or "unknown"
+        prompt = (
+            f"{ERROR_RECOVERY_MARKER} The previous code execution failed "
+            f"because package '{pkg}' is missing.\n\n"
+            f"Instructions:\n"
+            f"1. Call jupyter_cell_tool to install: !pip install {pkg}\n"
+            f"2. Then re-execute the original code\n"
+            f"3. Call write_todos to reflect the package installation step\n"
+            f"4. Do NOT call final_answer_tool until the error is resolved"
+            f"{todos_context}"
+        )
+    elif decision == ReplanDecision.REPLACE_STEP:
+        prompt = (
+            f"{ERROR_RECOVERY_MARKER} The previous approach completely failed "
+            f"with {error_type}: {error_preview}\n\n"
+            f"Root cause: {analysis.root_cause}\n\n"
+            f"Instructions:\n"
+            f"1. The current approach is not viable - try a completely "
+            f"different method\n"
+            f"2. {guidance}\n"
+            f"3. Call write_todos to replace the current step with a new "
+            f"approach\n"
+            f"4. Call jupyter_cell_tool with the new approach\n"
+            f"5. Do NOT call final_answer_tool until the error is resolved"
+            f"{todos_context}"
+        )
+    elif decision == ReplanDecision.REPLAN_REMAINING:
+        prompt = (
+            f"{ERROR_RECOVERY_MARKER} A critical error occurred: "
+            f"{error_type}: {error_preview}\n\n"
+            f"Root cause: {analysis.root_cause}\n\n"
+            f"Instructions:\n"
+            f"1. This is a fundamental issue requiring replanning all "
+            f"remaining steps\n"
+            f"2. {guidance}\n"
+            f"3. Call write_todos to redesign the remaining plan entirely\n"
+            f"4. Start executing the new plan with jupyter_cell_tool\n"
+            f"5. Do NOT call final_answer_tool until the error is resolved"
+            f"{todos_context}"
+        )
+    else:
+        prompt = (
+            f"{ERROR_RECOVERY_MARKER} An error occurred: "
+            f"{error_type}: {error_preview}\n\n"
+            f"Fix the error and continue. "
+            f"Do NOT call final_answer_tool until resolved."
+            f"{todos_context}"
+        )
+
+    return prompt
+
+
+def _get_error_guidance(error_type: str, error_msg: str) -> str:
+    """Get error-specific guidance for the LLM."""
+    guidance_map = {
+        "KeyError": (
+            "Check the actual column/key names by running "
+            "`df.columns.tolist()` or `print(list(dict.keys()))` first. "
+            "The key name may differ from what you expected "
+            "(case, spaces, encoding)."
+        ),
+        "FileNotFoundError": (
+            "Verify the file path by running "
+            "`import os; print(os.listdir('.'))` or checking subdirectories. "
+            "The file may be in a different location or have a different name."
+        ),
+        "NameError": (
+            "Ensure all imports and variable definitions are in the same cell. "
+            "Variables from previous cells may not be available if the kernel "
+            "was restarted."
+        ),
+        "ModuleNotFoundError": (
+            "Install the missing package using "
+            "`!pip install <package_name>` before importing."
+        ),
+        "ImportError": (
+            "Check if the package is installed and the import path is correct. "
+            "Some packages have different import names than install names "
+            "(e.g., `pip install scikit-learn` but `import sklearn`)."
+        ),
+        "TypeError": (
+            "Check function arguments and data types. Verify the correct "
+            "number and types of arguments are being passed."
+        ),
+        "ValueError": (
+            "Check input values. Verify data types, ranges, and formats "
+            "match what the function expects."
+        ),
+        "IndexError": (
+            "Check array/list bounds. Verify the data has the expected "
+            "number of elements before accessing by index."
+        ),
+        "AttributeError": (
+            "Check if the object has the expected type. Use `type(obj)` "
+            "and `dir(obj)` to verify available attributes and methods."
+        ),
+        "SyntaxError": (
+            "Check for missing colons, brackets, parentheses, or indentation errors."
+        ),
+    }
+
+    return guidance_map.get(
+        error_type,
+        f"Analyze the {error_type} carefully and fix the root cause before retrying.",
+    )

codebuff_core/middleware/logging_utils.py

@@ -0,0 +1,33 @@
+"""
+Logging utilities for middleware.
+
+Provides helper functions for structured logging of middleware execution.
+"""
+
+import logging
+from functools import wraps
+
+logger = logging.getLogger(__name__)
+
+LOG_SEPARATOR = "=" * 96
+
+
+def _format_middleware_marker(name: str, stage: str) -> str:
+    """Format a middleware execution marker."""
+    return "\n".join([LOG_SEPARATOR, f"MIDDLEWARE {stage}: {name}", LOG_SEPARATOR])
+
+
+def _with_middleware_logging(name: str):
+    """Decorator to add logging around middleware execution."""
+
+    def decorator(func):
+        @wraps(func)
+        def wrapped(request, handler):
+            logger.info("%s", _format_middleware_marker(name, "START"))
+            response = func(request, handler)
+            logger.info("%s", _format_middleware_marker(name, "END"))
+            return response
+
+        return wrapped
+
+    return decorator