From 241561165629f1cc35d362098f6847a8c9740db1 Mon Sep 17 00:00:00 2001 From: dmadisetti Date: Tue, 19 May 2026 16:13:10 -0700 Subject: [PATCH] fix: restore the trimmed execution trace --- marimo/_messaging/tracebacks.py | 31 ++---------- marimo/_runtime/executor/executor.py | 8 ++- .../_runtime/runner/hooks_post_execution.py | 5 +- tests/_messaging/test_tracebacks.py | 9 ---- tests/_runtime/test_executor_evaluator.py | 49 +++++++++++++++++++ tests/_server/test_scratchpad_integration.py | 2 +- 6 files changed, 60 insertions(+), 44 deletions(-) diff --git a/marimo/_messaging/tracebacks.py b/marimo/_messaging/tracebacks.py index cb153407758..29ae9669cdf 100644 --- a/marimo/_messaging/tracebacks.py +++ b/marimo/_messaging/tracebacks.py @@ -49,10 +49,8 @@ def write_traceback(traceback: str) -> None: # In run mode, only forward to the frontend if show_tracebacks is on. if in_run_mode and not _show_tracebacks_enabled(): return - # Strip marimo's internal executor.py frame and highlight for the UI - trimmed = _trim_traceback(traceback) sys.stderr._write_with_mimetype( - _highlight_traceback(trimmed), + _highlight_traceback(traceback), mimetype="application/vnd.marimo+traceback", ) else: @@ -64,16 +62,15 @@ def write_traceback(traceback: str) -> None: if in_run_mode and not _show_tracebacks_enabled(): sys.stderr.write(traceback) return - trimmed = _trim_traceback(traceback) broadcast_notification( CellNotification( cell_id=ctx.cell_id, console=CellOutput( channel=CellChannel.STDERR, mimetype="application/vnd.marimo+traceback", - data=trimmed + data=traceback if code_mode - else _highlight_traceback(trimmed), + else _highlight_traceback(traceback), ), ), ctx.stream, @@ -83,27 +80,5 @@ def write_traceback(traceback: str) -> None: sys.stderr.write(traceback) -def _trim_traceback(traceback: str) -> str: - """ - Skip first DefaultExecutor.execute_cell traceback item which all traces start with. - """ - - lines = traceback.split("\n") - if ( - len(lines) > 2 - and lines[0] == "Traceback (most recent call last):" - and ( - '/marimo/_runtime/executor.py", line ' in lines[1] - or '\\marimo\\_runtime\\executor.py", line ' in lines[1] - ) - and lines[1].endswith(", in execute_cell") - ): - for i in range(2, len(lines)): - if lines[i].startswith(" File "): - return "\n".join(lines[:1] + lines[i:]) - - return traceback - - def is_code_highlighting(value: str) -> bool: return 'class="codehilite"' in value diff --git a/marimo/_runtime/executor/executor.py b/marimo/_runtime/executor/executor.py index a4dc6c9254a..da52d6549f9 100644 --- a/marimo/_runtime/executor/executor.py +++ b/marimo/_runtime/executor/executor.py @@ -40,8 +40,10 @@ def execute_cell(self, cell: CellImpl, glbls: dict[str, Any]) -> Any: exec(cell.body, glbls) return eval(cell.last_expr, glbls) except BaseException as e: - # Raising from BaseException folds in the stack trace prior - # to execution; the Runner classifies via ``__cause__``. + # Strip our own frame so user-facing tracebacks start at user + # code. Runners classify via ``__cause__``. + if e.__traceback__ is not None: + e.__traceback__ = e.__traceback__.tb_next raise MarimoRuntimeException from e async def execute_cell_async( @@ -59,4 +61,6 @@ async def execute_cell_async( return await eval(cell.last_expr, glbls) return eval(cell.last_expr, glbls) except BaseException as e: + if e.__traceback__ is not None: + e.__traceback__ = e.__traceback__.tb_next raise MarimoRuntimeException from e diff --git a/marimo/_runtime/runner/hooks_post_execution.py b/marimo/_runtime/runner/hooks_post_execution.py index a2911ede32d..25b6ef3f593 100644 --- a/marimo/_runtime/runner/hooks_post_execution.py +++ b/marimo/_runtime/runner/hooks_post_execution.py @@ -35,7 +35,6 @@ ) from marimo._messaging.tracebacks import ( _highlight_traceback, - _trim_traceback, write_traceback, ) from marimo._messaging.variables import create_variable_value @@ -426,9 +425,7 @@ def _broadcast_outputs( and run_result.exception.__traceback__ ): tb_lines = tb.format_exception(run_result.exception) - formatted_traceback = _highlight_traceback( - _trim_traceback("".join(tb_lines)) - ) + formatted_traceback = _highlight_traceback("".join(tb_lines)) CellNotificationUtils.broadcast_error( data=[ diff --git a/tests/_messaging/test_tracebacks.py b/tests/_messaging/test_tracebacks.py index ddd4fe4350e..00fc5d906d8 100644 --- a/tests/_messaging/test_tracebacks.py +++ b/tests/_messaging/test_tracebacks.py @@ -6,7 +6,6 @@ from marimo._messaging.context import HTTP_REQUEST_CTX, is_code_mode_request from marimo._messaging.tracebacks import ( _highlight_traceback, - _trim_traceback, is_code_highlighting, write_traceback, ) @@ -221,11 +220,3 @@ def test_empty_url(self) -> None: assert is_code_mode_request() is False finally: HTTP_REQUEST_CTX.reset(token) - - def test_trim(self) -> None: - prefix = "Traceback (most recent call last):\n" - head = ' File ".../marimo/_runtime/executor.py", line 139, in execute_cell\n return eval(cell.last_expr, glbls)\n ^^^^^^^^^^^^^^^^^^^^^^^^^^^\n' - rest = ( - ' File ".../__marimo__cell_Hbol_.py", line 2, in \n...\n' - ) - assert _trim_traceback(f"{prefix}{head}{rest}") == f"{prefix}{rest}" diff --git a/tests/_runtime/test_executor_evaluator.py b/tests/_runtime/test_executor_evaluator.py index 3c26bd9592b..da214f64fd3 100644 --- a/tests/_runtime/test_executor_evaluator.py +++ b/tests/_runtime/test_executor_evaluator.py @@ -15,6 +15,8 @@ import asyncio from typing import Any +import pytest + from marimo._runtime.exceptions import MarimoRuntimeException from marimo._runtime.executor import ( DefaultExecutor, @@ -203,6 +205,53 @@ def is_coroutine(self) -> bool: assert isinstance(a.last_run_result.exception, MarimoRuntimeException) +def _cause_traceback_filenames(exc: BaseException) -> list[str]: + cause = exc.__cause__ + assert cause is not None + tb = cause.__traceback__ + files: list[str] = [] + while tb is not None: + files.append(tb.tb_frame.f_code.co_filename) + tb = tb.tb_next + return files + + +def test_default_executor_strips_own_frame_from_cause_sync() -> None: + """``DefaultExecutor.execute_cell`` must not leave its own frame on + the cause's ``__traceback__`` — user-facing tracebacks should begin + at user code (the compiled ```` source).""" + + class _FakeCell: + cell_id = "0" + body = compile("raise ValueError('user bomb')", "", "exec") + last_expr = compile("None", "", "eval") + + with pytest.raises(MarimoRuntimeException) as exc_info: + DefaultExecutor().execute_cell(_FakeCell(), {}) # type: ignore[arg-type] + + files = _cause_traceback_filenames(exc_info.value) + assert files, "cause traceback unexpectedly empty" + assert not any("executor/executor.py" in f for f in files), files + assert files[0] == "" + + +async def test_default_executor_strips_own_frame_from_cause_async() -> None: + """Same as the sync variant, for ``execute_cell_async``.""" + + class _FakeCell: + cell_id = "0" + body = compile("raise ValueError('user bomb')", "", "exec") + last_expr = compile("None", "", "eval") + + with pytest.raises(MarimoRuntimeException) as exc_info: + await DefaultExecutor().execute_cell_async(_FakeCell(), {}) # type: ignore[arg-type] + + files = _cause_traceback_filenames(exc_info.value) + assert files, "cause traceback unexpectedly empty" + assert not any("executor/executor.py" in f for f in files), files + assert files[0] == "" + + async def test_teardown_runs_for_completed_setups_when_later_setup_raises() -> ( None ): diff --git a/tests/_server/test_scratchpad_integration.py b/tests/_server/test_scratchpad_integration.py index b948223856c..86281316f9b 100644 --- a/tests/_server/test_scratchpad_integration.py +++ b/tests/_server/test_scratchpad_integration.py @@ -671,7 +671,7 @@ def test_ctx_create_cell_multiply_defined(session: _Session) -> None: assert lines == snapshot( [ "event: stderr", - 'data: {"data": "Traceback (most recent call last):\\n File \\"/marimo/_runtime/executor.py\\", line N, in execute_cell_async\\n await eval(cell.body, glbls)\\n File \\"\\", line 2, in \\n async with cm.get_context() as ctx:\\n File \\"/marimo/_code_mode/_context.py\\", line N, in __aexit__\\n self._dry_run_compile(ops)\\n File \\"/marimo/_code_mode/_context.py\\", line N, in _dry_run_compile\\n raise RuntimeError(\\nRuntimeError: Multiply-defined names:\\n - \'x\' is already defined in cell \'cell_a\' (cell_a)\\n\\nTo skip validation, use: async with cm.get_context(skip_validation=True) as ctx\\n"}', + 'data: {"data": "Traceback (most recent call last):\\n File \\"\\", line 2, in \\n async with cm.get_context() as ctx:\\n File \\"/marimo/_code_mode/_context.py\\", line N, in __aexit__\\n self._dry_run_compile(ops)\\n File \\"/marimo/_code_mode/_context.py\\", line N, in _dry_run_compile\\n raise RuntimeError(\\nRuntimeError: Multiply-defined names:\\n - \'x\' is already defined in cell \'cell_a\' (cell_a)\\n\\nTo skip validation, use: async with cm.get_context(skip_validation=True) as ctx\\n"}', "", "event: done", 'data: {"success": false, "output": {"mimetype": "text/plain", "data": ""}}',