better-dsh 0.3.0 → 0.3.1-a

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (76) hide show
  1. package/README.md +49 -177
  2. package/cordis.patch.yml +5 -36
  3. package/docs/20_specs/kernel-provisioning/spec.md +2 -0
  4. package/docs/adr/0015-independent-bundles-agent-remote.md +40 -0
  5. package/docs/remote-tools.md +56 -0
  6. package/docs/specs/kernel-provisioning/spec.md +12 -40
  7. package/docs/specs/omp-python-repl-port/runtime-guide.md +35 -0
  8. package/docs/specs/omp-python-repl-port/spec.md +68 -0
  9. package/docs/specs/remote-tool-enhancement/spec.md +203 -0
  10. package/dsh-docs/baseline.md +13 -0
  11. package/dsh-docs/subsystems/invariants.md +3 -0
  12. package/dsh-docs/subsystems/invariants.zh.md +3 -0
  13. package/eval-description.md +7 -3
  14. package/lib/LaunchOptions-0Mr6meJc.js +4 -0
  15. package/lib/{LaunchOptions-kGXcZOe0.js → LaunchOptions-iE9M-STS.js} +1 -1
  16. package/lib/{bidi-BalpKaJc.js → bidi-CaQeBZHY.js} +1 -1
  17. package/lib/client/index.js +30 -6
  18. package/lib/failover/index.js +1 -1
  19. package/lib/fs-aware/sandbox-plugin.js +3 -3
  20. package/lib/{index-browser-CysrjMyM.js → index-browser-CJzxT9Vi.js} +4 -4
  21. package/lib/index.d.ts +16 -50
  22. package/lib/index.js +482 -1720
  23. package/lib/kernel-env-x_WkG6Fr.js +76 -0
  24. package/lib/kernel-env.d.ts +5 -69
  25. package/lib/kernel-env.js +2 -2
  26. package/lib/{lsp-server-registry-BXvUTBY7.js → lsp-server-registry-BcdkHROD.js} +1 -1
  27. package/lib/{lsp-server-registry-CH9WNFFt.js → lsp-server-registry-D4VLs1Sm.js} +1 -1
  28. package/lib/{main-BQeFrBM1.js → main-taSMz1KX.js} +1 -1
  29. package/lib/{main-CupOhMp5.js → main-uqGqYYnA.js} +7 -5
  30. package/lib/mobile/plugin.d.ts +1 -1
  31. package/lib/omp-python/DILL-LICENSE +35 -0
  32. package/lib/omp-python/OMP-LICENSE +23 -0
  33. package/lib/omp-python/better_snapshot.py +391 -0
  34. package/lib/omp-python/better_support.py +425 -0
  35. package/lib/omp-python/dill-provenance.json +55 -0
  36. package/lib/omp-python/provenance.json +13 -0
  37. package/lib/omp-python/runner.py +2499 -0
  38. package/lib/omp-python/vendor/dill/__diff.py +234 -0
  39. package/lib/omp-python/vendor/dill/__info__.py +291 -0
  40. package/lib/omp-python/vendor/dill/__init__.py +119 -0
  41. package/lib/omp-python/vendor/dill/_dill.py +2265 -0
  42. package/lib/omp-python/vendor/dill/_objects.py +545 -0
  43. package/lib/omp-python/vendor/dill/_shims.py +193 -0
  44. package/lib/omp-python/vendor/dill/detect.py +287 -0
  45. package/lib/omp-python/vendor/dill/logger.py +285 -0
  46. package/lib/omp-python/vendor/dill/objtypes.py +24 -0
  47. package/lib/omp-python/vendor/dill/pointers.py +122 -0
  48. package/lib/omp-python/vendor/dill/session.py +612 -0
  49. package/lib/omp-python/vendor/dill/settings.py +25 -0
  50. package/lib/omp-python/vendor/dill/source.py +1025 -0
  51. package/lib/omp-python/vendor/dill/temp.py +252 -0
  52. package/lib/{puppeteer-core-CSk2PtSm.js → puppeteer-core-C-4NyOpY.js} +8 -8
  53. package/lib/remote/plugin.d.ts +459 -5
  54. package/lib/remote/plugin.js +2389 -210
  55. package/lib/url-schemes/index.d.ts +11 -11
  56. package/lib/url-schemes/index.js +11 -11
  57. package/lib/web-password.d.ts +8 -8
  58. package/lib/web-password.js +5 -1
  59. package/lib/{web-trust-Dej_RI2B.d.ts → web-trust-C_Y6o8Z2.d.ts} +11 -11
  60. package/lib/web-trust.d.ts +1 -2
  61. package/lib/welcome-notice.d.ts +5 -7
  62. package/lib/welcome-notice.js +9 -48
  63. package/lib/{wrap-T4I0nfPM.js → wrap-0VA0uFx8.js} +1 -1
  64. package/package.json +7 -24
  65. package/scripts/kernel-check.mjs +5 -0
  66. package/lib/LaunchOptions-B874eBu2.js +0 -4
  67. package/lib/bundle-delivery.d.ts +0 -20
  68. package/lib/bundle-delivery.js +0 -209
  69. package/lib/hmr.d.ts +0 -32
  70. package/lib/hmr.js +0 -49
  71. package/lib/kernel-bridge.py +0 -172
  72. package/lib/kernel-env-D8UOWeji.js +0 -195
  73. package/lib/webserver.d.ts +0 -10
  74. package/lib/webserver.js +0 -19
  75. package/scripts/kernel-provision.mjs +0 -35
  76. package/scripts/register-bundle.mjs +0 -80
