From 7cd05c374c1a7d693f6dc2bab101207747d74fb9 Mon Sep 17 00:00:00 2001
From: can1357 <me@can.ac>
Date: Fri, 15 May 2026 02:12:34 +0200
Subject: [PATCH] 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.
---
 packages/coding-agent/src/eval/py/runner.py   | 53 +++++++++++++++----
 .../core/python-runner.integration.test.ts    | 18 +++++++
 2 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/packages/coding-agent/src/eval/py/runner.py b/packages/coding-agent/src/eval/py/runner.py
index c947527a5..280590c6f 100644
--- a/packages/coding-agent/src/eval/py/runner.py
+++ b/packages/coding-agent/src/eval/py/runner.py
@@ -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)
 
 
 # ---------------------------------------------------------------------------
diff --git a/packages/coding-agent/test/core/python-runner.integration.test.ts b/packages/coding-agent/test/core/python-runner.integration.test.ts
index 6b5cb86c7..fb82cfc33 100644
--- a/packages/coding-agent/test/core/python-runner.integration.test.ts
+++ b/packages/coding-agent/test/core/python-runner.integration.test.ts
@@ -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() });