From 9a2cc3bda0415f91f95b3ae31228bfe8ea06e964 Mon Sep 17 00:00:00 2001 From: can1357 Date: Tue, 26 May 2026 15:26:29 +0200 Subject: [PATCH] feat(eval/py): added runtime environment and working directory support to Python kernel execution - Added `cwd` and `env` optional parameters to kernel execution API for runtime working directory and environment variable control. - Implemented runtime environment setup in Python runner with `_apply_request_runtime()` to apply cwd and env from request before code execution. - Enhanced SIGINT handler management with `active_executions` counter and `_begin_exec_sigint()` / `_end_exec_sigint()` functions to prevent state mutation during concurrent execution. - Changed `SearchRenderArgs.paths` parameter type from `string[]` to `string | string[]` to accept single string paths. - Added comprehensive test coverage for kernel cwd updates, timeout interruption safety, and SystemExit handling in shared executor sessions. --- .../eval/__tests__/shared-executors.test.ts | 60 ++++++++++++++ packages/coding-agent/src/eval/py/executor.ts | 39 +++++++-- packages/coding-agent/src/eval/py/kernel.ts | 8 ++ packages/coding-agent/src/eval/py/runner.py | 79 +++++++++++++++---- packages/coding-agent/src/tools/search.ts | 7 +- .../test/tools/search-path-lists.test.ts | 27 +++++++ 6 files changed, 195 insertions(+), 25 deletions(-) diff --git a/packages/coding-agent/src/eval/__tests__/shared-executors.test.ts b/packages/coding-agent/src/eval/__tests__/shared-executors.test.ts index 7711154d4..3a01565b4 100644 --- a/packages/coding-agent/src/eval/__tests__/shared-executors.test.ts +++ b/packages/coding-agent/src/eval/__tests__/shared-executors.test.ts @@ -151,6 +151,66 @@ describe("shared eval executors", () => { expect(result.output.trim()).toBe("42"); }); + it("updates Python cwd when one shared session id runs from multiple directories", async () => { + using tempDir = TempDir.createSync("@omp-eval-py-cwd-"); + const dirA = path.join(tempDir.path(), "a"); + const dirB = path.join(tempDir.path(), "b"); + await fs.mkdir(dirA); + await fs.mkdir(dirB); + const realDirA = await fs.realpath(dirA); + const realDirB = await fs.realpath(dirB); + const sessionFile = path.join(tempDir.path(), "session.jsonl"); + const sessionId = `py-cwd:${crypto.randomUUID()}`; + + const first = await executePython("import os\nprint(os.getcwd())", { cwd: dirA, sessionId, sessionFile }); + const second = await executePython("import os\nprint(os.getcwd())", { cwd: dirB, sessionId, sessionFile }); + + expect(first.exitCode).toBe(0); + expect(first.output.trim()).toBe(realDirA); + expect(second.exitCode).toBe(0); + expect(second.output.trim()).toBe(realDirB); + }); + + it("interrupts timed out synchronous Python cells before they mutate shared state", async () => { + using tempDir = TempDir.createSync("@omp-eval-py-sync-timeout-"); + const sessionFile = path.join(tempDir.path(), "session.jsonl"); + const sessionId = `py-sync-timeout:${crypto.randomUUID()}`; + + const timedOut = await executePython("import time\ntime.sleep(0.2)\nleaked_after_timeout = True", { + cwd: tempDir.path(), + sessionId, + sessionFile, + timeoutMs: 20, + }); + await Bun.sleep(250); + const probe = await executePython('print("leaked_after_timeout" in globals())', { + cwd: tempDir.path(), + sessionId, + sessionFile, + }); + + expect(timedOut.cancelled).toBe(true); + expect(probe.exitCode).toBe(0); + expect(probe.output.trim()).toBe("False"); + }); + + it("settles Python cells that raise SystemExit", async () => { + using tempDir = TempDir.createSync("@omp-eval-py-system-exit-"); + const sessionFile = path.join(tempDir.path(), "session.jsonl"); + const sessionId = `py-system-exit:${crypto.randomUUID()}`; + + const result = await executePython('raise SystemExit("bye")', { + cwd: tempDir.path(), + sessionId, + sessionFile, + timeoutMs: 500, + }); + + expect(result.exitCode).toBe(1); + expect(result.output).toContain("SystemExit"); + expect(result.output).toContain("bye"); + }); + it("lets a subagent inherit parent JavaScript and Python eval state", async () => { using tempDir = TempDir.createSync("@omp-eval-subagent-"); const sessionFile = path.join(tempDir.path(), "session.jsonl"); diff --git a/packages/coding-agent/src/eval/py/executor.ts b/packages/coding-agent/src/eval/py/executor.ts index 00cb2feeb..37d967420 100644 --- a/packages/coding-agent/src/eval/py/executor.ts +++ b/packages/coding-agent/src/eval/py/executor.ts @@ -4,11 +4,12 @@ import { OutputSink } from "../../session/streaming-output"; import type { ToolSession } from "../../tools"; import { resolveOutputMaxColumns, resolveOutputSinkHeadBytes } from "../../tools/output-meta"; import type { JsStatusEvent } from "../js/shared/types"; -import type { KernelDisplayOutput } from "./display"; import { checkPythonKernelAvailability, + type KernelDisplayOutput, type KernelExecuteOptions, type KernelExecuteResult, + type KernelRuntimeEnv, PythonKernel, } from "./kernel"; import { ensurePyToolBridge, registerPyToolBridge } from "./tool-bridge"; @@ -240,19 +241,40 @@ function createCancelledPythonResult(timedOut: boolean, timeoutMs?: number): Pyt // Kernel start helpers // --------------------------------------------------------------------------- +const MANAGED_KERNEL_ENV_KEYS = [ + "PI_SESSION_FILE", + "PI_ARTIFACTS_DIR", + "PI_TOOL_BRIDGE_URL", + "PI_TOOL_BRIDGE_TOKEN", + "PI_TOOL_BRIDGE_SESSION", +] as const; + +function buildKernelEnvPatch(options: { + sessionFile?: string; + artifactsDir?: string; + bridgeSessionId?: string; + bridge?: { url: string; token: string }; +}): KernelRuntimeEnv { + return { + PI_SESSION_FILE: options.sessionFile ?? null, + PI_ARTIFACTS_DIR: options.artifactsDir ?? null, + PI_TOOL_BRIDGE_URL: options.bridge?.url ?? null, + PI_TOOL_BRIDGE_TOKEN: options.bridge?.token ?? null, + PI_TOOL_BRIDGE_SESSION: options.bridge && options.bridgeSessionId ? options.bridgeSessionId : null, + }; +} + function buildKernelEnv(options: { sessionFile?: string; artifactsDir?: string; bridgeSessionId?: string; bridge?: { url: string; token: string }; }): Record | undefined { + const patch = buildKernelEnvPatch(options); const env: Record = {}; - if (options.sessionFile) env.PI_SESSION_FILE = options.sessionFile; - if (options.artifactsDir) env.PI_ARTIFACTS_DIR = options.artifactsDir; - if (options.bridge && options.bridgeSessionId) { - env.PI_TOOL_BRIDGE_URL = options.bridge.url; - env.PI_TOOL_BRIDGE_TOKEN = options.bridge.token; - env.PI_TOOL_BRIDGE_SESSION = options.bridgeSessionId; + for (const key of MANAGED_KERNEL_ENV_KEYS) { + const value = patch[key]; + if (value !== null) env[key] = value; } return Object.keys(env).length > 0 ? env : undefined; } @@ -415,6 +437,8 @@ async function executeWithKernel( try { executionTimeoutMs = requireRemainingTimeoutMs(deadlineMs); const result = await kernel.execute(code, { + cwd: options?.cwd, + env: buildKernelEnvPatch(options ?? {}), id: runId, signal: options?.signal, timeoutMs: executionTimeoutMs, @@ -563,6 +587,7 @@ export async function executePython(code: string, options?: PythonExecutorOption const deadlineMs = getExecutionDeadlineMs(options); const executionOptions: PythonExecutorOptions = { ...(options ?? {}), + cwd, deadlineMs, }; diff --git a/packages/coding-agent/src/eval/py/kernel.ts b/packages/coding-agent/src/eval/py/kernel.ts index a3ed9869d..60b93c73c 100644 --- a/packages/coding-agent/src/eval/py/kernel.ts +++ b/packages/coding-agent/src/eval/py/kernel.ts @@ -51,8 +51,14 @@ const STARTUP_TIMEOUT_MS = 10_000; // kernel's state, so we only kill as a last-resort recovery path. const INTERRUPT_ESCALATION_MS = 5_000; +export type KernelRuntimeEnv = Record; + export interface KernelExecuteOptions { id?: string; + /** Runtime working directory applied immediately before this request executes. */ + cwd?: string; + /** Managed runtime environment variables applied immediately before this request executes. */ + env?: KernelRuntimeEnv; signal?: AbortSignal; onChunk?: (text: string) => Promise | void; onDisplay?: (output: KernelDisplayOutput) => Promise | void; @@ -346,6 +352,8 @@ export class PythonKernel { const payload = JSON.stringify({ id: msgId, code, + cwd: options?.cwd, + env: options?.env, silent: options?.silent ?? false, storeHistory: options?.storeHistory ?? !(options?.silent ?? false), }); diff --git a/packages/coding-agent/src/eval/py/runner.py b/packages/coding-agent/src/eval/py/runner.py index 8ce83b1d2..25bbc11d9 100644 --- a/packages/coding-agent/src/eval/py/runner.py +++ b/packages/coding-agent/src/eval/py/runner.py @@ -5,6 +5,7 @@ wrapper writes typed frames back. Host -> wrapper: {"id": str, "code": str, "silent": bool?, "storeHistory": bool?} + {"id": str, "code": str, "silent": bool?, "storeHistory": bool?, "cwd": str?, "env": dict?} {"type": "exit"} # graceful shutdown Wrapper -> host: @@ -27,8 +28,8 @@ from __future__ import annotations import asyncio import ast -import base64 import contextvars +import base64 import builtins import inspect import io @@ -44,7 +45,7 @@ import threading import time import traceback from pathlib import Path -from typing import Any, Callable +from typing import Any # --------------------------------------------------------------------------- # Frame writer @@ -123,6 +124,7 @@ class _RunnerState: } self.last_install_marker: int = 0 self.loop: asyncio.AbstractEventLoop | None = None + self.active_executions: int = 0 _CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar("omp_current_rid", default=None) @@ -718,25 +720,20 @@ def _run_compiled_sync(code, ns: dict, *, want_value: bool) -> Any: return None -async def _run_in_executor_with_context(fn: Callable[[], Any]) -> Any: - loop = asyncio.get_running_loop() - ctx = contextvars.copy_context() - return await loop.run_in_executor(None, ctx.run, fn) - async def _run_compiled_async(code, ns: dict, *, want_value: bool) -> Any: """Execute a code object in the persistent event loop. Coroutine code is awaited in this task so top-level ``await`` interleaves - with sibling requests. Plain statement/expression code runs in the default - executor with the current ContextVar state copied into that thread. + with sibling requests. Plain statement/expression code runs on the main + runner thread so SIGINT can interrupt it reliably. """ if code.co_flags & inspect.CO_COROUTINE: result = await eval(code, ns) return result if want_value else None if want_value: - return await _run_in_executor_with_context(lambda: eval(code, ns)) - await _run_in_executor_with_context(lambda: exec(code, ns)) + return eval(code, ns) + exec(code, ns) return None @@ -804,6 +801,46 @@ def _install_exec_sigint() -> None: pass +def _begin_exec_sigint() -> None: + _STATE.active_executions += 1 + _install_exec_sigint() + + +def _end_exec_sigint() -> None: + if _STATE.active_executions > 0: + _STATE.active_executions -= 1 + if _STATE.active_executions == 0: + _install_idle_sigint() + + +_MANAGED_ENV_KEYS = ( + "PI_SESSION_FILE", + "PI_ARTIFACTS_DIR", + "PI_TOOL_BRIDGE_URL", + "PI_TOOL_BRIDGE_TOKEN", + "PI_TOOL_BRIDGE_SESSION", +) + + +def _apply_request_runtime(req: dict) -> None: + cwd = req.get("cwd") + if isinstance(cwd, str) and cwd: + os.chdir(cwd) + try: + sys.path.remove(cwd) + except ValueError: + pass + sys.path.insert(0, cwd) + + env = req.get("env") + if isinstance(env, dict): + for key in _MANAGED_ENV_KEYS: + value = env.get(key) + if isinstance(value, str): + os.environ[key] = value + elif value is None: + os.environ.pop(key, None) + def _start_parent_watchdog() -> None: """Self-terminate when the host process dies. @@ -853,6 +890,7 @@ async def _handle_request_async(req: dict) -> None: try: try: + _apply_request_runtime(req) transformed = transform_cell(req.get("code", "")) except SyntaxError as exc: _emit_error(rid, exc) @@ -864,21 +902,32 @@ async def _handle_request_async(req: dict) -> None: "cancelled": False, }) return + except BaseException as exc: # noqa: BLE001 - runtime setup errors must settle the request + _emit_error(rid, exc) + _emit({ + "type": "done", + "id": rid, + "status": "error", + "executionCount": execution_count, + "cancelled": False, + }) + return - _install_exec_sigint() + _begin_exec_sigint() try: await _exec_source_async(transformed, _STATE.user_ns) except KeyboardInterrupt: cancelled = True status = "error" _emit_error(rid, KeyboardInterrupt("Execution interrupted")) - except SystemExit: - raise + except SystemExit as exc: + status = "error" + _emit_error(rid, exc) except BaseException as exc: # noqa: BLE001 - we want to surface every user error status = "error" _emit_error(rid, exc) finally: - _install_idle_sigint() + _end_exec_sigint() try: _flush_matplotlib_figures() except Exception: diff --git a/packages/coding-agent/src/tools/search.ts b/packages/coding-agent/src/tools/search.ts index 9ca63b1ab..149360354 100644 --- a/packages/coding-agent/src/tools/search.ts +++ b/packages/coding-agent/src/tools/search.ts @@ -9,7 +9,7 @@ import { prompt, untilAborted } from "@oh-my-pi/pi-utils"; import * as z from "zod/v4"; import { getFileReadCache } from "../edit/file-read-cache"; import type { RenderResultOptions } from "../extensibility/custom-tools/types"; -import { formatHashlineHeader } from "../hashline/hash"; +import { computeFileHash, formatHashlineHeader } from "../hashline/hash"; import type { Theme } from "../modes/theme/theme"; import searchDescription from "../prompts/tools/search.md" with { type: "text" }; import { DEFAULT_MAX_COLUMN, type TruncationResult, truncateHead } from "../session/streaming-output"; @@ -634,7 +634,7 @@ export class SearchTool implements AgentTool 0) meta.push(`skip:${args.skip}`); diff --git a/packages/coding-agent/test/tools/search-path-lists.test.ts b/packages/coding-agent/test/tools/search-path-lists.test.ts index 077e56ecf..2f94a13cb 100644 --- a/packages/coding-agent/test/tools/search-path-lists.test.ts +++ b/packages/coding-agent/test/tools/search-path-lists.test.ts @@ -4,8 +4,12 @@ import * as os from "node:os"; import * as path from "node:path"; import { validateToolArguments } from "@oh-my-pi/pi-ai/utils/validation"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import type { RenderResultOptions } from "@oh-my-pi/pi-coding-agent/extensibility/custom-tools/types"; +import type { Theme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import { ToolChoiceQueue } from "@oh-my-pi/pi-coding-agent/session/tool-choice-queue"; import { createTools, type ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; +import { searchToolRenderer } from "@oh-my-pi/pi-coding-agent/tools/search"; +import { Text } from "@oh-my-pi/pi-tui"; function createTestSession(cwd: string, overrides: Partial = {}): ToolSession { return { @@ -18,6 +22,18 @@ function createTestSession(cwd: string, overrides: Partial = {}): T }; } +const plainTheme = { + fg: (_color: unknown, text: string) => text, + styledSymbol: () => "…", + sep: { dot: " • " }, + format: { bracketLeft: "[", bracketRight: "]" }, +} as unknown as Theme; + +const renderOptions: RenderResultOptions = { + expanded: false, + isPartial: true, +}; + function getText(result: { content: Array<{ type: string; text?: string }> }): string { return result.content .filter(entry => entry.type === "text") @@ -115,6 +131,17 @@ describe("tool path arrays", () => { expect(details?.scopePath).toBe("folder with spaces"); }); + it("search pending renderer accepts a single string path", () => { + const component = searchToolRenderer.renderCall( + { pattern: "space-needle", paths: "folder with spaces/" }, + renderOptions, + plainTheme, + ); + + expect(component).toBeInstanceOf(Text); + expect((component as Text).getText()).toContain("in folder with spaces/"); + }); + it("search keeps a single path that contains spaces", async () => { const tools = await createTools(createTestSession(tempDir)); const tool = tools.find(entry => entry.name === "search");