@@ -0,0 +1,2499 @@
1
+ """OMP Python runner — subprocess wrapper used by the coding-agent host.
2
+
3
+ NDJSON protocol over stdin/stdout. Host writes one JSON object per line;
4
+ wrapper writes typed frames back.
5
+
6
+ Host -> wrapper:
7
+ {"id": str, "code": str, "silent": bool?, "storeHistory": bool?}
8
+ {"id": str, "code": str, "silent": bool?, "storeHistory": bool?, "cwd": str?, "filename": str?, "env": dict?}
9
+ {"type": "tool", "id": str, "op": "describe", "names": [str]}
10
+ {"type": "tool", "id": str, "op": "call", "name": str, "args": dict}
11
+ {"type": "exit"} # graceful shutdown
12
+
13
+ Wrapper -> host:
14
+ {"type": "started", "id": ...}
15
+ {"type": "stdout", "id": ..., "data": str}
16
+ {"type": "stderr", "id": ..., "data": str}
17
+ {"type": "display", "id": ..., "bundle": {<mime>: <value>}}
18
+ {"type": "result", "id": ..., "bundle": {<mime>: <value>}}
19
+ {"type": "error", "id": ..., "ename": str, "evalue": str, "traceback": [str]}
20
+ {"type": "done", "id": ..., "status": "ok"|"error",
21
+ "executionCount": int, "cancelled": bool}
22
+
23
+ The runner is intentionally self-contained: no third-party imports, no IPython.
24
+ Magics are translated by a small line-scanner before AST parsing; rich display
25
+ falls back through `_repr_*_` methods so pandas/PIL/plotly etc. still render
26
+ when installed.
27
+ """
28
+
29
+ from __future__ import annotations
30
+
31
+ import ast
32
+ import asyncio
33
+ import base64
34
+ import builtins
35
+ import codecs
36
+ import contextvars
37
+ import inspect
38
+ import io
39
+ import itertools
40
+ import json
41
+ import hashlib
42
+ import linecache
43
+ import locale
44
+ import os
45
+ import re
46
+ import runpy
47
+ import shlex
48
+ import shutil
49
+ import signal
50
+ import subprocess
51
+ import sys
52
+ import threading
53
+ import time
54
+ import traceback
55
+ import tokenize
56
+ from pathlib import Path
57
+ from typing import Any, Callable
58
+
59
+ # ---------------------------------------------------------------------------
60
+ # Frame writer
61
+ # ---------------------------------------------------------------------------
62
+
63
+ # Frames travel on a private dup of the original stdout. fd 1 itself is then
64
+ # repointed at a capture pipe: child processes spawned by user code without
65
+ # stdout=PIPE inherit fd 1, and their output is forwarded to the host as
66
+ # regular stdout frames by a drain thread instead of being written raw into
67
+ # the NDJSON channel (where it would be dropped as invalid JSON — or worse,
68
+ # spoof a frame). The wire protocol is unchanged: the host still reads NDJSON
69
+ # frames from the subprocess stdout.
70
+ _RAW_STDERR = sys.__stderr__
71
+ try:
72
+ _FRAME_FD = os.dup(sys.__stdout__.fileno())
73
+ _RAW_STDOUT = os.fdopen(_FRAME_FD, "w", encoding="utf-8", errors="backslashreplace")
74
+ _CAPTURE_READ_FD, _capture_write_fd = os.pipe()
75
+ os.dup2(_capture_write_fd, sys.__stdout__.fileno())
76
+ os.close(_capture_write_fd)
77
+ except (AttributeError, OSError, ValueError, io.UnsupportedOperation):
78
+ _RAW_STDOUT = sys.__stdout__
79
+ _CAPTURE_READ_FD = None
80
+ _OUT_LOCK = threading.Lock()
81
+
82
+
83
+ def _json_default(o: Any) -> Any:
84
+ try:
85
+ return repr(o)
86
+ except Exception:
87
+ return f"<unrepr {type(o).__name__}>"
88
+
89
+
90
+ def _emit(frame: dict) -> None:
91
+ """Serialize a frame and write it to the host as a single NDJSON line."""
92
+ line = json.dumps(frame, ensure_ascii=False, default=_json_default)
93
+ with _OUT_LOCK:
94
+ _RAW_STDOUT.write(line)
95
+ _RAW_STDOUT.write("\n")
96
+ _RAW_STDOUT.flush()
97
+
98
+
99
+ # ---------------------------------------------------------------------------
100
+ # User stdout/stderr proxies
101
+ # ---------------------------------------------------------------------------
102
+
103
+
104
+ class _StreamProxy(io.TextIOBase):
105
+ """Emit ``write()`` data as typed frames tied to the current request.
106
+
107
+ Writes are coalesced per request: a frame is emitted once the buffer holds
108
+ a complete line (everything up to the last newline goes out together) or
109
+ grows past ``_MAX_BUFFER`` bytes, so the common ``print()`` pair of
110
+ ``write(text)`` + ``write("\\n")`` costs one frame instead of two. Partial
111
+ lines are bounded by ``flush()`` and the end-of-request flush.
112
+ """
113
+
114
+ _MAX_BUFFER = 8192
115
+
116
+ def __init__(self, kind: str) -> None:
117
+ super().__init__()
118
+ self._kind = kind
119
+ self._lock = threading.Lock()
120
+ self._buffers: dict[str, str] = {}
121
+
122
+ def writable(self) -> bool: # noqa: D401 - protocol method
123
+ return True
124
+
125
+ def isatty(self) -> bool: # noqa: D401 - protocol method
126
+ return False
127
+
128
+ def write(self, data: Any) -> int: # type: ignore[override]
129
+ if not isinstance(data, str):
130
+ data = str(data)
131
+ if not data:
132
+ return 0
133
+ rid = _CURRENT_RID.get()
134
+ if rid is None:
135
+ _RAW_STDERR.write(data)
136
+ _RAW_STDERR.flush()
137
+ return len(data)
138
+ emit_text = None
139
+ with self._lock:
140
+ buf = self._buffers.pop(rid, "") + data
141
+ if len(buf) >= self._MAX_BUFFER:
142
+ emit_text = buf
143
+ else:
144
+ nl = buf.rfind("\n")
145
+ if nl >= 0:
146
+ emit_text = buf[: nl + 1]
147
+ rest = buf[nl + 1 :]
148
+ if rest:
149
+ self._buffers[rid] = rest
150
+ else:
151
+ self._buffers[rid] = buf
152
+ if emit_text:
153
+ _emit({"type": self._kind, "id": rid, "data": emit_text})
154
+ return len(data)
155
+
156
+ def flush(self) -> None: # noqa: D401 - protocol method
157
+ rid = _CURRENT_RID.get()
158
+ if rid is not None:
159
+ self.flush_rid(rid)
160
+ return None
161
+
162
+ def flush_rid(self, rid: str) -> None:
163
+ """Flush any buffered partial line for ``rid`` as its own frame."""
164
+ with self._lock:
165
+ buf = self._buffers.pop(rid, None)
166
+ if buf:
167
+ _emit({"type": self._kind, "id": rid, "data": buf})
168
+
169
+
170
+ def _flush_stream_proxies(rid: str) -> None:
171
+ """Drain buffered proxy output for ``rid`` (called before its done frame)."""
172
+ for stream in (sys.stdout, sys.stderr):
173
+ if isinstance(stream, _StreamProxy):
174
+ stream.flush_rid(rid)
175
+
176
+
177
+ # ---------------------------------------------------------------------------
178
+ # Runner state
179
+ # ---------------------------------------------------------------------------
180
+
181
+
182
+ class _RunnerState:
183
+ def __init__(self) -> None:
184
+ self.execution_count: int = 0
185
+ self.namespace_revision: int = 0
186
+ self.cancel_requested: bool = False
187
+ # User globals — kept across requests when running in session mode.
188
+ self.user_ns: dict[str, Any] = {
189
+ "__name__": "__main__",
190
+ "__doc__": None,
191
+ "__builtins__": builtins,
192
+ }
193
+ self.last_install_marker: int = 0
194
+ self.loop: asyncio.AbstractEventLoop | None = None
195
+ self.active_executions: int = 0
196
+ # Best-effort attribution target for captured fd-1 bytes (child
197
+ # processes inheriting stdout). With overlapping requests the most
198
+ # recently started one wins — strictly better than dropping the bytes.
199
+ self.capture_rid: str | None = None
200
+
201
+
202
+ _CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar(
203
+ "omp_current_rid", default=None
204
+ )
205
+ _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS: contextvars.ContextVar[set[int] | None] = (
206
+ contextvars.ContextVar(
207
+ "omp_displayed_matplotlib_figure_ids",
208
+ default=None,
209
+ )
210
+ )
211
+
212
+
213
+ _STATE = _RunnerState()
214
+
215
+ _SHADOW_SNAPSHOT_MAX_DEPTH = 16
216
+ _SHADOW_SNAPSHOT_MAX_NODES = 2000
217
+ _SHADOW_SNAPSHOT_MAX_STRING_BYTES = 8 * 1024 * 1024
218
+ _SHADOW_UNSUPPORTED = object()
219
+
220
+ # Class attribute on the prelude `_ToolProxy` identifying the genuine tool
221
+ # bridge. A retained non-JSON-safe user `tool` binding is omitted from the
222
+ # shadow snapshot exactly like the genuine bridge, so snapshot absence alone
223
+ # cannot tell them apart (see `_shadow_tool_available`).
224
+ _TOOL_BRIDGE_MARKER = "__omp_tool_bridge__"
225
+
226
+ # Runner-owned reference to the genuine prelude tool bridge, captured on first
227
+ # sighting: kernel init executes the prelude before user cells, so the first
228
+ # `tool` binding observed by the planner is the installed bridge. Identity
229
+ # (`is`) plus the marker above defeats both retained shadows and spoofed
230
+ # markers; the residual is a hostile `%run` script that both runs before any
231
+ # plan and reproduces the marker.
232
+ _TRUSTED_TOOL_BRIDGE: dict[str, Any] = {}
233
+
234
+
235
+ def _tool_binding_state(user_ns: dict[str, Any]) -> str:
236
+ """Reserved `tool` binding identity: absent, intact, or shadowed.
237
+
238
+ Absent with nothing ever captured is the bare-harness norm (no bridge
239
+ installed at all). Absent after a capture means user deletion. Present
240
+ means intact only for the captured genuine bridge carrying the marker.
241
+ """
242
+ if "tool" not in user_ns:
243
+ return "absent" if _TRUSTED_TOOL_BRIDGE.get("bridge") is None else "shadowed"
244
+ current = user_ns["tool"]
245
+ trusted = _TRUSTED_TOOL_BRIDGE.get("bridge")
246
+ if trusted is None:
247
+ if getattr(current, _TOOL_BRIDGE_MARKER, False) is True:
248
+ _TRUSTED_TOOL_BRIDGE["bridge"] = current
249
+ return "intact"
250
+ return "shadowed"
251
+ return "intact" if current is trusted else "shadowed"
252
+
253
+
254
+ def _shadow_identity_state() -> dict[str, str]:
255
+ """Reserved-binding identity relevant to speculation, for the snapshot digest.
256
+
257
+ Out-of-band namespace changes (async tasks, background threads) that replace
258
+ `tool` or `str` may leave no JSON-safe trace and bump no revision; without
259
+ this the digest would still match and stale plans would verify. Predicates
260
+ mirror the planner gates exactly (`str` membership matches the seed in
261
+ `_emit_shadow_plan`; `tool` matches `_shadow_tool_available`).
262
+ """
263
+ user_ns = _STATE.user_ns
264
+ return {
265
+ "str": "shadowed" if "str" in user_ns else "absent",
266
+ "tool": _tool_binding_state(user_ns),
267
+ }
268
+
269
+ def _copy_shadow_value(value: Any, depth: int, state: dict[str, Any]) -> Any:
270
+ """Copy exact JSON-safe values without invoking user protocols."""
271
+ if depth > _SHADOW_SNAPSHOT_MAX_DEPTH:
272
+ return _SHADOW_UNSUPPORTED
273
+ state["nodes"] += 1
274
+ if state["nodes"] > _SHADOW_SNAPSHOT_MAX_NODES:
275
+ return _SHADOW_UNSUPPORTED
276
+ value_type = type(value)
277
+ if value is None or value_type is bool or value_type is int:
278
+ return value
279
+ if value_type is float:
280
+ return value if value == value and value not in (float("inf"), float("-inf")) else _SHADOW_UNSUPPORTED
281
+ if value_type is str:
282
+ state["bytes"] += len(value.encode("utf-8"))
283
+ return value if state["bytes"] <= _SHADOW_SNAPSHOT_MAX_STRING_BYTES else _SHADOW_UNSUPPORTED
284
+ if value_type not in (list, tuple, dict) or id(value) in state["seen"]:
285
+ return _SHADOW_UNSUPPORTED
286
+ state["seen"].add(id(value))
287
+ try:
288
+ if value_type in (list, tuple):
289
+ copied = [_copy_shadow_value(item, depth + 1, state) for item in value]
290
+ return copied if all(item is not _SHADOW_UNSUPPORTED for item in copied) else _SHADOW_UNSUPPORTED
291
+ copied_dict: dict[str, Any] = {}
292
+ for key, item in value.items():
293
+ if type(key) is not str:
294
+ return _SHADOW_UNSUPPORTED
295
+ copied = _copy_shadow_value(item, depth + 1, state)
296
+ if copied is _SHADOW_UNSUPPORTED:
297
+ return _SHADOW_UNSUPPORTED
298
+ copied_dict[key] = copied
299
+ return copied_dict
300
+ finally:
301
+ state["seen"].discard(id(value))
302
+
303
+
304
+ def _snapshot_user_namespace() -> dict[str, Any]:
305
+ values: dict[str, Any] = {}
306
+ snapshot_state: dict[str, Any] = {"nodes": 0, "bytes": 0, "seen": set()}
307
+ for key, value in _STATE.user_ns.items():
308
+ if key.startswith("_"):
309
+ continue
310
+ copied = _copy_shadow_value(value, 0, snapshot_state)
311
+ if copied is not _SHADOW_UNSUPPORTED:
312
+ values[key] = copied
313
+ return values
314
+
315
+
316
+ def _shadow_tool_available(snapshot: dict[str, Any], user_ns: dict[str, Any]) -> bool:
317
+ """Whether a speculative `tool.read` may reach the genuine bridge.
318
+
319
+ A JSON-safe user value is visible in the snapshot; anything else goes
320
+ through the binding identity, which also covers bare harnesses ("absent"
321
+ with nothing ever captured stays available as before).
322
+ """
323
+ if "tool" in snapshot:
324
+ return False
325
+ # A JSON-safe user value is visible above; anything else goes through the
326
+ # binding identity, which also covers bare harnesses ("absent" with nothing
327
+ # ever captured stays available as before).
328
+ return _tool_binding_state(user_ns) in ("intact", "absent")
329
+
330
+ def _shadow_snapshot_digest(values: dict[str, Any]) -> str:
331
+ payload = json.dumps(
332
+ {"identity": _shadow_identity_state(), "values": values},
333
+ ensure_ascii=False,
334
+ separators=(",", ":"),
335
+ sort_keys=True,
336
+ )
337
+ return hashlib.sha256(payload.encode("utf-8")).hexdigest()
338
+
339
+
340
+ def _emit_shadow_snapshot(req: dict) -> None:
341
+ rid = str(req.get("id"))
342
+ if _STATE.active_executions != 0:
343
+ _emit({"type": "shadow_snapshot", "id": rid, "eligible": False, "reason": "kernel is busy"})
344
+ return
345
+ values = _snapshot_user_namespace()
346
+ _emit(
347
+ {
348
+ "type": "shadow_snapshot",
349
+ "id": rid,
350
+ "eligible": True,
351
+ "revision": _STATE.namespace_revision,
352
+ "digest": _shadow_snapshot_digest(values),
353
+ "values": values,
354
+ }
355
+ )
356
+
357
+
358
+ def _shadow_span(node: ast.AST, line_offsets: list[int]) -> dict[str, int]:
359
+ start_line = max(1, int(getattr(node, "lineno", 1)))
360
+ end_line = max(start_line, int(getattr(node, "end_lineno", start_line)))
361
+ start = line_offsets[start_line - 1] + int(getattr(node, "col_offset", 0))
362
+ end = line_offsets[end_line - 1] + int(getattr(node, "end_col_offset", 0))
363
+ return {"start": start, "end": end}
364
+
365
+
366
+ def _shadow_dependencies(expression: dict[str, Any], output: set[str] | None = None) -> set[str]:
367
+ output = output if output is not None else set()
368
+ kind = expression.get("kind")
369
+ if kind == "operation_result":
370
+ output.add(str(expression["operationId"]))
371
+ elif kind == "property":
372
+ _shadow_dependencies(expression["target"], output)
373
+ elif kind in ("array", "concat"):
374
+ for item in expression["items"]:
375
+ _shadow_dependencies(item, output)
376
+ elif kind == "object":
377
+ for entry in expression["entries"]:
378
+ _shadow_dependencies(entry["value"], output)
379
+ elif kind == "transform":
380
+ _shadow_dependencies(expression["input"], output)
381
+ if expression.get("argument") is not None:
382
+ _shadow_dependencies(expression["argument"], output)
383
+ return output
384
+
385
+
386
+
387
+
388
+ def _shadow_expression_is_string(expression: dict[str, Any]) -> bool:
389
+ kind = expression.get("kind")
390
+ if kind == "literal":
391
+ return type(expression.get("value")) is str
392
+ if kind == "concat":
393
+ return True
394
+ return kind == "transform" and expression.get("name") in (
395
+ "Python.str",
396
+ "JSON.stringify",
397
+ "Array.join",
398
+ )
399
+
400
+
401
+ def _shadow_expression(
402
+ node: ast.AST,
403
+ environment: dict[str, dict[str, Any]],
404
+ snapshot: dict[str, Any] | None = None,
405
+ ) -> dict[str, Any] | None:
406
+ if isinstance(node, ast.Constant) and (
407
+ node.value is None or type(node.value) in (bool, int, float, str)
408
+ ):
409
+ if type(node.value) is float and (
410
+ node.value != node.value or node.value in (float("inf"), float("-inf"))
411
+ ):
412
+ return None
413
+ return {"kind": "literal", "value": node.value}
414
+ if isinstance(node, ast.Name):
415
+ return environment.get(node.id, {"kind": "snapshot", "name": node.id})
416
+ if isinstance(node, (ast.List, ast.Tuple)):
417
+ items = [_shadow_expression(item, environment, snapshot) for item in node.elts]
418
+ return (
419
+ {"kind": "array", "items": items}
420
+ if all(item is not None for item in items)
421
+ else None
422
+ )
423
+ if isinstance(node, ast.Dict):
424
+ entries: list[dict[str, Any]] = []
425
+ for key, value in zip(node.keys, node.values):
426
+ if not isinstance(key, ast.Constant) or type(key.value) is not str:
427
+ return None
428
+ projected = _shadow_expression(value, environment, snapshot)
429
+ if projected is None:
430
+ return None
431
+ entries.append({"key": key.value, "value": projected})
432
+ return {"kind": "object", "entries": entries}
433
+ if isinstance(node, ast.Subscript):
434
+ target = _shadow_expression(node.value, environment, snapshot)
435
+ key = _shadow_expression(node.slice, environment, snapshot)
436
+ if target is None or key is None or key.get("kind") != "literal":
437
+ return None
438
+ if type(key.get("value")) not in (str, int):
439
+ return None
440
+ return {"kind": "property", "target": target, "property": key["value"]}
441
+ if isinstance(node, ast.JoinedStr):
442
+ items: list[dict[str, Any]] = []
443
+ for value in node.values:
444
+ if isinstance(value, ast.Constant) and type(value.value) is str:
445
+ items.append({"kind": "literal", "value": value.value})
446
+ elif isinstance(value, ast.FormattedValue):
447
+ if value.conversion != -1 or value.format_spec is not None:
448
+ return None
449
+ projected = _shadow_expression(value.value, environment, snapshot)
450
+ if projected is None:
451
+ return None
452
+ items.append({"kind": "transform", "name": "Python.str", "input": projected})
453
+ else:
454
+ return None
455
+ return {"kind": "concat", "items": items}
456
+ if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
457
+ left = _shadow_expression(node.left, environment, snapshot)
458
+ right = _shadow_expression(node.right, environment, snapshot)
459
+ return (
460
+ {"kind": "concat", "items": [left, right]}
461
+ if left is not None
462
+ and right is not None
463
+ and _shadow_expression_is_string(left)
464
+ and _shadow_expression_is_string(right)
465
+ else None
466
+ )
467
+ if isinstance(node, ast.Call) and not node.keywords:
468
+ if (
469
+ isinstance(node.func, ast.Name)
470
+ and node.func.id == "str"
471
+ and node.func.id not in environment
472
+ and (snapshot is None or "str" not in snapshot)
473
+ and len(node.args) == 1
474
+ ):
475
+ projected = _shadow_expression(node.args[0], environment, snapshot)
476
+ return (
477
+ {"kind": "transform", "name": "Python.str", "input": projected}
478
+ if projected is not None
479
+ else None
480
+ )
481
+ return None
482
+
483
+
484
+ def _shadow_static_value(
485
+ expression: dict[str, Any], snapshot: dict[str, Any]
486
+ ) -> tuple[bool, Any]:
487
+ kind = expression.get("kind")
488
+ if kind == "literal":
489
+ return True, expression.get("value")
490
+ if kind == "snapshot":
491
+ name = str(expression.get("name"))
492
+ return (True, snapshot[name]) if name in snapshot else (False, None)
493
+ if kind == "array":
494
+ values = [_shadow_static_value(item, snapshot) for item in expression["items"]]
495
+ return (
496
+ (True, [value for _, value in values])
497
+ if all(ok for ok, _ in values)
498
+ else (False, None)
499
+ )
500
+ if kind == "object":
501
+ values = [
502
+ (entry["key"], _shadow_static_value(entry["value"], snapshot))
503
+ for entry in expression["entries"]
504
+ ]
505
+ return (
506
+ (True, {key: value for key, (_, value) in values})
507
+ if all(ok for _, (ok, _) in values)
508
+ else (False, None)
509
+ )
510
+ if kind == "property":
511
+ ok, target = _shadow_static_value(expression["target"], snapshot)
512
+ if not ok:
513
+ return False, None
514
+ property_name = expression["property"]
515
+ if type(target) is dict and type(property_name) is str and property_name in target:
516
+ return True, target[property_name]
517
+ if type(target) is list and type(property_name) is int and 0 <= property_name < len(target):
518
+ return True, target[property_name]
519
+ return False, None
520
+ if kind == "concat":
521
+ values = [_shadow_static_value(item, snapshot) for item in expression["items"]]
522
+ return (
523
+ (True, "".join(str(value) for _, value in values))
524
+ if all(ok for ok, _ in values)
525
+ else (False, None)
526
+ )
527
+ return False, None
528
+
529
+
530
+ def _shadow_call_kind(node: ast.AST, tool_available: bool = True) -> str | None:
531
+ if isinstance(node, ast.Await):
532
+ node = node.value
533
+ if not isinstance(node, ast.Call):
534
+ return None
535
+ if (
536
+ isinstance(node.func, ast.Attribute)
537
+ and isinstance(node.func.value, ast.Name)
538
+ and node.func.value.id == "tool"
539
+ and node.func.attr == "read"
540
+ and tool_available
541
+ ):
542
+ return "read"
543
+ return None
544
+
545
+
546
+ def _emit_shadow_plan(req: dict) -> None:
547
+ rid = str(req.get("id"))
548
+ if _STATE.active_executions != 0:
549
+ _emit({"type": "shadow_plan", "id": rid, "eligible": False, "reason": "kernel is busy"})
550
+ return
551
+ code = req.get("code")
552
+ if type(code) is not str:
553
+ _emit({"type": "shadow_plan", "id": rid, "eligible": False, "reason": "code is not a string"})
554
+ return
555
+ try:
556
+ module = ast.parse(code, mode="exec")
557
+ except SyntaxError:
558
+ _emit({"type": "shadow_plan", "id": rid, "eligible": False, "reason": "incomplete or invalid Python"})
559
+ return
560
+ snapshot = _snapshot_user_namespace()
561
+ try:
562
+ # Whole-cell validity gate: `ast.parse` accepts cells that `compile()`
563
+ # rejects (e.g. use-before-global), while authoritative
564
+ # `_compile_source()` raises SyntaxError before any bridge call.
565
+ # Mirror its filename/mode/flags (top-level await stays legal) and
566
+ # fail closed with zero operations plus an invalidating barrier.
567
+ compile(module, "<cell>", "exec", flags=_TLA_FLAG)
568
+ except SyntaxError:
569
+ _emit(
570
+ {
571
+ "type": "shadow_plan",
572
+ "id": rid,
573
+ "eligible": True,
574
+ "revision": _STATE.namespace_revision,
575
+ "digest": _shadow_snapshot_digest(snapshot),
576
+ "values": snapshot,
577
+ "operations": [],
578
+ "controls": [],
579
+ "barrier": {
580
+ "kind": "barrier",
581
+ "reason": "cell does not compile",
582
+ "span": {"start": 0, "end": len(code)},
583
+ },
584
+ }
585
+ )
586
+ return
587
+ # Mirror the `_compile_source` whole-cell instrumentation skip: when the
588
+ # cell binds `__omp_with_call_site__` anywhere, authoritative execution
589
+ # runs zero instrumented calls, so the planner must admit zero operations
590
+ # (per-call scoping is undecidable across call order; fail closed).
591
+ if _cell_binds_call_site_helper(module):
592
+ _emit(
593
+ {
594
+ "type": "shadow_plan",
595
+ "id": rid,
596
+ "eligible": True,
597
+ "revision": _STATE.namespace_revision,
598
+ "digest": _shadow_snapshot_digest(snapshot),
599
+ "values": snapshot,
600
+ "operations": [],
601
+ "controls": [],
602
+ "barrier": {
603
+ "kind": "barrier",
604
+ "reason": "Python call-site helper binding changed",
605
+ "span": {"start": 0, "end": len(code)},
606
+ },
607
+ }
608
+ )
609
+ return
610
+ tool_available = _shadow_tool_available(snapshot, _STATE.user_ns)
611
+ line_offsets = [0]
612
+ for line in code.splitlines(keepends=True):
613
+ line_offsets.append(line_offsets[-1] + len(line))
614
+ operations: list[dict[str, Any]] = []
615
+ controls: list[dict[str, Any]] = []
616
+ environment: dict[str, dict[str, Any]] = {}
617
+ # A retained `str` binding shadows the builtin even when the JSON-safe snapshot
618
+ # omits it (function values are never JSON-safe), so snapshot absence alone
619
+ # cannot prove the builtin is intact. Seed the same snapshot-ref shape Name
620
+ # resolution falls back to: the `str(...)` call gate below then fails closed,
621
+ # and any other use resolves unresolvable. (`del str` needs no seed: removing
622
+ # the binding restores builtin resolution authoritatively too.)
623
+ if "str" in _STATE.user_ns:
624
+ environment["str"] = {"kind": "snapshot", "name": "str"}
625
+ occurrences: dict[str, int] = {}
626
+ source_order = 0
627
+ barrier: dict[str, Any] | None = None
628
+
629
+ def add_operation(
630
+ expression_node: ast.AST,
631
+ dynamic_path: list[str],
632
+ control_dependencies: list[str],
633
+ ) -> dict[str, Any] | None:
634
+ nonlocal source_order
635
+ # Unawaited `tool.read(...)` never reaches the bridge: the kernel tool
636
+ # attribute is async, so a bare call only builds a coroutine. Admit an
637
+ # operation only for the awaited form; anything else fails closed via
638
+ # the caller's barrier path below.
639
+ if not isinstance(expression_node, ast.Await):
640
+ return None
641
+ call_node = expression_node.value
642
+ kind = _shadow_call_kind(call_node, tool_available)
643
+ if not isinstance(call_node, ast.Call) or kind != "read":
644
+ return None
645
+ if call_node.keywords or len(call_node.args) != 1:
646
+ return None
647
+ argument_ir = _shadow_expression(call_node.args[0], environment, snapshot)
648
+ if argument_ir is None:
649
+ return None
650
+ call_span = _shadow_span(call_node, line_offsets)
651
+ static_site = f"py:{call_span['start']}"
652
+ path_key = f"{static_site}:{'/'.join(dynamic_path)}"
653
+ occurrence = occurrences.get(path_key, 0)
654
+ occurrences[path_key] = occurrence + 1
655
+ operation_id = f"{path_key}:{occurrence}"
656
+ operation = {
657
+ "kind": "tool",
658
+ "call": {
659
+ "id": operation_id,
660
+ "siteId": static_site,
661
+ "dynamicPath": list(dynamic_path),
662
+ "occurrence": occurrence,
663
+ "name": kind,
664
+ "args": argument_ir,
665
+ "dependencies": sorted(_shadow_dependencies(argument_ir)),
666
+ "controlDependencies": list(control_dependencies),
667
+ "sourceOrder": source_order,
668
+ "span": call_span,
669
+ },
670
+ }
671
+ source_order += 1
672
+ operations.append(operation)
673
+ return operation
674
+
675
+ def project_statements(
676
+ statements: list[ast.stmt],
677
+ dynamic_path: list[str],
678
+ control_dependencies: list[str],
679
+ ) -> bool:
680
+ nonlocal barrier
681
+ for statement in statements:
682
+ if isinstance(statement, (ast.Assign, ast.AnnAssign)):
683
+ target = (
684
+ statement.targets[0]
685
+ if isinstance(statement, ast.Assign) and len(statement.targets) == 1
686
+ else getattr(statement, "target", None)
687
+ )
688
+ value_node = statement.value
689
+ if not isinstance(target, ast.Name) or value_node is None:
690
+ barrier = {
691
+ "kind": "barrier",
692
+ "reason": "unsupported Python assignment",
693
+ "span": _shadow_span(statement, line_offsets),
694
+ }
695
+ return False
696
+ if target.id == "tool":
697
+ barrier = {
698
+ "kind": "barrier",
699
+ "reason": "Python tool binding changed",
700
+ "span": _shadow_span(statement, line_offsets),
701
+ }
702
+ return False
703
+ operation = add_operation(value_node, dynamic_path, control_dependencies)
704
+ if operation is not None:
705
+ environment[target.id] = {"kind": "operation_result", "operationId": operation["call"]["id"]}
706
+ continue
707
+ projected = _shadow_expression(value_node, environment, snapshot)
708
+ if projected is None:
709
+ barrier = {
710
+ "kind": "barrier",
711
+ "reason": "unsupported Python assignment value",
712
+ "span": _shadow_span(value_node, line_offsets),
713
+ }
714
+ return False
715
+ if not _shadow_static_value(projected, snapshot)[0]:
716
+ barrier = {
717
+ "kind": "barrier",
718
+ "reason": "unsupported Python assignment value",
719
+ "span": _shadow_span(value_node, line_offsets),
720
+ }
721
+ return False
722
+ environment[target.id] = projected
723
+ continue
724
+ if isinstance(statement, ast.Expr):
725
+ if add_operation(statement.value, dynamic_path, control_dependencies) is not None:
726
+ continue
727
+ projected = _shadow_expression(statement.value, environment, snapshot)
728
+ if projected is not None and _shadow_static_value(projected, snapshot)[0]:
729
+ continue
730
+ barrier = {
731
+ "kind": "barrier",
732
+ "reason": "unsupported Python statement",
733
+ "span": _shadow_span(statement, line_offsets),
734
+ }
735
+ return False
736
+ if isinstance(statement, ast.If):
737
+ test = _shadow_expression(statement.test, environment, snapshot)
738
+ if test is None:
739
+ barrier = {
740
+ "kind": "barrier",
741
+ "reason": "unsupported Python condition",
742
+ "span": _shadow_span(statement.test, line_offsets),
743
+ }
744
+ return False
745
+ ok, selected = _shadow_static_value(test, snapshot)
746
+ conditional_id = f"py:{_shadow_span(statement, line_offsets)['start']}:if"
747
+ if ok:
748
+ branch = statement.body if selected else statement.orelse
749
+ if not project_statements(
750
+ branch,
751
+ [*dynamic_path, "if:true" if selected else "if:false"],
752
+ control_dependencies,
753
+ ):
754
+ return False
755
+ continue
756
+ controls.append(
757
+ {
758
+ "kind": "conditional",
759
+ "id": conditional_id,
760
+ "test": test,
761
+ "consequentPath": "if:true",
762
+ "alternatePath": "if:false",
763
+ "span": _shadow_span(statement, line_offsets),
764
+ }
765
+ )
766
+ if not project_statements(
767
+ statement.body,
768
+ [*dynamic_path, "if:true"],
769
+ [*control_dependencies, conditional_id],
770
+ ):
771
+ return False
772
+ if not project_statements(
773
+ statement.orelse,
774
+ [*dynamic_path, "if:false"],
775
+ [*control_dependencies, conditional_id],
776
+ ):
777
+ return False
778
+ continue
779
+ if isinstance(statement, ast.For) and isinstance(statement.target, ast.Name):
780
+ if statement.target.id == "tool":
781
+ barrier = {
782
+ "kind": "barrier",
783
+ "reason": "Python tool binding changed",
784
+ "span": _shadow_span(statement.target, line_offsets),
785
+ }
786
+ return False
787
+ iterable = _shadow_expression(statement.iter, environment, snapshot)
788
+ ok, values = _shadow_static_value(iterable, snapshot) if iterable is not None else (False, None)
789
+ if not ok or type(values) is not list or len(values) > 32:
790
+ barrier = {
791
+ "kind": "barrier",
792
+ "reason": "unbounded or dynamic Python loop",
793
+ "span": _shadow_span(statement, line_offsets),
794
+ }
795
+ return False
796
+ controls.append(
797
+ {
798
+ "kind": "loop",
799
+ "id": f"py:{_shadow_span(statement, line_offsets)['start']}:loop",
800
+ "iterable": iterable,
801
+ "iterations": len(values),
802
+ "span": _shadow_span(statement, line_offsets),
803
+ }
804
+ )
805
+ previous = environment.get(statement.target.id)
806
+ for index, value in enumerate(values):
807
+ environment[statement.target.id] = {"kind": "literal", "value": value}
808
+ if not project_statements(
809
+ statement.body,
810
+ [*dynamic_path, f"loop:{index}"],
811
+ control_dependencies,
812
+ ):
813
+ return False
814
+ if values:
815
+ environment[statement.target.id] = {"kind": "literal", "value": values[-1]}
816
+ elif previous is None:
817
+ environment.pop(statement.target.id, None)
818
+ else:
819
+ environment[statement.target.id] = previous
820
+ if statement.orelse and not project_statements(
821
+ statement.orelse,
822
+ [*dynamic_path, "loop:else"],
823
+ control_dependencies,
824
+ ):
825
+ return False
826
+ continue
827
+ barrier = {
828
+ "kind": "barrier",
829
+ "reason": "unsupported Python statement",
830
+ "span": _shadow_span(statement, line_offsets),
831
+ }
832
+ return False
833
+ return True
834
+
835
+ project_statements(module.body, [], [])
836
+ _emit(
837
+ {
838
+ "type": "shadow_plan",
839
+ "id": rid,
840
+ "eligible": True,
841
+ "revision": _STATE.namespace_revision,
842
+ "digest": _shadow_snapshot_digest(snapshot),
843
+ "values": snapshot,
844
+ "operations": operations,
845
+ "controls": controls,
846
+ "barrier": barrier,
847
+ }
848
+ )
849
+
850
+ def _drain_captured_stdout() -> None:
851
+ """Forward bytes written to the captured fd 1 as stdout frames.
852
+
853
+ Runs on a daemon thread for the life of the process. Child processes that
854
+ inherit fd 1 (any ``subprocess`` call without ``stdout=PIPE``) land here.
855
+ """
856
+ if _CAPTURE_READ_FD is None:
857
+ return
858
+ import codecs
859
+
860
+ decoder = codecs.getincrementaldecoder("utf-8")("replace")
861
+ while True:
862
+ try:
863
+ chunk = os.read(_CAPTURE_READ_FD, 65536)
864
+ except OSError:
865
+ return
866
+ if not chunk:
867
+ return
868
+ text = decoder.decode(chunk)
869
+ if not text:
870
+ continue
871
+ rid = _STATE.capture_rid
872
+ if rid is None:
873
+ _RAW_STDERR.write(text)
874
+ _RAW_STDERR.flush()
875
+ else:
876
+ _emit({"type": "stdout", "id": rid, "data": text})
877
+
878
+
879
+ def _start_capture_drain() -> None:
880
+ if _CAPTURE_READ_FD is None:
881
+ return
882
+ thread = threading.Thread(
883
+ target=_drain_captured_stdout, name="omp-fd1-capture", daemon=True
884
+ )
885
+ thread.start()
886
+
887
+
888
+ # ---------------------------------------------------------------------------
889
+ # Magic source transformer
890
+ # ---------------------------------------------------------------------------
891
+
892
+
893
+ _MAGIC_LINE_RE = re.compile(
894
+ r"^(?P<indent>[ \t]*)(?P<name>[A-Za-z_][A-Za-z_0-9]*)(?:[ \t]+(?P<args>.*))?$"
895
+ )
896
+ _ASSIGN_LINE_RE = re.compile(
897
+ r"^(?P<indent>[ \t]*)(?P<lhs>[A-Za-z_][A-Za-z_0-9.\[\], ]*?)\s*=\s*(?P<rhs>.+)$"
898
+ )
899
+
900
+
901
+ def _fold_continuations(lines: list[str], start: int) -> tuple[str, int]:
902
+ """Fold trailing backslash continuations starting at ``start``. Returns
903
+ ``(folded_text, lines_consumed)``."""
904
+ parts: list[str] = []
905
+ i = start
906
+ while i < len(lines):
907
+ line = lines[i]
908
+ if line.endswith("\\"):
909
+ parts.append(line[:-1])
910
+ i += 1
911
+ continue
912
+ parts.append(line)
913
+ i += 1
914
+ break
915
+ return ("".join(parts), i - start)
916
+
917
+
918
+ def _quote_arg(text: str) -> str:
919
+ """Return a Python string literal that round-trips ``text`` exactly."""
920
+ return json.dumps(text, ensure_ascii=False)
921
+
922
+
923
+ _STRING_START_TOKENS = {
924
+ getattr(tokenize, name, -1) for name in ("FSTRING_START", "TSTRING_START")
925
+ }
926
+ _STRING_END_TOKENS = {
927
+ getattr(tokenize, name, -1) for name in ("FSTRING_END", "TSTRING_END")
928
+ }
929
+
930
+
931
+ def _magic_line_indices(lines: list[str]) -> set[int]:
932
+ """Find command lines outside Python strings without lexing shell payloads."""
933
+ indices: set[int] = set()
934
+ offset = 0
935
+ while offset < len(lines):
936
+ # A synthetic enclosing expression suppresses indentation checks when
937
+ # restarting in the middle of a suite; this pass only identifies tokens.
938
+ stream = itertools.chain(
939
+ ("(\n",),
940
+ (line + "\n" for line in itertools.islice(lines, offset, None)),
941
+ (")\n",),
942
+ )
943
+ last_row = -1
944
+ string_depth = 0
945
+ try:
946
+ for item in tokenize.generate_tokens(lambda: next(stream, "")):
947
+ if item.type in _STRING_START_TOKENS:
948
+ string_depth += 1
949
+ elif item.type in _STRING_END_TOKENS:
950
+ string_depth -= 1
951
+ continue
952
+ if string_depth or item.type in (
953
+ tokenize.NL, tokenize.NEWLINE, tokenize.INDENT,
954
+ tokenize.DEDENT, tokenize.ENDMARKER, tokenize.COMMENT,
955
+ ) or (item.type == tokenize.ERRORTOKEN and item.string.isspace()):
956
+ # Python < 3.12 emits indentation before `!` as one
957
+ # whitespace ERRORTOKEN per column.
958
+ continue
959
+ row = offset + item.start[0] - 2
960
+ if row < offset or row >= len(lines) or row == last_row:
961
+ continue
962
+ last_row = row
963
+ stripped = lines[row].lstrip()
964
+ if item.start[1] != len(lines[row]) - len(stripped):
965
+ continue
966
+ assignment = _ASSIGN_LINE_RE.match(lines[row])
967
+ rhs = assignment.group("rhs").strip() if assignment else ""
968
+ if not stripped.startswith(("%", "!")) and not rhs.startswith(("%", "!")):
969
+ continue
970
+ indices.add(row)
971
+ if stripped.startswith("%%"):
972
+ return indices
973
+ # Shell quotes and cell/line magic arguments are not Python.
974
+ # Restart after the command instead of tokenizing its payload.
975
+ consumed = _fold_continuations(lines, row)[1] if stripped.startswith(("%", "!")) else 1
976
+ offset = row + consumed
977
+ break
978
+ else:
979
+ return indices
980
+ except (tokenize.TokenError, IndentationError):
981
+ # Leave incomplete Python source to the ordinary syntax-error path.
982
+ return indices
983
+ return indices
984
+
985
+
986
+ def transform_cell(source: str) -> str:
987
+ """Translate IPython-style magics + shell escapes into plain Python.
988
+
989
+ Rules
990
+ -----
991
+ * ``%name args`` -> ``__omp_magic("name", "args")``
992
+ * ``%load path`` -> ``await __omp_magic_async("load", "path")``
993
+ * ``var = %name args`` -> ``var = __omp_magic("name", "args")``
994
+ * ``!cmd`` -> ``__omp_shell("cmd")``
995
+ * ``var = !cmd`` -> ``var = __omp_shell("cmd")``
996
+ * ``%%name args\\n<body>`` -> ``__omp_magic_cell("name", "args", "<body>")``
997
+ (cell magic must be the first non-whitespace token of a top-level line and
998
+ consumes the remainder of the cell)
999
+
1000
+ Python tokenization identifies command lines outside strings and comments.
1001
+ Real command payloads are skipped before resuming lexical scanning so shell
1002
+ quoting cannot change how subsequent Python source is interpreted.
1003
+ """
1004
+
1005
+ if "%" not in source and "!" not in source:
1006
+ return source
1007
+
1008
+ lines = source.splitlines()
1009
+ magic_lines = _magic_line_indices(lines)
1010
+ out: list[str] = []
1011
+ i = 0
1012
+ while i < len(lines):
1013
+ line = lines[i]
1014
+ if i not in magic_lines:
1015
+ out.append(line)
1016
+ i += 1
1017
+ continue
1018
+ stripped = line.lstrip()
1019
+ indent = line[: len(line) - len(stripped)]
1020
+
1021
+ # Cell magic — consumes from here to EOF.
1022
+ if stripped.startswith("%%"):
1023
+ head, _ = _split_magic_head(stripped[2:])
1024
+ name, args = head
1025
+ body_lines = lines[i + 1 :]
1026
+ body = "\n".join(body_lines)
1027
+ out.append(
1028
+ f"{indent}__omp_magic_cell({_quote_arg(name)}, {_quote_arg(args)}, {_quote_arg(body)})"
1029
+ )
1030
+ return "\n".join(out)
1031
+
1032
+ # Line magic / shell at start of line.
1033
+ if stripped.startswith("%") and not stripped.startswith("%%"):
1034
+ folded, consumed = _fold_continuations(lines, i)
1035
+ stripped_folded = folded.lstrip()
1036
+ indent = folded[: len(folded) - len(stripped_folded)]
1037
+ head, _ = _split_magic_head(stripped_folded[1:])
1038
+ name, args = head
1039
+ call = "__omp_magic_async" if name == "load" else "__omp_magic"
1040
+ prefix = "await " if name == "load" else ""
1041
+ out.append(f"{indent}{prefix}{call}({_quote_arg(name)}, {_quote_arg(args)})")
1042
+ i += consumed
1043
+ continue
1044
+
1045
+ if stripped.startswith("!"):
1046
+ folded, consumed = _fold_continuations(lines, i)
1047
+ stripped_folded = folded.lstrip()
1048
+ indent = folded[: len(folded) - len(stripped_folded)]
1049
+ cmd = stripped_folded[1:].strip()
1050
+ out.append(f"{indent}__omp_shell({_quote_arg(cmd)})")
1051
+ i += consumed
1052
+ continue
1053
+
1054
+ # Assignment forms: var = %magic / var = !cmd
1055
+ m = _ASSIGN_LINE_RE.match(line)
1056
+ if m:
1057
+ rhs = m.group("rhs").strip()
1058
+ if rhs.startswith("!"):
1059
+ cmd = rhs[1:].strip()
1060
+ out.append(
1061
+ f"{m.group('indent')}{m.group('lhs').rstrip()} = __omp_shell({_quote_arg(cmd)})"
1062
+ )
1063
+ i += 1
1064
+ continue
1065
+ if rhs.startswith("%") and not rhs.startswith("%%"):
1066
+ head, _ = _split_magic_head(rhs[1:])
1067
+ name, args = head
1068
+ call = "__omp_magic_async" if name == "load" else "__omp_magic"
1069
+ prefix = "await " if name == "load" else ""
1070
+ out.append(
1071
+ f"{m.group('indent')}{m.group('lhs').rstrip()} = {prefix}{call}({_quote_arg(name)}, {_quote_arg(args)})"
1072
+ )
1073
+ i += 1
1074
+ continue
1075
+
1076
+ out.append(line)
1077
+ i += 1
1078
+
1079
+ return "\n".join(out)
1080
+
1081
+
1082
+ def _split_magic_head(text: str) -> tuple[tuple[str, str], str]:
1083
+ """Split ``"name rest"`` into ``("name", "rest")``."""
1084
+ text = text.lstrip()
1085
+ if not text:
1086
+ return ("", ""), ""
1087
+ m = re.match(r"([A-Za-z_][A-Za-z_0-9]*)(?:\s+(.*))?$", text)
1088
+ if not m:
1089
+ return ("", text), ""
1090
+ return (m.group(1), (m.group(2) or "").rstrip()), ""
1091
+
1092
+
1093
+ # ---------------------------------------------------------------------------
1094
+ # Magic registry
1095
+ # ---------------------------------------------------------------------------
1096
+
1097
+
1098
+ _LINE_MAGICS: dict[str, Callable[[str], Any]] = {}
1099
+ _CELL_MAGICS: dict[str, Callable[[str, str], Any]] = {}
1100
+
1101
+
1102
+ def line_magic(name: str) -> Callable[[Callable[[str], Any]], Callable[[str], Any]]:
1103
+ def decorator(fn: Callable[[str], Any]) -> Callable[[str], Any]:
1104
+ _LINE_MAGICS[name] = fn
1105
+ return fn
1106
+
1107
+ return decorator
1108
+
1109
+
1110
+ def cell_magic(
1111
+ name: str,
1112
+ ) -> Callable[[Callable[[str, str], Any]], Callable[[str, str], Any]]:
1113
+ def decorator(fn: Callable[[str, str], Any]) -> Callable[[str, str], Any]:
1114
+ _CELL_MAGICS[name] = fn
1115
+ return fn
1116
+
1117
+ return decorator
1118
+
1119
+
1120
+ def _emit_status(op: str, **data: Any) -> None:
1121
+ bundle = {"application/x-omp-status": {"op": op, **data}}
1122
+ rid = _CURRENT_RID.get()
1123
+ if rid is None:
1124
+ return
1125
+ _emit({"type": "display", "id": rid, "bundle": bundle})
1126
+
1127
+
1128
+ _SHELL_READ_CHUNK_BYTES = 8192
1129
+ _SHELL_OUTPUT_MAX_BYTES = 1024 * 1024
1130
+ _SHELL_OUTPUT_MAX_LINES = 3000
1131
+ _SHELL_RESULT_CAPTURE_BYTES = _SHELL_OUTPUT_MAX_BYTES
1132
+ _PIP_LINE_SCAN_CHARS = 64 * 1024
1133
+ _SHELL_TRUNCATION_NOTICE = (
1134
+ f"[output truncated: shell helper exceeded {_SHELL_OUTPUT_MAX_BYTES} bytes "
1135
+ f"or {_SHELL_OUTPUT_MAX_LINES} lines; remaining output discarded]\n"
1136
+ )
1137
+
1138
+
1139
+ def _process_output_encoding() -> str:
1140
+ return locale.getpreferredencoding(False) or "utf-8"
1141
+
1142
+
1143
+ def _process_output_decoder(encoding: str) -> codecs.IncrementalDecoder:
1144
+ return codecs.getincrementaldecoder(encoding)(errors="strict")
1145
+
1146
+
1147
+ def _take_prefix_by_lines(text: str, max_lines: int) -> str:
1148
+ if max_lines <= 0:
1149
+ return ""
1150
+ cursor = 0
1151
+ for _ in range(max_lines):
1152
+ newline = text.find("\n", cursor)
1153
+ if newline < 0:
1154
+ return text
1155
+ cursor = newline + 1
1156
+ return text[:cursor]
1157
+
1158
+
1159
+ def _take_prefix_by_encoded_bytes(text: str, max_bytes: int, encoding: str) -> str:
1160
+ if max_bytes <= 0:
1161
+ return ""
1162
+ if len(text.encode(encoding, errors="strict")) <= max_bytes:
1163
+ return text
1164
+ lo = 0
1165
+ hi = len(text)
1166
+ while lo < hi:
1167
+ mid = (lo + hi + 1) // 2
1168
+ if len(text[:mid].encode(encoding, errors="strict")) <= max_bytes:
1169
+ lo = mid
1170
+ else:
1171
+ hi = mid - 1
1172
+ return text[:lo]
1173
+
1174
+
1175
+ class _ShellOutputLimiter:
1176
+ def __init__(self, *, max_bytes: int, max_lines: int, encoding: str) -> None:
1177
+ self._remaining_bytes = max_bytes
1178
+ self._remaining_lines = max_lines
1179
+ self._encoding = encoding
1180
+ self._truncated = False
1181
+ self._at_line_start = True
1182
+
1183
+ def write(self, text: str) -> None:
1184
+ if not text or self._truncated:
1185
+ return
1186
+ limited = _take_prefix_by_lines(text, self._remaining_lines)
1187
+ truncated = limited != text
1188
+ byte_limited = _take_prefix_by_encoded_bytes(
1189
+ limited, self._remaining_bytes, self._encoding
1190
+ )
1191
+ truncated = truncated or byte_limited != limited
1192
+ if byte_limited:
1193
+ sys.stdout.write(byte_limited)
1194
+ sys.stdout.flush()
1195
+ self._remaining_bytes -= len(
1196
+ byte_limited.encode(self._encoding, errors="strict")
1197
+ )
1198
+ self._remaining_lines -= byte_limited.count("\n")
1199
+ self._at_line_start = byte_limited.endswith("\n")
1200
+ if truncated:
1201
+ self._emit_truncation_notice()
1202
+
1203
+ def _emit_truncation_notice(self) -> None:
1204
+ if self._truncated:
1205
+ return
1206
+ prefix = "" if self._at_line_start else "\n"
1207
+ sys.stdout.write(prefix + _SHELL_TRUNCATION_NOTICE)
1208
+ sys.stdout.flush()
1209
+ self._truncated = True
1210
+
1211
+
1212
+ def _stream_process_output(
1213
+ proc: subprocess.Popen, on_text: Callable[[str], None] | None = None
1214
+ ) -> None:
1215
+ assert proc.stdout is not None
1216
+ encoding = _process_output_encoding()
1217
+ decoder = _process_output_decoder(encoding)
1218
+ limiter = _ShellOutputLimiter(
1219
+ max_bytes=_SHELL_OUTPUT_MAX_BYTES,
1220
+ max_lines=_SHELL_OUTPUT_MAX_LINES,
1221
+ encoding=encoding,
1222
+ )
1223
+ while True:
1224
+ chunk = os.read(proc.stdout.fileno(), _SHELL_READ_CHUNK_BYTES)
1225
+ if not chunk:
1226
+ break
1227
+ text = decoder.decode(chunk)
1228
+ if text:
1229
+ limiter.write(text)
1230
+ if on_text is not None:
1231
+ on_text(text)
1232
+ tail = decoder.decode(b"", final=True)
1233
+ if tail:
1234
+ limiter.write(tail)
1235
+ if on_text is not None:
1236
+ on_text(tail)
1237
+
1238
+
1239
+ class _BoundedTextCapture:
1240
+ def __init__(self, max_bytes: int, max_lines: int, encoding: str) -> None:
1241
+ self._remaining_bytes = max_bytes
1242
+ self._remaining_lines = max_lines
1243
+ self._encoding = encoding
1244
+ self._parts: list[str] = []
1245
+
1246
+ def add(self, text: str) -> None:
1247
+ if self._remaining_bytes <= 0 or self._remaining_lines <= 0:
1248
+ return
1249
+ line_limited = _take_prefix_by_lines(text, self._remaining_lines)
1250
+ part = _take_prefix_by_encoded_bytes(
1251
+ line_limited, self._remaining_bytes, self._encoding
1252
+ )
1253
+ if not part:
1254
+ return
1255
+ self._parts.append(part)
1256
+ self._remaining_bytes -= len(part.encode(self._encoding, errors="strict"))
1257
+ self._remaining_lines -= part.count("\n")
1258
+
1259
+ def text(self) -> str:
1260
+ return "".join(self._parts)
1261
+
1262
+
1263
+ class _BoundedLineScanner:
1264
+ def __init__(self, max_chars: int, on_line: Callable[[str], None]) -> None:
1265
+ self._max_chars = max_chars
1266
+ self._on_line = on_line
1267
+ self._partial = ""
1268
+
1269
+ def add(self, text: str) -> None:
1270
+ data = self._partial + text
1271
+ lines = data.splitlines(keepends=True)
1272
+ if not lines:
1273
+ return
1274
+ if lines[-1].endswith(("\n", "\r")):
1275
+ self._partial = ""
1276
+ else:
1277
+ self._partial = lines.pop()
1278
+ for line in lines:
1279
+ self._on_line(line)
1280
+ if len(self._partial) > self._max_chars:
1281
+ self._partial = self._partial[-self._max_chars :]
1282
+
1283
+ def finish(self) -> None:
1284
+ if self._partial:
1285
+ self._on_line(self._partial)
1286
+ self._partial = ""
1287
+
1288
+
1289
+ @line_magic("pip")
1290
+ def _magic_pip(args: str) -> None:
1291
+ argv = shlex.split(args) if args else ["--help"]
1292
+ cmd = [sys.executable, "-m", "pip", *argv]
1293
+ # Installs are host work, not cell compute: suspend the eval watchdog (same
1294
+ # protocol as bridge waits in bridge-timeout.ts). Caller aborts still apply.
1295
+ _emit_status("timeout-pause")
1296
+ try:
1297
+ _run_pip(cmd, args)
1298
+ finally:
1299
+ _emit_status("timeout-resume")
1300
+
1301
+
1302
+ def _run_pip(cmd: list[str], args: str) -> None:
1303
+ # stdin=DEVNULL: see _run_shell_body.
1304
+ proc = subprocess.Popen(
1305
+ cmd,
1306
+ stdin=subprocess.DEVNULL,
1307
+ stdout=subprocess.PIPE,
1308
+ stderr=subprocess.STDOUT,
1309
+ )
1310
+ installed_packages: list[str] = []
1311
+
1312
+ def scan_pip_line(raw_line: str) -> None:
1313
+ m = re.search(r"Successfully installed\s+(.+)$", raw_line)
1314
+ if m:
1315
+ for token in m.group(1).split():
1316
+ # Token is name-version; drop the version suffix.
1317
+ pkg = token.rsplit("-", 1)[0]
1318
+ installed_packages.append(pkg.replace("_", "-"))
1319
+
1320
+ scanner = _BoundedLineScanner(_PIP_LINE_SCAN_CHARS, scan_pip_line)
1321
+ _stream_process_output(proc, scanner.add)
1322
+ scanner.finish()
1323
+ proc.wait()
1324
+ if installed_packages:
1325
+ import importlib
1326
+
1327
+ importlib.invalidate_caches()
1328
+ prefixes = {pkg.lower().replace("-", "_") for pkg in installed_packages}
1329
+ for mod_name in list(sys.modules):
1330
+ head = mod_name.split(".", 1)[0].lower()
1331
+ if head in prefixes:
1332
+ sys.modules.pop(mod_name, None)
1333
+ _emit_status(
1334
+ "pip", args=args, installed=installed_packages, exit_code=proc.returncode
1335
+ )
1336
+
1337
+
1338
+ @line_magic("cd")
1339
+ def _magic_cd(args: str) -> str:
1340
+ path = os.path.expanduser(args.strip()) or os.path.expanduser("~")
1341
+ os.chdir(path)
1342
+ cwd = os.getcwd()
1343
+ _emit_status("cd", path=cwd)
1344
+ return cwd
1345
+
1346
+
1347
+ @line_magic("pwd")
1348
+ def _magic_pwd(_args: str) -> str:
1349
+ cwd = os.getcwd()
1350
+ _emit_status("pwd", path=cwd)
1351
+ return cwd
1352
+
1353
+
1354
+ @line_magic("ls")
1355
+ def _magic_ls(args: str) -> list[str]:
1356
+ target = os.path.expanduser(args.strip()) or "."
1357
+ entries = sorted(os.listdir(target))
1358
+ _emit_status("ls", path=os.path.abspath(target), count=len(entries))
1359
+ return entries
1360
+
1361
+
1362
+ @line_magic("env")
1363
+ def _magic_env(args: str) -> Any:
1364
+ args = args.strip()
1365
+ if not args:
1366
+ return dict(sorted(os.environ.items()))
1367
+ if "=" in args:
1368
+ key, value = args.split("=", 1)
1369
+ os.environ[key.strip()] = value.strip()
1370
+ return value.strip()
1371
+ return os.environ.get(args)
1372
+
1373
+
1374
+ @line_magic("set_env")
1375
+ def _magic_set_env(args: str) -> str:
1376
+ parts = args.split(None, 1)
1377
+ if len(parts) != 2:
1378
+ raise ValueError("Usage: %set_env KEY VALUE")
1379
+ key, value = parts
1380
+ os.environ[key] = value
1381
+ return value
1382
+
1383
+
1384
+ @line_magic("time")
1385
+ def _magic_time(args: str) -> Any:
1386
+ start = time.perf_counter()
1387
+ result = eval(args, _STATE.user_ns)
1388
+ elapsed = time.perf_counter() - start
1389
+ sys.stdout.write(f"Wall time: {elapsed * 1000:.2f} ms\n")
1390
+ _emit_status("time", elapsed_ms=round(elapsed * 1000, 3))
1391
+ return result
1392
+
1393
+
1394
+ @line_magic("timeit")
1395
+ def _magic_timeit(args: str) -> None:
1396
+ import timeit as _timeit
1397
+
1398
+ timer = _timeit.Timer(stmt=args, globals=_STATE.user_ns)
1399
+ iters, total = timer.autorange()
1400
+ per = total / iters
1401
+ sys.stdout.write(f"{iters} loops, best of 1: {per * 1e6:.2f} us per loop\n")
1402
+ _emit_status("timeit", loops=iters, total_ms=round(total * 1000, 3))
1403
+
1404
+
1405
+ @line_magic("who")
1406
+ def _magic_who(_args: str) -> list[str]:
1407
+ names = sorted(
1408
+ name
1409
+ for name, value in _STATE.user_ns.items()
1410
+ if not name.startswith("_")
1411
+ and not callable(value)
1412
+ or hasattr(value, "__class__")
1413
+ )
1414
+ return [n for n in names if not n.startswith("__")]
1415
+
1416
+
1417
+ @line_magic("whos")
1418
+ def _magic_whos(_args: str) -> list[tuple[str, str]]:
1419
+ rows = []
1420
+ for name in sorted(_STATE.user_ns):
1421
+ if name.startswith("__"):
1422
+ continue
1423
+ value = _STATE.user_ns[name]
1424
+ rows.append((name, type(value).__name__))
1425
+ return rows
1426
+
1427
+
1428
+ @line_magic("reset")
1429
+ def _magic_reset(_args: str) -> None:
1430
+ _STATE.user_ns.clear()
1431
+ _STATE.user_ns.update(
1432
+ {"__name__": "__main__", "__doc__": None, "__builtins__": builtins}
1433
+ )
1434
+ _install_builtins(_STATE.user_ns)
1435
+ _emit_status("reset")
1436
+
1437
+
1438
+ @line_magic("load")
1439
+ async def _magic_load(args: str) -> None:
1440
+ try:
1441
+ parts = shlex.split(args)
1442
+ except ValueError as exc:
1443
+ raise ValueError(f"Usage: %load <path> ({exc})") from exc
1444
+ if len(parts) != 1:
1445
+ raise ValueError("Usage: %load <path> (quote paths containing spaces)")
1446
+ path = Path(os.path.expanduser(parts[0])).resolve()
1447
+ source = path.read_text(encoding="utf-8")
1448
+ await _exec_source_async(
1449
+ transform_cell(source), _STATE.user_ns,
1450
+ filename=str(path), linecache_source=source,
1451
+ )
1452
+
1453
+
1454
+ @line_magic("run")
1455
+ def _magic_run(args: str) -> None:
1456
+ parts = shlex.split(args) if args else []
1457
+ if not parts:
1458
+ raise ValueError("Usage: %run <path>")
1459
+ target = os.path.expanduser(parts[0])
1460
+ saved_argv = sys.argv
1461
+ try:
1462
+ sys.argv = [target, *parts[1:]]
1463
+ result_ns = runpy.run_path(target, run_name="__main__")
1464
+ finally:
1465
+ sys.argv = saved_argv
1466
+ for name, value in result_ns.items():
1467
+ if name.startswith("__"):
1468
+ continue
1469
+ _STATE.user_ns[name] = value
1470
+
1471
+
1472
+ def _resolve_bash() -> str:
1473
+ if os.name != "nt":
1474
+ return "/bin/bash"
1475
+ # Prefer Git Bash over WSL's System32 bash.exe, which runs inside a
1476
+ # separate Linux environment and does not share the Windows filesystem
1477
+ # layout or PATH.
1478
+ for env_var, suffix in (
1479
+ ("ProgramFiles", r"Git\bin\bash.exe"),
1480
+ ("ProgramFiles(x86)", r"Git\bin\bash.exe"),
1481
+ ("LOCALAPPDATA", r"Programs\Git\bin\bash.exe"),
1482
+ ):
1483
+ root = os.environ.get(env_var)
1484
+ if root:
1485
+ candidate = os.path.join(root, suffix)
1486
+ if os.path.isfile(candidate):
1487
+ return candidate
1488
+ found = shutil.which("bash")
1489
+ if found and "system32" not in found.lower():
1490
+ return found
1491
+ # WSL's System32 bash.exe runs in a separate Linux environment, so
1492
+ # silently falling back to it would execute the cell somewhere the user
1493
+ # did not intend; fail loudly instead.
1494
+ raise RuntimeError(
1495
+ "%%bash requires a POSIX bash, but none was found. "
1496
+ "Install Git for Windows or add a non-WSL bash to PATH."
1497
+ )
1498
+
1499
+
1500
+ @cell_magic("bash")
1501
+ def _magic_cell_bash(args: str, body: str) -> int:
1502
+ return _run_shell_body(body, shell_arg=_resolve_bash())
1503
+
1504
+
1505
+ @cell_magic("capture")
1506
+ def _magic_cell_capture(args: str, body: str) -> str:
1507
+ """Capture stdout/stderr of body; bind to ``args`` (a name) if provided."""
1508
+ captured = io.StringIO()
1509
+ saved_stdout, saved_stderr = sys.stdout, sys.stderr
1510
+ sys.stdout = sys.stderr = captured
1511
+ try:
1512
+ _exec_source(body, _STATE.user_ns)
1513
+ finally:
1514
+ sys.stdout, sys.stderr = saved_stdout, saved_stderr
1515
+ text = captured.getvalue()
1516
+ name = args.strip()
1517
+ if name:
1518
+ _STATE.user_ns[name] = text
1519
+ return text
1520
+
1521
+
1522
+ @cell_magic("timeit")
1523
+ def _magic_cell_timeit(args: str, body: str) -> None:
1524
+ import timeit as _timeit
1525
+
1526
+ timer = _timeit.Timer(stmt=body, globals=_STATE.user_ns)
1527
+ iters, total = timer.autorange()
1528
+ per = total / iters
1529
+ sys.stdout.write(f"{iters} loops, best of 1: {per * 1e6:.2f} us per loop\n")
1530
+ _emit_status("timeit", loops=iters, total_ms=round(total * 1000, 3))
1531
+
1532
+
1533
+ @cell_magic("writefile")
1534
+ def _magic_cell_writefile(args: str, body: str) -> str:
1535
+ path = Path(os.path.expanduser(args.strip()))
1536
+ path.parent.mkdir(parents=True, exist_ok=True)
1537
+ path.write_text(body, encoding="utf-8")
1538
+ _emit_status("writefile", path=str(path), bytes=len(body))
1539
+ return str(path)
1540
+
1541
+
1542
+ def _run_shell_body(body: str, *, shell_arg: str) -> int:
1543
+ # stdin=DEVNULL: children must not inherit the runner's stdin, which is
1544
+ # the host's NDJSON control channel (a reading child would steal frames,
1545
+ # and inheriting the pipe deadlocks nested interpreters on Windows).
1546
+ proc = subprocess.Popen(
1547
+ [shell_arg, "-c", body],
1548
+ stdin=subprocess.DEVNULL,
1549
+ stdout=subprocess.PIPE,
1550
+ stderr=subprocess.STDOUT,
1551
+ )
1552
+ _stream_process_output(proc)
1553
+ proc.wait()
1554
+ return proc.returncode
1555
+
1556
+
1557
+ def __omp_magic(name: str, args: str) -> Any:
1558
+ fn = _LINE_MAGICS.get(name)
1559
+ if fn is None:
1560
+ raise NameError(f"UsageError: Line magic function '%{name}' not found.")
1561
+ return fn(args)
1562
+
1563
+
1564
+ async def __omp_magic_async(name: str, args: str) -> Any:
1565
+ result = __omp_magic(name, args)
1566
+ return await result if inspect.isawaitable(result) else result
1567
+
1568
+
1569
+ def __omp_magic_cell(name: str, args: str, body: str) -> Any:
1570
+ fn = _CELL_MAGICS.get(name)
1571
+ if fn is None:
1572
+ raise NameError(f"UsageError: Cell magic function '%%{name}' not found.")
1573
+ return fn(args, body)
1574
+
1575
+
1576
+ class _ShellResult(list):
1577
+ """Result of ``!cmd`` — list of stripped output lines."""
1578
+
1579
+ def __init__(self, lines: list[str], returncode: int) -> None:
1580
+ super().__init__(lines)
1581
+ self.returncode = returncode
1582
+
1583
+ @property
1584
+ def n(self) -> str: # IPython compat
1585
+ return "\n".join(self)
1586
+
1587
+ @property
1588
+ def s(self) -> str: # IPython compat
1589
+ return " ".join(self)
1590
+
1591
+
1592
+ def __omp_shell(cmd: str) -> _ShellResult:
1593
+ # stdin=DEVNULL: see _run_shell_body.
1594
+ proc = subprocess.Popen(
1595
+ cmd,
1596
+ shell=True,
1597
+ stdin=subprocess.DEVNULL,
1598
+ stdout=subprocess.PIPE,
1599
+ stderr=subprocess.STDOUT,
1600
+ )
1601
+ capture = _BoundedTextCapture(
1602
+ _SHELL_RESULT_CAPTURE_BYTES, _SHELL_OUTPUT_MAX_LINES, _process_output_encoding()
1603
+ )
1604
+ _stream_process_output(proc, capture.add)
1605
+ proc.wait()
1606
+ lines = [line for line in capture.text().splitlines()]
1607
+ return _ShellResult(lines, proc.returncode)
1608
+
1609
+
1610
+ # ---------------------------------------------------------------------------
1611
+ # Display dispatch
1612
+ # ---------------------------------------------------------------------------
1613
+
1614
+
1615
+ _REPR_MIMES = [
1616
+ ("_repr_html_", "text/html"),
1617
+ ("_repr_markdown_", "text/markdown"),
1618
+ ("_repr_svg_", "image/svg+xml"),
1619
+ ("_repr_png_", "image/png"),
1620
+ ("_repr_jpeg_", "image/jpeg"),
1621
+ ("_repr_json_", "application/json"),
1622
+ ("_repr_latex_", "text/latex"),
1623
+ ]
1624
+
1625
+
1626
+ def _is_matplotlib_figure(value: Any) -> bool:
1627
+ figure_module = sys.modules.get("matplotlib.figure")
1628
+ figure_cls = getattr(figure_module, "Figure", None)
1629
+ if isinstance(figure_cls, type) and isinstance(value, figure_cls):
1630
+ return True
1631
+
1632
+ value_type = type(value)
1633
+ return (
1634
+ value_type.__module__ == "matplotlib.figure" and value_type.__name__ == "Figure"
1635
+ )
1636
+
1637
+
1638
+ def _matplotlib_figure_png(value: Any) -> str | None:
1639
+ if not _is_matplotlib_figure(value):
1640
+ return None
1641
+
1642
+ savefig = getattr(value, "savefig", None)
1643
+ if not callable(savefig):
1644
+ return None
1645
+
1646
+ try:
1647
+ buf = io.BytesIO()
1648
+ savefig(buf, format="png", bbox_inches="tight")
1649
+ except Exception:
1650
+ return None
1651
+
1652
+ displayed_ids = _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.get()
1653
+ if displayed_ids is not None:
1654
+ displayed_ids.add(id(value))
1655
+ return base64.b64encode(buf.getvalue()).decode("ascii")
1656
+
1657
+
1658
+ def _coerce_image_bytes(value: Any) -> str:
1659
+ if isinstance(value, (bytes, bytearray)):
1660
+ return base64.b64encode(bytes(value)).decode("ascii")
1661
+ if isinstance(value, str):
1662
+ return value
1663
+ return base64.b64encode(repr(value).encode("utf-8")).decode("ascii")
1664
+
1665
+
1666
+ def _mime_bundle(value: Any) -> dict:
1667
+ """Build a Jupyter-style MIME bundle for ``value``.
1668
+
1669
+ Honors ``_repr_mimebundle_`` first, falls back to individual ``_repr_*_``
1670
+ accessors, and always provides ``text/plain``.
1671
+ """
1672
+ # Match JS display({type: "image", ...}) without losing detail or remote
1673
+ # references in a plain PNG/JPEG MIME representation.
1674
+ if (
1675
+ isinstance(value, dict)
1676
+ and value.get("type") == "image"
1677
+ and isinstance(value.get("mimeType"), str)
1678
+ and isinstance(value.get("data"), (str, bytes, bytearray))
1679
+ ):
1680
+ image = dict(value)
1681
+ image["data"] = _coerce_image_bytes(image["data"])
1682
+ return {"application/x-omp-image": image}
1683
+
1684
+ bundle: dict[str, Any] = {}
1685
+ matplotlib_png = _matplotlib_figure_png(value)
1686
+ if matplotlib_png is not None:
1687
+ bundle["image/png"] = matplotlib_png
1688
+
1689
+ mimebundle = getattr(value, "_repr_mimebundle_", None)
1690
+ if callable(mimebundle):
1691
+ try:
1692
+ data = mimebundle()
1693
+ except Exception:
1694
+ data = None
1695
+ if isinstance(data, tuple):
1696
+ data = data[0]
1697
+ if isinstance(data, dict):
1698
+ bundle.update({str(k): v for k, v in data.items()})
1699
+
1700
+ for attr, mime in _REPR_MIMES:
1701
+ if mime in bundle:
1702
+ continue
1703
+ repr_fn = getattr(value, attr, None)
1704
+ if not callable(repr_fn):
1705
+ continue
1706
+ try:
1707
+ data = repr_fn()
1708
+ except Exception:
1709
+ continue
1710
+ if data is None:
1711
+ continue
1712
+ if mime in ("image/png", "image/jpeg"):
1713
+ bundle[mime] = _coerce_image_bytes(data)
1714
+ else:
1715
+ bundle[mime] = data
1716
+
1717
+ if "text/plain" not in bundle:
1718
+ try:
1719
+ bundle["text/plain"] = repr(value)
1720
+ except Exception:
1721
+ bundle["text/plain"] = f"<unrepr {type(value).__name__}>"
1722
+
1723
+ return bundle
1724
+
1725
+
1726
+ def _emit_display(bundle: dict, *, kind: str = "display") -> None:
1727
+ rid = _CURRENT_RID.get()
1728
+ if rid is None:
1729
+ return
1730
+ _emit({"type": kind, "id": rid, "bundle": bundle})
1731
+
1732
+
1733
+ def __omp_display(value: Any, *, raw: bool = False, kind: str = "display") -> None:
1734
+ if raw:
1735
+ if not isinstance(value, dict):
1736
+ raise TypeError("display(..., raw=True) requires a MIME bundle dict")
1737
+ bundle = {str(k): v for k, v in value.items()}
1738
+ if "text/plain" not in bundle:
1739
+ bundle["text/plain"] = ""
1740
+ _emit_display(bundle, kind=kind)
1741
+ return
1742
+ _emit_display(_mime_bundle(value), kind=kind)
1743
+
1744
+
1745
+ # ---------------------------------------------------------------------------
1746
+ # Matplotlib post-cell flush
1747
+ # ---------------------------------------------------------------------------
1748
+
1749
+
1750
+ def _flush_matplotlib_figures() -> None:
1751
+ plt = sys.modules.get("matplotlib.pyplot")
1752
+ if plt is None:
1753
+ return
1754
+ try:
1755
+ fignums = list(plt.get_fignums())
1756
+ except Exception:
1757
+ return
1758
+ for num in fignums:
1759
+ try:
1760
+ fig = plt.figure(num)
1761
+ if id(fig) in (_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.get() or set()):
1762
+ plt.close(fig)
1763
+ continue
1764
+ buf = io.BytesIO()
1765
+ fig.savefig(buf, format="png", bbox_inches="tight")
1766
+ data = base64.b64encode(buf.getvalue()).decode("ascii")
1767
+ _emit_display({"image/png": data, "text/plain": f"<Figure {num}>"})
1768
+ plt.close(fig)
1769
+ except Exception:
1770
+ continue
1771
+
1772
+
1773
+ # Force a non-interactive backend before user code imports matplotlib. Set as
1774
+ # environ default so the user can still override it explicitly.
1775
+ os.environ.setdefault("MPLBACKEND", "Agg")
1776
+
1777
+
1778
+ # ---------------------------------------------------------------------------
1779
+ # Builtin injection
1780
+ # ---------------------------------------------------------------------------
1781
+
1782
+
1783
+ def _install_builtins(ns: dict) -> None:
1784
+ ns["display"] = __omp_display
1785
+ ns["__omp_display"] = __omp_display
1786
+ ns["__omp_magic"] = __omp_magic
1787
+ ns["__omp_magic_async"] = __omp_magic_async
1788
+ ns["__omp_magic_cell"] = __omp_magic_cell
1789
+ ns["__omp_shell"] = __omp_shell
1790
+ ns["__omp_current_run_id__"] = lambda: _CURRENT_RID.get()
1791
+
1792
+
1793
+ _install_builtins(_STATE.user_ns)
1794
+
1795
+
1796
+ # ---------------------------------------------------------------------------
1797
+ # Source execution (split last expression for rich display)
1798
+ # ---------------------------------------------------------------------------
1799
+
1800
+
1801
+ _TLA_FLAG = getattr(ast, "PyCF_ALLOW_TOP_LEVEL_AWAIT", 0x2000)
1802
+
1803
+
1804
+ def _await_sync(coro) -> Any:
1805
+ try:
1806
+ running_loop = asyncio.get_running_loop()
1807
+ except RuntimeError:
1808
+ running_loop = None
1809
+ if running_loop is not None and running_loop.is_running():
1810
+ raise RuntimeError(
1811
+ "top-level await is not supported from synchronous magic execution"
1812
+ )
1813
+ return asyncio.run(coro)
1814
+
1815
+
1816
+ def _run_compiled_sync(code, ns: dict, *, want_value: bool) -> Any:
1817
+ """Synchronous execution path used by nested magic helpers."""
1818
+ if code.co_flags & inspect.CO_COROUTINE:
1819
+ result = _await_sync(eval(code, ns))
1820
+ return result if want_value else None
1821
+ if want_value:
1822
+ return eval(code, ns)
1823
+ exec(code, ns)
1824
+ return None
1825
+
1826
+
1827
+ async def _run_compiled_async(code, ns: dict, *, want_value: bool) -> Any:
1828
+ """Execute a code object in the persistent event loop.
1829
+
1830
+ Coroutine code is awaited in this task so top-level ``await`` interleaves
1831
+ with sibling requests. Plain statement/expression code runs on the main
1832
+ runner thread so SIGINT can interrupt it reliably.
1833
+ """
1834
+ if code.co_flags & inspect.CO_COROUTINE:
1835
+ result = await eval(code, ns)
1836
+ return result if want_value else None
1837
+ if want_value:
1838
+ return eval(code, ns)
1839
+ exec(code, ns)
1840
+ return None
1841
+
1842
+
1843
+
1844
+ _CALL_SITE_HELPER_NAME = "__omp_with_call_site__"
1845
+
1846
+ # Runner-owned reference to the genuine instrumentation helper. The helper is
1847
+ # defined by the eval prelude inside the user namespace, so the first callable
1848
+ # observed there is trusted: kernel init always executes the prelude before
1849
+ # user cells. Any later replacement is pollution from a retained cell.
1850
+ _TRUSTED_CALL_SITE_HELPERS: dict[str, Any] = {}
1851
+
1852
+
1853
+ def _restore_call_site_helper(ns: dict) -> None:
1854
+ """Reinstall the genuine helper after retained-cell pollution.
1855
+
1856
+ A previous cell can assign or delete `__omp_with_call_site__`; without
1857
+ this, the next rewritten `tool.read(...)` would resolve to
1858
+ user-controlled state and change ordinary eval behavior even with
1859
+ speculation disabled.
1860
+ """
1861
+ current = ns.get(_CALL_SITE_HELPER_NAME)
1862
+ trusted = _TRUSTED_CALL_SITE_HELPERS.get(_CALL_SITE_HELPER_NAME)
1863
+ if trusted is None:
1864
+ if callable(current):
1865
+ _TRUSTED_CALL_SITE_HELPERS[_CALL_SITE_HELPER_NAME] = current
1866
+ return
1867
+ if current is not trusted:
1868
+ ns[_CALL_SITE_HELPER_NAME] = trusted
1869
+
1870
+
1871
+ def _cell_binds_call_site_helper(module: ast.Module) -> bool:
1872
+ """Whether the cell binds the reserved helper name in any scope."""
1873
+ for node in ast.walk(module):
1874
+ if isinstance(node, ast.Name):
1875
+ if node.id == _CALL_SITE_HELPER_NAME and isinstance(node.ctx, (ast.Store, ast.Del)):
1876
+ return True
1877
+ elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
1878
+ if node.name == _CALL_SITE_HELPER_NAME:
1879
+ return True
1880
+ elif isinstance(node, ast.arg):
1881
+ if node.arg == _CALL_SITE_HELPER_NAME:
1882
+ return True
1883
+ elif isinstance(node, ast.alias):
1884
+ if node.asname == _CALL_SITE_HELPER_NAME:
1885
+ return True
1886
+ elif isinstance(node, ast.ExceptHandler):
1887
+ if node.name == _CALL_SITE_HELPER_NAME:
1888
+ return True
1889
+ return False
1890
+
1891
+ class _ShadowCallSiteTransformer(ast.NodeTransformer):
1892
+ def __init__(self, line_offsets: list[int]) -> None:
1893
+ self._line_offsets = line_offsets
1894
+
1895
+ def visit_Call(self, node: ast.Call) -> ast.AST:
1896
+ transformed = self.generic_visit(node)
1897
+ if (
1898
+ not isinstance(transformed, ast.Call)
1899
+ or _shadow_call_kind(transformed) != "read"
1900
+ or transformed.keywords
1901
+ or len(transformed.args) != 1
1902
+ ):
1903
+ return transformed
1904
+ if any(
1905
+ isinstance(child, (ast.Await, ast.Yield, ast.YieldFrom, ast.NamedExpr))
1906
+ for child in ast.walk(transformed)
1907
+ ):
1908
+ return transformed
1909
+ # Passing the callable and evaluated argument keeps `locals()` in the
1910
+ # user cell's frame; a thunk would introduce a lambda frame.
1911
+ site_id = f"py:{_shadow_span(node, self._line_offsets)['start']}"
1912
+ wrapped = ast.Call(
1913
+ func=ast.Name(id=_CALL_SITE_HELPER_NAME, ctx=ast.Load()),
1914
+ args=[
1915
+ ast.Constant(value=site_id),
1916
+ transformed.func,
1917
+ transformed.args[0],
1918
+ ],
1919
+ keywords=[],
1920
+ )
1921
+ return ast.copy_location(wrapped, node)
1922
+
1923
+
1924
+ def _compile_source(
1925
+ source: str, filename: str = "<cell>"
1926
+ ) -> tuple[Any, Any | None, bool]:
1927
+ module = ast.parse(source, filename, "exec")
1928
+ line_offsets = [0]
1929
+ for line in source.splitlines(keepends=True):
1930
+ line_offsets.append(line_offsets[-1] + len(line))
1931
+ if _CALL_SITE_HELPER_NAME not in source or not _cell_binds_call_site_helper(module):
1932
+ module = _ShadowCallSiteTransformer(line_offsets).visit(module)
1933
+ ast.fix_missing_locations(module)
1934
+ if not module.body:
1935
+ return None, None, False
1936
+
1937
+ last = module.body[-1]
1938
+ if isinstance(last, ast.Expr):
1939
+ body_module = ast.Module(body=module.body[:-1], type_ignores=[])
1940
+ expr_module = ast.Expression(body=last.value)
1941
+ ast.copy_location(expr_module, last)
1942
+ body_code = compile(body_module, filename, "exec", flags=_TLA_FLAG)
1943
+ expr_code = compile(expr_module, filename, "eval", flags=_TLA_FLAG)
1944
+ return body_code, expr_code, True
1945
+
1946
+ return compile(module, filename, "exec", flags=_TLA_FLAG), None, False
1947
+
1948
+
1949
+ def _prepare_file_source(
1950
+ source: str, ns: dict, filename: str, linecache_source: str | None
1951
+ ) -> None:
1952
+ display_source = source if linecache_source is None else linecache_source
1953
+ lines = display_source.splitlines(keepends=True)
1954
+ if display_source and (not lines or not lines[-1].endswith(("\n", "\r"))):
1955
+ lines[-1] += "\n"
1956
+ linecache.cache[filename] = (len(display_source), None, lines, filename)
1957
+ ns["__file__"] = filename
1958
+ script_dir = os.path.dirname(filename)
1959
+ if script_dir:
1960
+ try:
1961
+ sys.path.remove(script_dir)
1962
+ except ValueError:
1963
+ pass
1964
+ sys.path.insert(0, script_dir)
1965
+
1966
+
1967
+ def _exec_source(source: str, ns: dict) -> None:
1968
+ """Synchronous source execution for legacy magic helpers."""
1969
+ _restore_call_site_helper(ns)
1970
+ body_code, expr_code, has_expr = _compile_source(source)
1971
+ if body_code is None:
1972
+ return
1973
+ _run_compiled_sync(body_code, ns, want_value=False)
1974
+ if has_expr and expr_code is not None:
1975
+ value = _run_compiled_sync(expr_code, ns, want_value=True)
1976
+ if value is not None:
1977
+ __omp_display(value, kind="result")
1978
+
1979
+
1980
+ async def _exec_source_async(
1981
+ source: str,
1982
+ ns: dict,
1983
+ *,
1984
+ filename: str = "<cell>",
1985
+ linecache_source: str | None = None,
1986
+ ) -> None:
1987
+ """Compile + execute ``source``; if the last node is an expression, route
1988
+ its value through ``__omp_display`` so dataframes/figures render rich.
1989
+ Top-level ``await`` / ``async for`` / ``async with`` is permitted; awaited
1990
+ regions yield to other requests in the runner's persistent event loop."""
1991
+ _restore_call_site_helper(ns)
1992
+ body_code, expr_code, has_expr = _compile_source(source, filename)
1993
+ if filename != "<cell>":
1994
+ _prepare_file_source(source, ns, filename, linecache_source)
1995
+ if body_code is None:
1996
+ return
1997
+ await _run_compiled_async(body_code, ns, want_value=False)
1998
+ if has_expr and expr_code is not None:
1999
+ value = await _run_compiled_async(expr_code, ns, want_value=True)
2000
+ if value is not None:
2001
+ __omp_display(value, kind="result")
2002
+
2003
+
2004
+ # ---------------------------------------------------------------------------
2005
+ # Signal handling
2006
+ # ---------------------------------------------------------------------------
2007
+
2008
+
2009
+ def _install_idle_sigint() -> None:
2010
+ try:
2011
+ signal.signal(signal.SIGINT, signal.SIG_IGN)
2012
+ except (OSError, ValueError):
2013
+ # Some platforms (Windows in non-console mode) reject this; fine.
2014
+ pass
2015
+
2016
+
2017
+ def _install_exec_sigint() -> None:
2018
+ try:
2019
+ signal.signal(signal.SIGINT, signal.default_int_handler)
2020
+ except (OSError, ValueError):
2021
+ pass
2022
+
2023
+
2024
+ def _begin_exec_sigint() -> None:
2025
+ _STATE.active_executions += 1
2026
+ _install_exec_sigint()
2027
+
2028
+
2029
+ def _end_exec_sigint() -> None:
2030
+ if _STATE.active_executions > 0:
2031
+ _STATE.active_executions -= 1
2032
+ if _STATE.active_executions == 0:
2033
+ _install_idle_sigint()
2034
+
2035
+
2036
+ _MANAGED_ENV_KEYS = (
2037
+ "PI_SESSION_FILE",
2038
+ "PI_ARTIFACTS_DIR",
2039
+ "PI_TOOL_BRIDGE_URL",
2040
+ "PI_TOOL_BRIDGE_TOKEN",
2041
+ "PI_TOOL_BRIDGE_SESSION",
2042
+ "PI_EVAL_LOCAL_ROOTS",
2043
+ )
2044
+
2045
+
2046
+ def _apply_request_runtime(req: dict) -> None:
2047
+ cwd = req.get("cwd")
2048
+ if isinstance(cwd, str) and cwd:
2049
+ os.chdir(cwd)
2050
+ try:
2051
+ sys.path.remove(cwd)
2052
+ except ValueError:
2053
+ pass
2054
+ sys.path.insert(0, cwd)
2055
+
2056
+ env = req.get("env")
2057
+ if isinstance(env, dict):
2058
+ for key in _MANAGED_ENV_KEYS:
2059
+ value = env.get(key)
2060
+ if isinstance(value, str):
2061
+ os.environ[key] = value
2062
+ elif value is None:
2063
+ os.environ.pop(key, None)
2064
+
2065
+
2066
+ def _start_parent_watchdog() -> None:
2067
+ """Self-terminate when the host process dies.
2068
+
2069
+ The main loop only exits when stdin EOFs, which only happens once user
2070
+ code finishes and the next ``readline`` call returns. If the host gets
2071
+ SIGKILL mid-execution (or any way that skips graceful shutdown) the
2072
+ runner would otherwise outlive its parent and keep holding kernel
2073
+ state. Poll ``os.getppid()`` instead and ``os._exit`` the moment we get
2074
+ reparented \u2014 covers POSIX hosts. Windows has no reliable ppid
2075
+ equivalent; there we still bail out on the next stdin read.
2076
+ """
2077
+ if os.name != "posix":
2078
+ return
2079
+ original_ppid = os.getppid()
2080
+ if original_ppid <= 1:
2081
+ return
2082
+
2083
+ def watch() -> None:
2084
+ while True:
2085
+ try:
2086
+ if os.getppid() != original_ppid:
2087
+ os._exit(0)
2088
+ except Exception:
2089
+ return
2090
+ time.sleep(10)
2091
+
2092
+ thread = threading.Thread(target=watch, name="omp-parent-watchdog", daemon=True)
2093
+ thread.start()
2094
+
2095
+
2096
+ # ---------------------------------------------------------------------------
2097
+ # Request dispatch
2098
+ # ---------------------------------------------------------------------------
2099
+
2100
+
2101
+ def _jsonable(value: Any) -> Any:
2102
+ """Convert a tool result into a JSON-safe value without failing the request."""
2103
+ if value is None or isinstance(value, (str, int, float, bool)):
2104
+ return value
2105
+ if isinstance(value, dict):
2106
+ return {str(key): _jsonable(item) for key, item in value.items()}
2107
+ if isinstance(value, (list, tuple, set)):
2108
+ return [_jsonable(item) for item in value]
2109
+ return repr(value)
2110
+
2111
+
2112
+ def _handle_tool_request(req: dict) -> None:
2113
+ """Describe or invoke a kernel-defined tool on a dedicated runner thread."""
2114
+ rid = str(req.get("id"))
2115
+ token = _CURRENT_RID.set(rid)
2116
+ status = "ok"
2117
+ _emit({"type": "started", "id": rid})
2118
+ try:
2119
+ tools = _STATE.user_ns.get("__omp_tools__") or {}
2120
+ op = req.get("op")
2121
+ if op == "describe":
2122
+ requested = req.get("names") or list(tools)
2123
+ envelope = {
2124
+ "ok": True,
2125
+ "tools": [tools[name].describe() for name in requested if name in tools],
2126
+ "missing": [name for name in requested if name not in tools],
2127
+ }
2128
+ elif op == "call":
2129
+ name = str(req.get("name") or "")
2130
+ spec = tools.get(name)
2131
+ if spec is None:
2132
+ raise KeyError(f"tool {name!r} is not defined")
2133
+ result = spec.fn(**(req.get("args") or {}))
2134
+ if inspect.isawaitable(result):
2135
+ result = asyncio.run(result)
2136
+ envelope = {"ok": True, "value": _jsonable(result)}
2137
+ else:
2138
+ raise ValueError(f"unknown tool operation {op!r}")
2139
+ _emit_display({"application/json": envelope})
2140
+ except BaseException as exc: # noqa: BLE001 - protocol errors must settle the request
2141
+ status = "error"
2142
+ _emit_error(rid, exc)
2143
+ finally:
2144
+ _flush_stream_proxies(rid)
2145
+ _emit(
2146
+ {
2147
+ "type": "done",
2148
+ "id": rid,
2149
+ "status": status,
2150
+ "executionCount": _STATE.execution_count,
2151
+ "cancelled": False,
2152
+ }
2153
+ )
2154
+ _CURRENT_RID.reset(token)
2155
+
2156
+
2157
+ async def _handle_request_async(req: dict) -> None:
2158
+ rid = str(req.get("id"))
2159
+ execution_reserved = bool(req.pop("_execution_reserved", False))
2160
+ execution_started = False
2161
+ token = _CURRENT_RID.set(rid)
2162
+ displayed_matplotlib_token = _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.set(set())
2163
+ _STATE.capture_rid = rid
2164
+ reset_call_occurrences = _STATE.user_ns.get("__omp_reset_call_occurrences__")
2165
+ if callable(reset_call_occurrences):
2166
+ reset_call_occurrences()
2167
+ _STATE.namespace_revision += 1
2168
+ _STATE.user_ns["__omp_run_id__"] = rid
2169
+ _STATE.cancel_requested = False
2170
+ _STATE.execution_count += 1
2171
+ execution_count = _STATE.execution_count
2172
+ _emit({"type": "started", "id": rid})
2173
+
2174
+ status: str = "ok"
2175
+ cancelled = False
2176
+
2177
+ try:
2178
+ try:
2179
+ _apply_request_runtime(req)
2180
+ source = req.get("code", "")
2181
+ transformed = transform_cell(source)
2182
+ requested_filename = req.get("filename")
2183
+ filename = (
2184
+ requested_filename
2185
+ if isinstance(requested_filename, str) and requested_filename
2186
+ else "<cell>"
2187
+ )
2188
+ except SyntaxError as exc:
2189
+ _emit_error(rid, exc)
2190
+ _emit(
2191
+ {
2192
+ "type": "done",
2193
+ "id": rid,
2194
+ "status": "error",
2195
+ "executionCount": execution_count,
2196
+ "cancelled": False,
2197
+ }
2198
+ )
2199
+ return
2200
+ except BaseException as exc: # noqa: BLE001 - runtime setup errors must settle the request
2201
+ _emit_error(rid, exc)
2202
+ _emit(
2203
+ {
2204
+ "type": "done",
2205
+ "id": rid,
2206
+ "status": "error",
2207
+ "executionCount": execution_count,
2208
+ "cancelled": False,
2209
+ }
2210
+ )
2211
+ return
2212
+
2213
+ if execution_reserved:
2214
+ _install_exec_sigint()
2215
+ else:
2216
+ _begin_exec_sigint()
2217
+ execution_started = True
2218
+ try:
2219
+ await _exec_source_async(
2220
+ transformed,
2221
+ _STATE.user_ns,
2222
+ filename=filename,
2223
+ linecache_source=source,
2224
+ )
2225
+ except KeyboardInterrupt:
2226
+ cancelled = True
2227
+ status = "error"
2228
+ _emit_error(rid, KeyboardInterrupt("Execution interrupted"))
2229
+ except SystemExit as exc:
2230
+ status = "error"
2231
+ _emit_error(rid, exc)
2232
+ except BaseException as exc: # noqa: BLE001 - we want to surface every user error
2233
+ status = "error"
2234
+ _emit_error(rid, exc, source_filename=filename)
2235
+ finally:
2236
+ _end_exec_sigint()
2237
+ try:
2238
+ _flush_matplotlib_figures()
2239
+ except Exception:
2240
+ pass
2241
+
2242
+ _flush_stream_proxies(rid)
2243
+ _emit(
2244
+ {
2245
+ "type": "done",
2246
+ "id": rid,
2247
+ "status": status,
2248
+ "executionCount": execution_count,
2249
+ "cancelled": cancelled,
2250
+ }
2251
+ )
2252
+ finally:
2253
+ if _STATE.capture_rid == rid:
2254
+ _STATE.capture_rid = None
2255
+ _flush_stream_proxies(rid)
2256
+ _CURRENT_RID.reset(token)
2257
+ _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.reset(displayed_matplotlib_token)
2258
+ if execution_reserved and not execution_started:
2259
+ _end_exec_sigint()
2260
+
2261
+
2262
+ def _emit_error(rid: str, exc: BaseException, source_filename: str = "<cell>") -> None:
2263
+ if isinstance(exc, SyntaxError) and exc.filename in ("<cell>", source_filename):
2264
+ # Syntax error in the executed source itself: every stack frame is
2265
+ # runner machinery, so emit only the caret display, like a REPL.
2266
+ tb_lines = traceback.format_exception_only(type(exc), exc)
2267
+ else:
2268
+ # Magics execute between the cell and loaded script. Omit runner
2269
+ # frames throughout that chain, not just before the first user frame.
2270
+ # Internal-only failures retain their full traceback for diagnosis.
2271
+ formatted = traceback.TracebackException.from_exception(exc)
2272
+ user_frames = [frame for frame in formatted.stack if frame.filename != __file__]
2273
+ if user_frames:
2274
+ formatted.stack = traceback.StackSummary.from_list(user_frames)
2275
+ tb_lines = list(formatted.format())
2276
+ if isinstance(exc, ModuleNotFoundError) and exc.name:
2277
+ tb_lines.append(
2278
+ f"Install the distribution that provides module {exc.name!r} with "
2279
+ "`%pip install <distribution-name>` in eval. "
2280
+ "Distribution names can differ from import names."
2281
+ )
2282
+ _emit(
2283
+ {
2284
+ "type": "error",
2285
+ "id": rid,
2286
+ "ename": type(exc).__name__,
2287
+ "evalue": str(exc),
2288
+ "traceback": [line.rstrip("\n") for line in tb_lines],
2289
+ }
2290
+ )
2291
+
2292
+
2293
+ def _admit_shadow_run(req: dict) -> bool:
2294
+ expected_revision = req.get("expectedShadowRevision")
2295
+ expected_digest = req.get("expectedShadowDigest")
2296
+ if expected_revision is not None or expected_digest is not None:
2297
+ values = _snapshot_user_namespace()
2298
+ eligible = (
2299
+ _STATE.active_executions == 0
2300
+ and expected_revision == _STATE.namespace_revision
2301
+ and expected_digest == _shadow_snapshot_digest(values)
2302
+ )
2303
+ if not eligible:
2304
+ _emit(
2305
+ {
2306
+ "type": "done",
2307
+ "id": str(req.get("id")),
2308
+ "status": "ok",
2309
+ "executionCount": _STATE.execution_count,
2310
+ "cancelled": False,
2311
+ "admissionRejected": True,
2312
+ }
2313
+ )
2314
+ return False
2315
+ _STATE.active_executions += 1
2316
+ req["_execution_reserved"] = True
2317
+ return True
2318
+
2319
+ # ---------------------------------------------------------------------------
2320
+ # Main loop
2321
+ # ---------------------------------------------------------------------------
2322
+
2323
+
2324
+ def _parse_request(line: str) -> dict | None:
2325
+ """Parse one NDJSON control line, or ``None`` when it should be skipped.
2326
+
2327
+ Emits a ``ProtocolError`` frame for malformed JSON so the host sees the
2328
+ same diagnostic the reader has always produced.
2329
+ """
2330
+ try:
2331
+ return json.loads(line)
2332
+ except json.JSONDecodeError as exc:
2333
+ _emit(
2334
+ {
2335
+ "type": "error",
2336
+ "id": "",
2337
+ "ename": "ProtocolError",
2338
+ "evalue": f"Invalid JSON request: {exc}",
2339
+ "traceback": [],
2340
+ }
2341
+ )
2342
+ return None
2343
+
2344
+
2345
+ def _read_stdin(loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, stdin) -> None:
2346
+ """Feed control-channel requests to the event loop from a reader thread.
2347
+
2348
+ POSIX only -- see ``_serve_windows`` for why Windows cannot run a
2349
+ perpetually blocked stdin reader.
2350
+ """
2351
+ for raw_line in stdin:
2352
+ line = raw_line.strip()
2353
+ if not line:
2354
+ continue
2355
+ req = _parse_request(line)
2356
+ if req is None:
2357
+ continue
2358
+ if req.get("type") == "tool":
2359
+ threading.Thread(
2360
+ target=_handle_tool_request,
2361
+ args=(req,),
2362
+ name=f"omp-tool-{req.get('id')}",
2363
+ daemon=True,
2364
+ ).start()
2365
+ continue
2366
+ loop.call_soon_threadsafe(queue.put_nowait, req)
2367
+ loop.call_soon_threadsafe(queue.put_nowait, {"type": "exit"})
2368
+
2369
+
2370
+ async def _serve_posix(loop: asyncio.AbstractEventLoop, stdin) -> None:
2371
+ """Dispatch requests concurrently, one asyncio task per request.
2372
+
2373
+ A background thread reads stdin and enqueues requests so a cell parked on
2374
+ a top-level ``await`` (an ``await agent(...)`` bridge call, say) does not
2375
+ block sibling requests: auto-backgrounded cells, user Python shortcuts, and
2376
+ kernel-defined tool calls from subagents can all be in flight on one kernel
2377
+ at once. The reader thread stays blocked in a
2378
+ ``sys.stdin`` read for its whole life, which is safe on POSIX but wedges
2379
+ native-extension imports on Windows (see ``_serve_windows``).
2380
+ """
2381
+ queue: asyncio.Queue = asyncio.Queue()
2382
+ reader = threading.Thread(
2383
+ target=_read_stdin,
2384
+ args=(loop, queue, stdin),
2385
+ name="omp-stdin-reader",
2386
+ daemon=True,
2387
+ )
2388
+ reader.start()
2389
+
2390
+ tasks: set[asyncio.Task] = set()
2391
+
2392
+ def _task_done(task: asyncio.Task) -> None:
2393
+ tasks.discard(task)
2394
+ try:
2395
+ exc = task.exception()
2396
+ except asyncio.CancelledError:
2397
+ return
2398
+ if exc is not None:
2399
+ _emit_error("", exc)
2400
+
2401
+ try:
2402
+ while True:
2403
+ req = await queue.get()
2404
+ if req.get("type") == "exit":
2405
+ break
2406
+ if req.get("type") == "shadow_snapshot":
2407
+ _emit_shadow_snapshot(req)
2408
+ continue
2409
+ if req.get("type") == "shadow_plan":
2410
+ _emit_shadow_plan(req)
2411
+ continue
2412
+ if not _admit_shadow_run(req):
2413
+ continue
2414
+ task = asyncio.create_task(_handle_request_async(req))
2415
+ tasks.add(task)
2416
+ task.add_done_callback(_task_done)
2417
+ finally:
2418
+ for task in tasks:
2419
+ task.cancel()
2420
+ if tasks:
2421
+ await asyncio.gather(*tasks, return_exceptions=True)
2422
+
2423
+
2424
+ async def _serve_windows(loop: asyncio.AbstractEventLoop, stdin) -> None:
2425
+ """Dispatch requests serially, serving tool requests between cells.
2426
+
2427
+ A thread perpetually parked in a blocking ``sys.stdin`` read deadlocks
2428
+ native-extension imports (NumPy in particular) under a pipe-backed
2429
+ subprocess on Windows: the native DLL load and the concurrent stdin read
2430
+ wedge each other (numpy#24290, issue #7985). Reading each line via
2431
+ ``run_in_executor`` confines the stdin read to the gap *between* requests
2432
+ -- once a line arrives the reader returns, so no thread is reading stdin
2433
+ while a cell (and its imports) run. The trade-off is that requests are
2434
+ handled one at a time instead of interleaved; on Windows that is strictly
2435
+ better than the alternative, since any concurrent reader deadlocks the
2436
+ import outright.
2437
+ """
2438
+ while True:
2439
+ raw_line = await loop.run_in_executor(None, stdin.readline)
2440
+ if not raw_line:
2441
+ break # EOF: the host closed the control channel.
2442
+ line = raw_line.strip()
2443
+ if not line:
2444
+ continue
2445
+ req = _parse_request(line)
2446
+ if req is None:
2447
+ continue
2448
+ if req.get("type") == "exit":
2449
+ break
2450
+ if req.get("type") == "shadow_snapshot":
2451
+ _emit_shadow_snapshot(req)
2452
+ continue
2453
+ if req.get("type") == "shadow_plan":
2454
+ _emit_shadow_plan(req)
2455
+ continue
2456
+ if req.get("type") == "tool":
2457
+ threading.Thread(
2458
+ target=_handle_tool_request,
2459
+ args=(req,),
2460
+ name=f"omp-tool-{req.get('id')}",
2461
+ daemon=True,
2462
+ ).start()
2463
+ continue
2464
+ if not _admit_shadow_run(req):
2465
+ continue
2466
+ try:
2467
+ await _handle_request_async(req)
2468
+ except asyncio.CancelledError:
2469
+ raise # task cancellation must propagate; never swallow it and spin
2470
+ except BaseException as exc: # noqa: BLE001 - one bad request must not wedge the loop
2471
+ _emit_error(str(req.get("id", "")), exc)
2472
+
2473
+
2474
+ async def _main_async() -> None:
2475
+ sys.stdout = _StreamProxy("stdout")
2476
+ sys.stderr = _StreamProxy("stderr")
2477
+ _install_idle_sigint()
2478
+ _start_parent_watchdog()
2479
+ _start_capture_drain()
2480
+
2481
+ stdin = sys.__stdin__
2482
+ if stdin is None:
2483
+ return
2484
+
2485
+ loop = asyncio.get_running_loop()
2486
+ _STATE.loop = loop
2487
+
2488
+ if os.name == "nt":
2489
+ await _serve_windows(loop, stdin)
2490
+ else:
2491
+ await _serve_posix(loop, stdin)
2492
+
2493
+
2494
+ def main() -> None:
2495
+ asyncio.run(_main_async())
2496
+
2497
+
2498
+ if __name__ == "__main__":
2499
+ main()