diff --git a/mempalace/mcp_server.py b/mempalace/mcp_server.py index 10f9b6a..f257c36 100644 --- a/mempalace/mcp_server.py +++ b/mempalace/mcp_server.py @@ -3036,6 +3036,10 @@ def tool_delete_drawer(drawer_id: str): return {"success": False, "error": str(e)} +class _ProtocolStdoutRestoreFailure(BaseException): + """Fatal loss of the MCP protocol stream after fd-level redirection.""" + + def _capture_fd_stdout(fn): """Run ``fn()`` with its stdout captured at both the Python and fd level. @@ -3063,29 +3067,67 @@ def _capture_fd_stdout(fn): import tempfile buf = io.StringIO() - sys.stdout.flush() - sys.stderr.flush() - try: - saved_fd = os.dup(1) - except (OSError, AttributeError): + + def _capture_python_stdout(): with contextlib.redirect_stdout(buf): result = fn() return result, buf.getvalue() try: - with tempfile.TemporaryFile() as tmp: - os.dup2(tmp.fileno(), 1) + sys.stdout.flush() + sys.stderr.flush() + except (OSError, AttributeError, ValueError): + return _capture_python_stdout() + + try: + saved_fd = os.dup(1) + except (OSError, AttributeError, ValueError): + return _capture_python_stdout() + + redirected = False + try: + try: + tmp_file = tempfile.TemporaryFile() + except (OSError, AttributeError, ValueError): + return _capture_python_stdout() + + with tmp_file as tmp: + try: + os.dup2(tmp.fileno(), 1) + except (OSError, AttributeError, ValueError): + # No callback has run and fd 1 was not replaced. Use the + # documented Python-level fallback. + return _capture_python_stdout() + redirected = True try: with contextlib.redirect_stdout(buf): result = fn() finally: - sys.stdout.flush() - os.dup2(saved_fd, 1) + flush_error = None + try: + sys.stdout.flush() + except (OSError, AttributeError, ValueError) as exc: + flush_error = exc + try: + os.dup2(saved_fd, 1) + except (OSError, AttributeError, ValueError) as exc: + # Ordinary tool and protocol handlers catch Exception. A + # failed restore is process-fatal instead: continuing could + # emit JSON-RPC into the temporary file and hang the client. + # Keep saved_fd open for diagnostics/emergency recovery; + # process exit will release it. + raise _ProtocolStdoutRestoreFailure( + "failed to restore MCP protocol stdout" + ) from exc + redirected = False + if flush_error is not None: + raise flush_error tmp.seek(0) fd_text = tmp.read().decode("utf-8", "replace") return result, buf.getvalue() + fd_text finally: - os.close(saved_fd) + if not redirected: + os.close(saved_fd) def tool_mine( diff --git a/tests/test_mcp_mine.py b/tests/test_mcp_mine.py index 0f73add..d0f4fca 100644 --- a/tests/test_mcp_mine.py +++ b/tests/test_mcp_mine.py @@ -15,8 +15,11 @@ stdout isolation holds. """ import os +import sys +from types import SimpleNamespace import chromadb +import pytest def _patch(monkeypatch, config): @@ -147,6 +150,103 @@ def test_stdout_captured_not_leaked_to_fd(monkeypatch, config, tmp_dir, capfd): assert "Done." not in captured.out +def test_fd_redirect_unavailable_falls_back_to_python_capture(monkeypatch): + """A host that rejects fd-level redirection still gets one safe callback. + + Windows MCP hosts can expose a valid protocol stdout that ``os.dup`` can + copy while rejecting a later ``os.dup2`` to a temporary-file descriptor. + The documented Python-only fallback must cover that setup failure too. + """ + from mempalace import mcp_server + + calls = [] + + def _reject_redirect(_source_fd, _target_fd): + raise OSError(22, "Invalid argument") + + def _callback(): + calls.append("called") + print("python fallback output") + return "result" + + fake_os = SimpleNamespace(dup=os.dup, dup2=_reject_redirect, close=os.close) + monkeypatch.setattr(mcp_server, "os", fake_os) + + result, output = mcp_server._capture_fd_stdout(_callback) + + assert result == "result" + assert calls == ["called"] + assert output == "python fallback output\n" + + +def test_fd_restore_failure_remains_fail_closed(monkeypatch): + """Once fd 1 was redirected, a failed restore becomes a fatal transport error.""" + from mempalace import mcp_server + + calls = 0 + + def _fail_restore(_source_fd, _target_fd): + nonlocal calls + calls += 1 + if calls == 2: + raise OSError(22, "Invalid argument") + + fake_os = SimpleNamespace(dup=os.dup, dup2=_fail_restore, close=os.close) + monkeypatch.setattr(mcp_server, "os", fake_os) + + with pytest.raises(mcp_server._ProtocolStdoutRestoreFailure, match="restore"): + mcp_server._capture_fd_stdout(lambda: print("captured")) + + assert calls == 2 + + +def test_callback_flush_failure_still_restores_fd(monkeypatch): + """A failed post-callback flush cannot skip protocol-fd restoration.""" + from mempalace import mcp_server + + flushes = 0 + dup2_calls = [] + + class _Stdout: + def flush(self): + nonlocal flushes + flushes += 1 + if flushes == 2: + raise OSError(22, "flush failed") + + fake_sys = SimpleNamespace(stdout=_Stdout(), stderr=sys.stderr) + fake_os = SimpleNamespace( + dup=os.dup, + dup2=lambda source_fd, target_fd: dup2_calls.append((source_fd, target_fd)), + close=os.close, + ) + monkeypatch.setattr(mcp_server, "sys", fake_sys) + monkeypatch.setattr(mcp_server, "os", fake_os) + + with pytest.raises(OSError, match="flush failed"): + mcp_server._capture_fd_stdout(lambda: print("captured")) + + assert flushes == 2 + assert len(dup2_calls) == 2 + + +def test_tool_mine_does_not_swallow_fatal_stdout_restore_failure(monkeypatch, config, tmp_dir): + """The normal tool error contract cannot continue after protocol-fd loss.""" + from mempalace import mcp_server + + _patch(monkeypatch, config) + src = os.path.join(tmp_dir, "proj") + os.makedirs(src) + + def _fatal(_fn): + raise mcp_server._ProtocolStdoutRestoreFailure("cannot restore protocol stdout") + + monkeypatch.setattr(mcp_server, "_capture_fd_stdout", _fatal) + + with pytest.raises(mcp_server._ProtocolStdoutRestoreFailure, match="restore"): + mcp_server.tool_mine(source=src, mode="projects", dry_run=True) + + def test_mine_already_running_surfaces_structured_error(monkeypatch, config, tmp_dir): """A held palace lock (MineAlreadyRunning) surfaces as a structured already-running error, mirroring tool_sync."""