codebuff-agent 2.1.74__py3-none-any.whl

codebuff_core/tests/test_error_recovery.py

@@ -0,0 +1,647 @@
+"""
+Tests for error recovery middleware.
+
+Tests cover:
+- ErrorClassifier integration (REFINE, INSERT_STEPS decisions)
+- Middleware passthrough for non-error and non-HITL messages
+- HumanMessage injection on error
+- Max retries enforcement
+- Prompt content (write_todos, jupyter_cell_tool, final_answer_tool ban)
+- Consecutive error count reset on success
+"""
+
+from unittest.mock import MagicMock
+
+from langchain_core.messages import HumanMessage
+
+from codebuff_core.middleware.error_recovery import (
+    ERROR_RECOVERY_MARKER,
+    _build_replan_prompt,
+    _count_error_recoveries,
+    _extract_error_parts,
+    _get_error_guidance,
+    _has_error_indicators,
+    _is_hitl_tool_error,
+    create_error_recovery_middleware,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+class MockMessage:
+    """Minimal message mock for testing."""
+
+    def __init__(self, type="", name="", content=""):
+        self.type = type
+        self.name = name
+        self.content = content
+
+
+class MockRequest:
+    """Minimal request mock for testing."""
+
+    def __init__(self, messages=None, state=None):
+        self.messages = messages or []
+        self.state = state or {}
+
+    def override(self, **kwargs):
+        return MockRequest(
+            messages=kwargs.get("messages", self.messages),
+            state=kwargs.get("state", self.state),
+        )
+
+
+def mock_wrap_model_call(fn):
+    """Identity decorator mimicking wrap_model_call for testing."""
+    return fn
+
+
+def make_handler(response=None):
+    """Create a mock handler that records calls."""
+    resp = response or MagicMock()
+    handler = MagicMock(return_value=resp)
+    return handler
+
+
+def make_tool_msg(name, content):
+    """Create a ToolMessage mock."""
+    return MockMessage(type="tool", name=name, content=content)
+
+
+def make_error_content(error_type="KeyError", error_msg="'Survived'"):
+    """Create typical Python error output."""
+    return (
+        f"Traceback (most recent call last):\n"
+        f'  File "<stdin>", line 1, in <module>\n'
+        f"{error_type}: {error_msg}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# ErrorClassifier integration tests
+# ---------------------------------------------------------------------------
+
+
+class TestErrorClassifierIntegration:
+    """Test that ErrorClassifier correctly classifies errors for the middleware."""
+
+    def test_classify_keyerror_refine(self):
+        """KeyError should be classified as REFINE with correct guidance."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            get_error_classifier,
+        )
+
+        classifier = get_error_classifier()
+        analysis = classifier.classify("KeyError", "'Survived'", "")
+
+        assert analysis.decision == ReplanDecision.REFINE
+        assert "키" in analysis.root_cause or "key" in analysis.root_cause.lower()
+
+    def test_classify_module_not_found_insert_steps(self):
+        """ModuleNotFoundError should be classified as INSERT_STEPS."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            get_error_classifier,
+        )
+
+        classifier = get_error_classifier()
+        analysis = classifier.classify(
+            "ModuleNotFoundError",
+            "No module named 'seaborn'",
+            "ModuleNotFoundError: No module named 'seaborn'",
+        )
+
+        assert analysis.decision == ReplanDecision.INSERT_STEPS
+        assert analysis.missing_package == "seaborn"
+
+    def test_classify_file_not_found_refine(self):
+        """FileNotFoundError should be classified as REFINE."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            get_error_classifier,
+        )
+
+        classifier = get_error_classifier()
+        analysis = classifier.classify(
+            "FileNotFoundError",
+            "[Errno 2] No such file or directory: 'data.csv'",
+            "",
+        )
+
+        assert analysis.decision == ReplanDecision.REFINE
+
+
+# ---------------------------------------------------------------------------
+# Helper function tests
+# ---------------------------------------------------------------------------
+
+
+class TestHelperFunctions:
+    """Test helper functions used by the middleware."""
+
+    def test_is_hitl_tool_error_true(self):
+        """ToolMessage from jupyter_cell_tool with error should return True."""
+        msg = make_tool_msg("jupyter_cell_tool", make_error_content())
+        assert _is_hitl_tool_error(msg) is True
+
+    def test_is_hitl_tool_error_no_error(self):
+        """ToolMessage from jupyter_cell_tool without error should return False."""
+        msg = make_tool_msg("jupyter_cell_tool", "   x  y\n0  1  2\n1  3  4")
+        assert _is_hitl_tool_error(msg) is False
+
+    def test_is_hitl_tool_error_non_hitl_tool(self):
+        """ToolMessage from non-HITL tool should return False."""
+        msg = make_tool_msg("markdown_tool", make_error_content())
+        assert _is_hitl_tool_error(msg) is False
+
+    def test_is_hitl_tool_error_non_tool_message(self):
+        """Non-tool message should return False."""
+        msg = MockMessage(type="human", content=make_error_content())
+        assert _is_hitl_tool_error(msg) is False
+
+    def test_has_error_indicators_with_traceback(self):
+        """Content with traceback should be detected."""
+        content = "Traceback (most recent call last):\n  File ...\nKeyError: 'x'"
+        assert _has_error_indicators(content) is True
+
+    def test_has_error_indicators_with_error_line(self):
+        """Content with error line should be detected."""
+        assert _has_error_indicators("ValueError: invalid literal") is True
+
+    def test_has_error_indicators_clean_output(self):
+        """Clean output without errors should return False."""
+        assert _has_error_indicators("   x  y\n0  1  2") is False
+
+    def test_extract_error_parts_keyerror(self):
+        """Should extract KeyError type and message from traceback."""
+        content = make_error_content("KeyError", "'Survived'")
+        error_type, error_msg, tb = _extract_error_parts(content)
+
+        assert error_type == "KeyError"
+        assert "'Survived'" in error_msg
+        assert "Traceback" in tb
+
+    def test_extract_error_parts_module_error(self):
+        """Should extract ModuleNotFoundError."""
+        content = (
+            "Traceback (most recent call last):\n"
+            '  File "<stdin>", line 1, in <module>\n'
+            "ModuleNotFoundError: No module named 'pandas'"
+        )
+        error_type, error_msg, tb = _extract_error_parts(content)
+
+        assert error_type == "ModuleNotFoundError"
+        assert "pandas" in error_msg
+
+    def test_extract_error_parts_no_error(self):
+        """Should return None for clean output."""
+        error_type, error_msg, tb = _extract_error_parts("Hello world\n42")
+        assert error_type is None
+
+    def test_extract_error_parts_ansi_codes(self):
+        """Should handle ANSI color codes in traceback."""
+        content = (
+            "Traceback (most recent call last):\n\x1b[31mKeyError\x1b[0m: 'missing_col'"
+        )
+        error_type, error_msg, tb = _extract_error_parts(content)
+        assert error_type == "KeyError"
+
+    def test_get_error_guidance_keyerror(self):
+        """KeyError guidance should mention columns/keys."""
+        guidance = _get_error_guidance("KeyError", "'col'")
+        assert "column" in guidance.lower() or "key" in guidance.lower()
+
+    def test_get_error_guidance_unknown(self):
+        """Unknown error type should get generic guidance."""
+        guidance = _get_error_guidance("CustomError", "something")
+        assert "CustomError" in guidance
+
+
+# ---------------------------------------------------------------------------
+# Count error recoveries tests
+# ---------------------------------------------------------------------------
+
+
+class TestCountErrorRecoveries:
+    """Test consecutive error recovery counting logic."""
+
+    def test_no_recoveries(self):
+        """Empty messages should return 0."""
+        assert _count_error_recoveries([]) == 0
+
+    def test_one_recovery(self):
+        """Single recovery marker should return 1."""
+        messages = [
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix it"),
+        ]
+        assert _count_error_recoveries(messages) == 1
+
+    def test_multiple_recoveries(self):
+        """Multiple recovery markers in error chain should be counted."""
+        messages = [
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 1"),
+            MockMessage(type="ai", content="I'll fix it"),
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 2"),
+            MockMessage(type="ai", content="trying again"),
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 3"),
+        ]
+        assert _count_error_recoveries(messages) == 3
+
+    def test_consecutive_count_resets_on_success(self):
+        """Count should reset when a successful HITL tool result appears."""
+        messages = [
+            # First error cycle
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 1"),
+            MockMessage(type="ai", content="fixing"),
+            # SUCCESS - this should reset the count
+            make_tool_msg("jupyter_cell_tool", "   x  y\n0  1  2"),
+            MockMessage(type="ai", content="next step"),
+            # New error cycle - only these count
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 2"),
+        ]
+        assert _count_error_recoveries(messages) == 1
+
+    def test_non_hitl_tools_dont_reset(self):
+        """Non-HITL tool results should not affect the count."""
+        messages = [
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 1"),
+            MockMessage(type="ai", content="reading file"),
+            make_tool_msg("read_file_tool", "file content here"),
+            MockMessage(type="ai", content="got it"),
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 2"),
+        ]
+        assert _count_error_recoveries(messages) == 2
+
+
+# ---------------------------------------------------------------------------
+# Middleware integration tests
+# ---------------------------------------------------------------------------
+
+
+class TestMiddlewareIntegration:
+    """Test the full middleware function via create_error_recovery_middleware."""
+
+    def _create_middleware(self, max_retries=3):
+        """Create middleware with mock wrap_model_call."""
+        return create_error_recovery_middleware(
+            mock_wrap_model_call, max_retries=max_retries
+        )
+
+    def test_passthrough_no_error(self):
+        """Normal ToolMessage (no error) should pass through unchanged."""
+        middleware = self._create_middleware()
+        handler = make_handler()
+        request = MockRequest(
+            messages=[make_tool_msg("jupyter_cell_tool", "   x  y\n0  1  2\n1  3  4")]
+        )
+
+        middleware(request, handler)
+
+        # Handler called with original request (no extra messages)
+        called_request = handler.call_args[0][0]
+        assert len(called_request.messages) == 1
+
+    def test_passthrough_non_hitl_tool(self):
+        """ToolMessage from non-HITL tool should pass through."""
+        middleware = self._create_middleware()
+        handler = make_handler()
+        request = MockRequest(
+            messages=[make_tool_msg("markdown_tool", make_error_content())]
+        )
+
+        middleware(request, handler)
+
+        called_request = handler.call_args[0][0]
+        assert len(called_request.messages) == 1
+
+    def test_passthrough_empty_messages(self):
+        """Empty messages should pass through."""
+        middleware = self._create_middleware()
+        handler = make_handler()
+        request = MockRequest(messages=[])
+
+        middleware(request, handler)
+        handler.assert_called_once()
+
+    def test_injects_on_error(self):
+        """HITL tool error should inject a HumanMessage."""
+        middleware = self._create_middleware()
+        handler = make_handler()
+        request = MockRequest(
+            messages=[make_tool_msg("jupyter_cell_tool", make_error_content())],
+            state={"todos": []},
+        )
+
+        middleware(request, handler)
+
+        called_request = handler.call_args[0][0]
+        assert len(called_request.messages) == 2
+        injected = called_request.messages[1]
+        assert isinstance(injected, HumanMessage)
+        assert ERROR_RECOVERY_MARKER in injected.content
+
+    def test_max_retries_passthrough(self):
+        """After max_retries, should pass through without injection."""
+        middleware = self._create_middleware(max_retries=2)
+        handler = make_handler()
+
+        # Build messages with 2 existing recovery markers
+        messages = [
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 1"),
+            MockMessage(type="ai", content="fixing"),
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+            MockMessage(type="human", content=f"{ERROR_RECOVERY_MARKER} fix 2"),
+            MockMessage(type="ai", content="trying"),
+            # This is the 3rd error - should hit max_retries=2
+            make_tool_msg("jupyter_cell_tool", make_error_content()),
+        ]
+        request = MockRequest(messages=messages, state={"todos": []})
+
+        middleware(request, handler)
+
+        # Should pass through with original messages (no new injection)
+        called_request = handler.call_args[0][0]
+        assert len(called_request.messages) == len(messages)
+
+    def test_injects_for_execute_command_tool(self):
+        """execute_command_tool errors should also trigger recovery."""
+        middleware = self._create_middleware()
+        handler = make_handler()
+        request = MockRequest(
+            messages=[
+                make_tool_msg(
+                    "execute_command_tool",
+                    "Traceback (most recent call last):\n"
+                    "FileNotFoundError: [Errno 2] No such file: 'missing.sh'",
+                )
+            ],
+            state={"todos": []},
+        )
+
+        middleware(request, handler)
+
+        called_request = handler.call_args[0][0]
+        assert len(called_request.messages) == 2
+        assert ERROR_RECOVERY_MARKER in called_request.messages[1].content
+
+
+# ---------------------------------------------------------------------------
+# Prompt content tests
+# ---------------------------------------------------------------------------
+
+
+class TestPromptContent:
+    """Test that generated prompts contain required instructions."""
+
+    def _get_refine_prompt(self):
+        """Get a REFINE prompt for KeyError."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            ErrorAnalysis,
+        )
+
+        analysis = ErrorAnalysis(
+            decision=ReplanDecision.REFINE,
+            root_cause="딕셔너리/데이터프레임 키 없음: 'Survived'",
+            reasoning="KeyError는 코드 수정으로 해결 가능합니다.",
+        )
+        return _build_replan_prompt(analysis, [], "KeyError", "'Survived'")
+
+    def _get_insert_steps_prompt(self):
+        """Get an INSERT_STEPS prompt for ModuleNotFoundError."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            ErrorAnalysis,
+        )
+
+        analysis = ErrorAnalysis(
+            decision=ReplanDecision.INSERT_STEPS,
+            root_cause="'seaborn' 모듈이 설치되지 않음",
+            reasoning="ModuleNotFoundError는 항상 패키지 설치로 해결합니다.",
+            missing_package="seaborn",
+        )
+        return _build_replan_prompt(
+            analysis, [], "ModuleNotFoundError", "No module named 'seaborn'"
+        )
+
+    def test_prompt_contains_replan_instructions(self):
+        """Prompt should instruct to call write_todos and jupyter_cell_tool."""
+        prompt = self._get_refine_prompt()
+        assert "write_todos" in prompt
+        assert "jupyter_cell_tool" in prompt
+
+    def test_prompt_forbids_final_answer(self):
+        """Prompt should forbid calling final_answer_tool."""
+        prompt = self._get_refine_prompt()
+        assert "final_answer_tool" in prompt
+        assert "NOT" in prompt
+
+    def test_prompt_contains_error_recovery_marker(self):
+        """Prompt should contain the error recovery marker."""
+        prompt = self._get_refine_prompt()
+        assert ERROR_RECOVERY_MARKER in prompt
+
+    def test_refine_prompt_contains_guidance(self):
+        """REFINE prompt should contain error-specific guidance."""
+        prompt = self._get_refine_prompt()
+        # KeyError guidance mentions columns/keys
+        assert "column" in prompt.lower() or "key" in prompt.lower()
+
+    def test_insert_steps_prompt_contains_pip(self):
+        """INSERT_STEPS prompt should contain pip install instruction."""
+        prompt = self._get_insert_steps_prompt()
+        assert "pip install" in prompt
+        assert "seaborn" in prompt
+
+    def test_prompt_includes_todos_context(self):
+        """Prompt should include pending todos when available."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            ErrorAnalysis,
+        )
+
+        analysis = ErrorAnalysis(
+            decision=ReplanDecision.REFINE,
+            root_cause="KeyError",
+            reasoning="코드 수정",
+        )
+        todos = [
+            {"status": "completed", "content": "Load data"},
+            {"status": "in_progress", "content": "Analyze survival rate"},
+            {"status": "pending", "content": "Create visualization"},
+        ]
+        prompt = _build_replan_prompt(analysis, todos, "KeyError", "'col'")
+
+        assert "pending tasks" in prompt.lower()
+        assert "Analyze survival" in prompt
+        assert "Create visualization" in prompt
+        # Completed task should not appear in pending
+        assert "Load data" not in prompt
+
+    def test_replace_step_prompt(self):
+        """REPLACE_STEP prompt should instruct a completely different method."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            ErrorAnalysis,
+        )
+
+        analysis = ErrorAnalysis(
+            decision=ReplanDecision.REPLACE_STEP,
+            root_cause="Approach failed",
+            reasoning="Need different method",
+        )
+        prompt = _build_replan_prompt(analysis, [], "RuntimeError", "approach failed")
+        assert "different" in prompt.lower()
+        assert "write_todos" in prompt
+
+    def test_replan_remaining_prompt(self):
+        """REPLAN_REMAINING prompt should instruct full replanning."""
+        from codebuff_core.core.error_classifier import (
+            ReplanDecision,
+            ErrorAnalysis,
+        )
+
+        analysis = ErrorAnalysis(
+            decision=ReplanDecision.REPLAN_REMAINING,
+            root_cause="Critical system error",
+            reasoning="Full replan needed",
+        )
+        prompt = _build_replan_prompt(analysis, [], "OSError", "system error")
+        assert "replan" in prompt.lower() or "remaining" in prompt.lower()
+        assert "write_todos" in prompt
+
+
+# ---------------------------------------------------------------------------
+# JSON structured content tests (HITL execution_result)
+# ---------------------------------------------------------------------------
+
+
+class TestStructuredContent:
+    """Test handling of JSON-structured tool output with execution_result."""
+
+    def _make_json_content(
+        self, success, output="", error=None, error_type=None, traceback=None
+    ):
+        """Create JSON content mimicking jupyter_cell_tool output with execution_result."""
+        import json
+
+        return json.dumps(
+            {
+                "tool": "jupyter_cell",
+                "parameters": {"code": "print('hello')", "description": None},
+                "status": "complete" if success else "complete",
+                "message": "Code cell executed with client-reported results",
+                "execution_result": {
+                    "success": success,
+                    "output": output,
+                    "error": error,
+                    "error_type": error_type,
+                    "traceback": traceback,
+                },
+            }
+        )
+
+    def test_is_hitl_tool_error_json_with_error(self):
+        """JSON content with failed execution_result should be detected as error."""
+        content = self._make_json_content(
+            success=False,
+            error="KeyError: 'Survived'",
+            error_type="KeyError",
+        )
+        msg = make_tool_msg("jupyter_cell_tool", content)
+        assert _is_hitl_tool_error(msg) is True
+
+    def test_is_hitl_tool_error_json_success(self):
+        """JSON content with successful execution_result should not be error."""
+        content = self._make_json_content(
+            success=True,
+            output="hello\n",
+        )
+        msg = make_tool_msg("jupyter_cell_tool", content)
+        assert _is_hitl_tool_error(msg) is False
+
+    def test_extract_error_parts_json_keyerror(self):
+        """Should extract error from JSON execution_result."""
+        content = self._make_json_content(
+            success=False,
+            error="KeyError: 'Survived'",
+            error_type="KeyError",
+            traceback="Traceback (most recent call last):\n  File ...\nKeyError: 'Survived'",
+        )
+        error_type, error_msg, tb = _extract_error_parts(content)
+        assert error_type == "KeyError"
+        assert "'Survived'" in error_msg
+        assert "Traceback" in tb
+
+    def test_extract_error_parts_json_nameerror(self):
+        """Should extract NameError from JSON execution_result."""
+        content = self._make_json_content(
+            success=False,
+            error="NameError: name 'corr_matrix' is not defined",
+            error_type="NameError",
+            traceback="Traceback...\nNameError: name 'corr_matrix' is not defined",
+        )
+        error_type, error_msg, tb = _extract_error_parts(content)
+        assert error_type == "NameError"
+        assert "corr_matrix" in error_msg
+
+    def test_extract_error_parts_json_success_no_error(self):
+        """Successful JSON content should return no error."""
+        content = self._make_json_content(
+            success=True,
+            output="hello\n",
+        )
+        error_type, error_msg, tb = _extract_error_parts(content)
+        assert error_type is None
+
+    def test_middleware_injects_on_json_error(self):
+        """Middleware should inject re-plan prompt for JSON error content."""
+        content = self._make_json_content(
+            success=False,
+            error="KeyError: 'Survived'",
+            error_type="KeyError",
+            traceback="Traceback...\nKeyError: 'Survived'",
+        )
+        msg = make_tool_msg("jupyter_cell_tool", content)
+
+        def mock_wrap_model_call(fn):
+            return fn
+
+        middleware = create_error_recovery_middleware(mock_wrap_model_call)
+        request = MockRequest(messages=[msg])
+        handler_called = []
+
+        def handler(req):
+            handler_called.append(req)
+            return "ok"
+
+        middleware(request, handler)
+
+        assert len(handler_called) == 1
+        modified_req = handler_called[0]
+        # Should have injected a HumanMessage
+        assert len(modified_req.messages) == 2
+        injected = modified_req.messages[1]
+        assert isinstance(injected, HumanMessage)
+        assert ERROR_RECOVERY_MARKER in injected.content
+        assert "KeyError" in injected.content
+
+    def test_extract_error_parts_re_search_embedded(self):
+        """re.search should find errors embedded in non-standard format."""
+        content = "[Execution Error] NameError: name 'df' is not defined"
+        error_type, error_msg, tb = _extract_error_parts(content)
+        assert error_type == "NameError"
+        assert "df" in error_msg