feat(coding-agent/eval): enabled top-level await execution in Python runner cells

- Added a persistent asyncio event loop and coroutine-aware compiled-code execution for runner cells.
- Enabled compiling notebook cells with top-level await flags and awaiting coroutine results before rendering.
- Added an integration test confirming top-level await works across kernel cells with preserved state.
This commit is contained in:
can1357
2026-05-15 02:12:34 +02:00
parent 92c42f44f7
commit 7cd05c374c
2 changed files with 60 additions and 11 deletions
+42 -11
View File
@@ -25,9 +25,11 @@ when installed.
from __future__ import annotations
import asyncio
import ast
import base64
import builtins
import inspect
import io
import json
import os
@@ -120,6 +122,7 @@ class _RunnerState:
"__builtins__": builtins,
}
self.last_install_marker: int = 0
self.loop: asyncio.AbstractEventLoop | None = None
_STATE = _RunnerState()
@@ -688,13 +691,41 @@ _install_builtins(_STATE.user_ns)
# ---------------------------------------------------------------------------
_TLA_FLAG = getattr(ast, "PyCF_ALLOW_TOP_LEVEL_AWAIT", 0x2000)
def _get_event_loop() -> asyncio.AbstractEventLoop:
loop = _STATE.loop
if loop is None or loop.is_closed():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
_STATE.loop = loop
return loop
def _run_compiled(code, ns: dict, *, want_value: bool) -> Any:
"""Execute a code object, awaiting it if compiled as a coroutine.
``want_value`` is True for the trailing expression — we return ``eval``'s
result (or the awaited coroutine's value). For statement blocks the
return is always ``None``.
"""
if code.co_flags & inspect.CO_COROUTINE:
coro = eval(code, ns)
result = _get_event_loop().run_until_complete(coro)
return result if want_value else None
if want_value:
return eval(code, ns)
exec(code, ns)
return None
def _exec_source(source: str, ns: dict) -> None:
"""Compile + execute ``source``; if the last node is an expression, route
its value through ``__omp_display`` so dataframes/figures render rich."""
try:
module = ast.parse(source, mode="exec")
except SyntaxError:
raise
its value through ``__omp_display`` so dataframes/figures render rich.
Top-level ``await`` / ``async for`` / ``async with`` is permitted; the
cell is driven through the runner's persistent event loop."""
module = ast.parse(source, mode="exec")
if not module.body:
return
@@ -704,16 +735,16 @@ def _exec_source(source: str, ns: dict) -> None:
body_module = ast.Module(body=module.body[:-1], type_ignores=[])
expr_module = ast.Expression(body=last.value)
ast.copy_location(expr_module, last)
body_code = compile(body_module, "<cell>", "exec")
expr_code = compile(expr_module, "<cell>", "eval")
exec(body_code, ns)
value = eval(expr_code, ns)
body_code = compile(body_module, "<cell>", "exec", flags=_TLA_FLAG)
expr_code = compile(expr_module, "<cell>", "eval", flags=_TLA_FLAG)
_run_compiled(body_code, ns, want_value=False)
value = _run_compiled(expr_code, ns, want_value=True)
if value is not None:
__omp_display(value, kind="result")
return
code = compile(module, "<cell>", "exec")
exec(code, ns)
code = compile(module, "<cell>", "exec", flags=_TLA_FLAG)
_run_compiled(code, ns, want_value=False)
# ---------------------------------------------------------------------------
@@ -90,6 +90,24 @@ describe.skipIf(!SHOULD_RUN)("python runner subprocess", () => {
}
});
it("supports top-level await across cells", async () => {
using tempDir = TempDir.createSync("@python-runner-await-");
const kernel = await PythonKernel.start({ cwd: tempDir.path() });
try {
const first = await executePythonWithKernel(
kernel,
["import asyncio", "x = await asyncio.sleep(0, result=21)", "x * 2"].join("\n"),
);
expect(first.exitCode).toBe(0);
expect(first.output).toContain("42");
const second = await executePythonWithKernel(kernel, "x + 1");
expect(second.exitCode).toBe(0);
expect(second.output).toContain("22");
} finally {
await kernel.shutdown();
}
});
it("translates %pwd magic to the user namespace", async () => {
using tempDir = TempDir.createSync("@python-runner-magic-");
const kernel = await PythonKernel.start({ cwd: tempDir.path() });