From 53e75e33c5609d6f8fbc908503008bb6e35e6e8e Mon Sep 17 00:00:00 2001 From: can1357 Date: Sun, 18 Jan 2026 21:36:44 +0100 Subject: [PATCH] feat(coding-agent/core): loaded Python modules from user and project directories with priority - Added Python prelude extension module loading from user and project directories. - Implemented automatic discovery of Python modules from `.omp/modules` and `.pi/modules` directories. - Added prioritized module loading with project-level modules overriding user-level modules. - Updated kernel initialization methods to load Python modules during startup. --- packages/coding-agent/CHANGELOG.md | 5 + .../src/core/python-kernel.test.ts | 380 +++++++----------- .../coding-agent/src/core/python-kernel.ts | 10 +- .../src/core/python-modules.test.ts | 102 +++++ .../coding-agent/src/core/python-modules.ts | 110 +++++ .../test/fixtures/python-modules/alpha.py | 2 + .../test/fixtures/python-modules/beta.py | 2 + .../test/fixtures/python-modules/delta.py | 2 + .../test/fixtures/python-modules/gamma.py | 2 + 9 files changed, 378 insertions(+), 237 deletions(-) create mode 100644 packages/coding-agent/src/core/python-modules.test.ts create mode 100644 packages/coding-agent/src/core/python-modules.ts create mode 100644 packages/coding-agent/test/fixtures/python-modules/alpha.py create mode 100644 packages/coding-agent/test/fixtures/python-modules/beta.py create mode 100644 packages/coding-agent/test/fixtures/python-modules/delta.py create mode 100644 packages/coding-agent/test/fixtures/python-modules/gamma.py diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index e91bbbb85..3484529c3 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,6 +1,11 @@ # Changelog ## [Unreleased] +### Added + +- Added support for loading Python prelude extension modules from user and project directories +- Added automatic discovery of Python modules from `.omp/modules` and `.pi/modules` directories +- Added prioritized module loading with project-level modules overriding user-level modules ## [5.6.7] - 2026-01-18 diff --git a/packages/coding-agent/src/core/python-kernel.test.ts b/packages/coding-agent/src/core/python-kernel.test.ts index e5a675281..68f5c1078 100644 --- a/packages/coding-agent/src/core/python-kernel.test.ts +++ b/packages/coding-agent/src/core/python-kernel.test.ts @@ -67,6 +67,39 @@ function decodeMessage(data: ArrayBuffer): JupyterMessage { return JSON.parse(msgText) as JupyterMessage; } +function sendOkExecution(ws: FakeWebSocket, msgId: string, executionCount = 1) { + const reply: JupyterMessage = { + channel: "shell", + header: { + msg_id: `reply-${msgId}`, + session: "session", + username: "omp", + date: new Date().toISOString(), + msg_type: "execute_reply", + version: "5.5", + }, + parent_header: { msg_id: msgId }, + metadata: {}, + content: { status: "ok", execution_count: executionCount }, + }; + const status: JupyterMessage = { + channel: "iopub", + header: { + msg_id: `status-${msgId}`, + session: "session", + username: "omp", + date: new Date().toISOString(), + msg_type: "status", + version: "5.5", + }, + parent_header: { msg_id: msgId }, + metadata: {}, + content: { execution_state: "idle" }, + }; + ws.onmessage?.({ data: encodeMessage(reply) }); + ws.onmessage?.({ data: encodeMessage(status) }); +} + class FakeWebSocket { static OPEN = 1; static CLOSED = 3; @@ -149,106 +182,7 @@ describe("PythonKernel (external gateway)", () => { }); globalThis.fetch = fetchMock as unknown as typeof fetch; - const responseQueue: Array<(msgId: string, ws: FakeWebSocket) => void> = []; - responseQueue.push((msgId, ws) => { - const reply: JupyterMessage = { - channel: "shell", - header: { - msg_id: "reply-1", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "execute_reply", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { status: "ok", execution_count: 1 }, - }; - const status: JupyterMessage = { - channel: "iopub", - header: { - msg_id: "status-1", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "status", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { execution_state: "idle" }, - }; - ws.onmessage?.({ data: encodeMessage(reply) }); - ws.onmessage?.({ data: encodeMessage(status) }); - }); - responseQueue.push((msgId, ws) => { - const stream: JupyterMessage = { - channel: "iopub", - header: { - msg_id: "stream-1", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "stream", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { text: "hello\n" }, - }; - const display: JupyterMessage = { - channel: "iopub", - header: { - msg_id: "display-1", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "execute_result", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { - data: { - "text/plain": "result", - "application/json": { answer: 42 }, - }, - }, - }; - const reply: JupyterMessage = { - channel: "shell", - header: { - msg_id: "reply-2", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "execute_reply", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { status: "ok", execution_count: 2 }, - }; - const status: JupyterMessage = { - channel: "iopub", - header: { - msg_id: "status-2", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "status", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { execution_state: "idle" }, - }; - ws.onmessage?.({ data: encodeMessage(stream) }); - ws.onmessage?.({ data: encodeMessage(display) }); - ws.onmessage?.({ data: encodeMessage(reply) }); - ws.onmessage?.({ data: encodeMessage(status) }); - }); + let preludeSeen = false; const kernelPromise = PythonKernel.start({ cwd: "/" }); await Bun.sleep(10); @@ -256,11 +190,84 @@ describe("PythonKernel (external gateway)", () => { if (!ws) throw new Error("WebSocket not initialized"); ws.setSendHandler((data) => { const msg = typeof data === "string" ? (JSON.parse(data) as JupyterMessage) : decodeMessage(data); - const handler = responseQueue.shift(); - if (!handler) { - throw new Error(`Unexpected message: ${msg.header.msg_type}`); + const code = String(msg.content.code ?? ""); + if (!preludeSeen) { + expect(code).toBe(PYTHON_PRELUDE); + preludeSeen = true; + sendOkExecution(ws, msg.header.msg_id); + return; } - handler(msg.header.msg_id, ws); + + if (code === "print('hello')") { + const stream: JupyterMessage = { + channel: "iopub", + header: { + msg_id: "stream-1", + session: "session", + username: "omp", + date: new Date().toISOString(), + msg_type: "stream", + version: "5.5", + }, + parent_header: { msg_id: msg.header.msg_id }, + metadata: {}, + content: { text: "hello\n" }, + }; + const display: JupyterMessage = { + channel: "iopub", + header: { + msg_id: "display-1", + session: "session", + username: "omp", + date: new Date().toISOString(), + msg_type: "execute_result", + version: "5.5", + }, + parent_header: { msg_id: msg.header.msg_id }, + metadata: {}, + content: { + data: { + "text/plain": "result", + "application/json": { answer: 42 }, + }, + }, + }; + const reply: JupyterMessage = { + channel: "shell", + header: { + msg_id: "reply-2", + session: "session", + username: "omp", + date: new Date().toISOString(), + msg_type: "execute_reply", + version: "5.5", + }, + parent_header: { msg_id: msg.header.msg_id }, + metadata: {}, + content: { status: "ok", execution_count: 2 }, + }; + const status: JupyterMessage = { + channel: "iopub", + header: { + msg_id: "status-2", + session: "session", + username: "omp", + date: new Date().toISOString(), + msg_type: "status", + version: "5.5", + }, + parent_header: { msg_id: msg.header.msg_id }, + metadata: {}, + content: { execution_state: "idle" }, + }; + ws.onmessage?.({ data: encodeMessage(stream) }); + ws.onmessage?.({ data: encodeMessage(display) }); + ws.onmessage?.({ data: encodeMessage(reply) }); + ws.onmessage?.({ data: encodeMessage(status) }); + return; + } + + sendOkExecution(ws, msg.header.msg_id); }); const kernel = await kernelPromise; @@ -300,40 +307,7 @@ describe("PythonKernel (external gateway)", () => { }); globalThis.fetch = fetchMock as unknown as typeof fetch; - const responseQueue: Array<(msgId: string, ws: FakeWebSocket) => void> = [ - (msgId, ws) => { - const reply: JupyterMessage = { - channel: "shell", - header: { - msg_id: "reply-prelude", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "execute_reply", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { status: "ok", execution_count: 1 }, - }; - const status: JupyterMessage = { - channel: "iopub", - header: { - msg_id: "status-prelude", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "status", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { execution_state: "idle" }, - }; - ws.onmessage?.({ data: encodeMessage(reply) }); - ws.onmessage?.({ data: encodeMessage(status) }); - }, - ]; + let preludeSeen = false; const kernelPromise = PythonKernel.start({ cwd: "/" }); await Bun.sleep(10); @@ -341,11 +315,12 @@ describe("PythonKernel (external gateway)", () => { if (!ws) throw new Error("WebSocket not initialized"); ws.setSendHandler((data) => { const msg = typeof data === "string" ? (JSON.parse(data) as JupyterMessage) : decodeMessage(data); - const handler = responseQueue.shift(); - if (!handler) { - throw new Error(`Unexpected message: ${msg.header.msg_type}`); + const code = String(msg.content.code ?? ""); + if (!preludeSeen) { + expect(code).toBe(PYTHON_PRELUDE); + preludeSeen = true; } - handler(msg.header.msg_id, ws); + sendOkExecution(ws, msg.header.msg_id); }); const kernel = await kernelPromise; @@ -368,40 +343,7 @@ describe("PythonKernel (external gateway)", () => { }); globalThis.fetch = fetchMock as unknown as typeof fetch; - const responseQueue: Array<(msgId: string, ws: FakeWebSocket) => void> = [ - (msgId, ws) => { - const reply: JupyterMessage = { - channel: "shell", - header: { - msg_id: "reply-prelude", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "execute_reply", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { status: "ok", execution_count: 1 }, - }; - const status: JupyterMessage = { - channel: "iopub", - header: { - msg_id: "status-prelude", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "status", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { execution_state: "idle" }, - }; - ws.onmessage?.({ data: encodeMessage(reply) }); - ws.onmessage?.({ data: encodeMessage(status) }); - }, - ]; + let preludeSeen = false; const kernelPromise = PythonKernel.start({ cwd: "/" }); await Bun.sleep(10); @@ -409,12 +351,12 @@ describe("PythonKernel (external gateway)", () => { if (!ws) throw new Error("WebSocket not initialized"); ws.setSendHandler((data) => { const msg = typeof data === "string" ? (JSON.parse(data) as JupyterMessage) : decodeMessage(data); - const handler = responseQueue.shift(); - if (!handler) { - throw new Error(`Unexpected message: ${msg.header.msg_type}`); + const code = String(msg.content.code ?? ""); + if (!preludeSeen) { + expect(code).toBe(PYTHON_PRELUDE); + preludeSeen = true; } - expect(msg.content.code).toBe(PYTHON_PRELUDE); - handler(msg.header.msg_id, ws); + sendOkExecution(ws, msg.header.msg_id); }); const kernel = await kernelPromise; @@ -441,40 +383,23 @@ describe("PythonKernel (external gateway)", () => { ]; const payload = JSON.stringify(docs); - const responseQueue: Array<(msgId: string, ws: FakeWebSocket) => void> = [ - (msgId, ws) => { - const reply: JupyterMessage = { - channel: "shell", - header: { - msg_id: "reply-prelude", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "execute_reply", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { status: "ok", execution_count: 1 }, - }; - const status: JupyterMessage = { - channel: "iopub", - header: { - msg_id: "status-prelude", - session: "session", - username: "omp", - date: new Date().toISOString(), - msg_type: "status", - version: "5.5", - }, - parent_header: { msg_id: msgId }, - metadata: {}, - content: { execution_state: "idle" }, - }; - ws.onmessage?.({ data: encodeMessage(reply) }); - ws.onmessage?.({ data: encodeMessage(status) }); - }, - (msgId, ws) => { + let preludeSeen = false; + + const kernelPromise = PythonKernel.start({ cwd: "/" }); + await Bun.sleep(10); + const ws = FakeWebSocket.lastInstance; + if (!ws) throw new Error("WebSocket not initialized"); + ws.setSendHandler((data) => { + const msg = typeof data === "string" ? (JSON.parse(data) as JupyterMessage) : decodeMessage(data); + const code = String(msg.content.code ?? ""); + if (!preludeSeen) { + expect(code).toBe(PYTHON_PRELUDE); + preludeSeen = true; + sendOkExecution(ws, msg.header.msg_id); + return; + } + + if (code.includes("__omp_prelude_docs__")) { const stream: JupyterMessage = { channel: "iopub", header: { @@ -485,7 +410,7 @@ describe("PythonKernel (external gateway)", () => { msg_type: "stream", version: "5.5", }, - parent_header: { msg_id: msgId }, + parent_header: { msg_id: msg.header.msg_id }, metadata: {}, content: { text: `${payload}\n` }, }; @@ -499,7 +424,7 @@ describe("PythonKernel (external gateway)", () => { msg_type: "execute_reply", version: "5.5", }, - parent_header: { msg_id: msgId }, + parent_header: { msg_id: msg.header.msg_id }, metadata: {}, content: { status: "ok", execution_count: 2 }, }; @@ -513,30 +438,17 @@ describe("PythonKernel (external gateway)", () => { msg_type: "status", version: "5.5", }, - parent_header: { msg_id: msgId }, + parent_header: { msg_id: msg.header.msg_id }, metadata: {}, content: { execution_state: "idle" }, }; ws.onmessage?.({ data: encodeMessage(stream) }); ws.onmessage?.({ data: encodeMessage(reply) }); ws.onmessage?.({ data: encodeMessage(status) }); - }, - ]; + return; + } - const kernelPromise = PythonKernel.start({ cwd: "/" }); - await Bun.sleep(10); - const ws = FakeWebSocket.lastInstance; - if (!ws) throw new Error("WebSocket not initialized"); - ws.setSendHandler((data) => { - const msg = typeof data === "string" ? (JSON.parse(data) as JupyterMessage) : decodeMessage(data); - const handler = responseQueue.shift(); - if (!handler) { - throw new Error(`Unexpected message: ${msg.header.msg_type}`); - } - if (msg.content.code !== PYTHON_PRELUDE) { - expect(String(msg.content.code)).toContain("__omp_prelude_docs__"); - } - handler(msg.header.msg_id, ws); + sendOkExecution(ws, msg.header.msg_id); }); const kernel = await kernelPromise; diff --git a/packages/coding-agent/src/core/python-kernel.ts b/packages/coding-agent/src/core/python-kernel.ts index b4b160f55..db32d387e 100644 --- a/packages/coding-agent/src/core/python-kernel.ts +++ b/packages/coding-agent/src/core/python-kernel.ts @@ -6,6 +6,7 @@ import { getShellConfig, killProcessTree } from "../utils/shell"; import { getOrCreateSnapshot } from "../utils/shell-snapshot"; import { logger } from "./logger"; import { acquireSharedGateway, releaseSharedGateway } from "./python-gateway-coordinator"; +import { loadPythonModules } from "./python-modules"; import { PYTHON_PRELUDE } from "./python-prelude"; import { htmlToBasicMarkdown } from "./tools/web-scrapers/types"; import { ScopeSignal } from "./utils"; @@ -515,7 +516,7 @@ export class PythonKernel { const externalConfig = getExternalGatewayConfig(); if (externalConfig) { - return PythonKernel.startWithExternalGateway(externalConfig); + return PythonKernel.startWithExternalGateway(externalConfig, options.cwd); } // Try shared gateway first (unless explicitly disabled) @@ -535,7 +536,7 @@ export class PythonKernel { return PythonKernel.startWithLocalGateway(options); } - private static async startWithExternalGateway(config: ExternalGatewayConfig): Promise { + private static async startWithExternalGateway(config: ExternalGatewayConfig, cwd: string): Promise { const headers: Record = { "Content-Type": "application/json" }; if (config.token) { headers.Authorization = `token ${config.token}`; @@ -563,6 +564,7 @@ export class PythonKernel { if (preludeResult.cancelled || preludeResult.status === "error") { throw new Error("Failed to initialize Python kernel prelude"); } + await loadPythonModules(kernel, { cwd }); return kernel; } catch (err: unknown) { await kernel.shutdown(); @@ -570,7 +572,7 @@ export class PythonKernel { } } - private static async startWithSharedGateway(gatewayUrl: string, _cwd: string): Promise { + private static async startWithSharedGateway(gatewayUrl: string, cwd: string): Promise { const createResponse = await fetch(`${gatewayUrl}/api/kernels`, { method: "POST", headers: { "Content-Type": "application/json" }, @@ -594,6 +596,7 @@ export class PythonKernel { if (preludeResult.cancelled || preludeResult.status === "error") { throw new Error("Failed to initialize Python kernel prelude"); } + await loadPythonModules(kernel, { cwd }); return kernel; } catch (err: unknown) { await kernel.shutdown(); @@ -709,6 +712,7 @@ export class PythonKernel { if (preludeResult.cancelled || preludeResult.status === "error") { throw new Error("Failed to initialize Python kernel prelude"); } + await loadPythonModules(kernel, { cwd: options.cwd }); return kernel; } catch (err: unknown) { await kernel.shutdown(); diff --git a/packages/coding-agent/src/core/python-modules.test.ts b/packages/coding-agent/src/core/python-modules.test.ts new file mode 100644 index 000000000..c157e1598 --- /dev/null +++ b/packages/coding-agent/src/core/python-modules.test.ts @@ -0,0 +1,102 @@ +import { afterEach, describe, expect, it } from "bun:test"; +import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { basename, join, resolve } from "node:path"; +import { discoverPythonModules, loadPythonModules, type PythonModuleExecutor } from "./python-modules"; + +const fixturesDir = resolve(__dirname, "../../test/fixtures/python-modules"); + +const readFixture = (name: string): string => readFileSync(join(fixturesDir, name), "utf-8"); + +const writeModule = (dir: string, name: string, tag: string) => { + mkdirSync(dir, { recursive: true }); + const base = readFixture(name); + writeFileSync(join(dir, name), `${base}\n# ${tag}`); +}; + +const createTempRoot = () => mkdtempSync(join(tmpdir(), "omp-python-modules-")); + +describe("python modules", () => { + let tempRoot: string | null = null; + + afterEach(() => { + if (tempRoot) { + rmSync(tempRoot, { recursive: true, force: true }); + } + tempRoot = null; + }); + + it("discovers modules with project override and sorted order", async () => { + tempRoot = createTempRoot(); + const homeDir = join(tempRoot, "home"); + const cwd = join(tempRoot, "project"); + + writeModule(join(homeDir, ".omp", "agent", "modules"), "alpha.py", "user-omp"); + writeModule(join(homeDir, ".pi", "agent", "modules"), "beta.py", "user-pi"); + writeModule(join(homeDir, ".pi", "agent", "modules"), "delta.py", "user-pi"); + + writeModule(join(cwd, ".omp", "modules"), "alpha.py", "project-omp"); + writeModule(join(cwd, ".omp", "modules"), "beta.py", "project-omp"); + writeModule(join(cwd, ".pi", "modules"), "gamma.py", "project-pi"); + + const modules = await discoverPythonModules({ cwd, homeDir }); + const names = modules.map((module) => basename(module.path)); + expect(names).toEqual(["alpha.py", "beta.py", "delta.py", "gamma.py"]); + expect(modules.map((module) => ({ name: basename(module.path), source: module.source }))).toEqual([ + { name: "alpha.py", source: "project" }, + { name: "beta.py", source: "project" }, + { name: "delta.py", source: "user" }, + { name: "gamma.py", source: "project" }, + ]); + expect(modules.find((module) => module.path.endsWith("alpha.py"))?.content).toContain("project-omp"); + expect(modules.find((module) => module.path.endsWith("delta.py"))?.content).toContain("user-pi"); + }); + + it("loads modules in sorted order with silent execution", async () => { + tempRoot = createTempRoot(); + const homeDir = join(tempRoot, "home"); + const cwd = join(tempRoot, "project"); + + writeModule(join(homeDir, ".omp", "agent", "modules"), "beta.py", "user-omp"); + writeModule(join(homeDir, ".omp", "agent", "modules"), "alpha.py", "user-omp"); + + const calls: Array<{ name: string; options?: { silent?: boolean; storeHistory?: boolean } }> = []; + const executor: PythonModuleExecutor = { + execute: async (code: string, options?: { silent?: boolean; storeHistory?: boolean }) => { + const name = code.includes("def alpha") ? "alpha" : "beta"; + calls.push({ name, options }); + return { status: "ok", cancelled: false }; + }, + }; + + await loadPythonModules(executor, { cwd, homeDir }); + expect(calls.map((call) => call.name)).toEqual(["alpha", "beta"]); + for (const call of calls) { + expect(call.options).toEqual({ silent: true, storeHistory: false }); + } + }); + + it("fails fast when a module fails to execute", async () => { + tempRoot = createTempRoot(); + const homeDir = join(tempRoot, "home"); + const cwd = join(tempRoot, "project"); + + writeModule(join(homeDir, ".omp", "agent", "modules"), "alpha.py", "user-omp"); + writeModule(join(cwd, ".omp", "modules"), "beta.py", "project-omp"); + + const executor: PythonModuleExecutor = { + execute: async (code: string) => { + if (code.includes("def beta")) { + return { + status: "error", + cancelled: false, + error: { name: "Error", value: "boom", traceback: [] }, + }; + } + return { status: "ok", cancelled: false }; + }, + }; + + await expect(loadPythonModules(executor, { cwd, homeDir })).rejects.toThrow("Failed to load Python module"); + }); +}); diff --git a/packages/coding-agent/src/core/python-modules.ts b/packages/coding-agent/src/core/python-modules.ts new file mode 100644 index 000000000..bd9b61a60 --- /dev/null +++ b/packages/coding-agent/src/core/python-modules.ts @@ -0,0 +1,110 @@ +import { readdir } from "node:fs/promises"; +import { homedir } from "node:os"; +import { join, resolve } from "node:path"; + +export type PythonModuleSource = "user" | "project"; + +export interface PythonModuleEntry { + path: string; + content: string; + source: PythonModuleSource; +} + +export interface PythonModuleExecuteResult { + status: "ok" | "error"; + cancelled: boolean; + error?: { name: string; value: string; traceback: string[] }; +} + +export interface PythonModuleExecutor { + execute: ( + code: string, + options?: { silent?: boolean; storeHistory?: boolean }, + ) => Promise; +} + +export interface DiscoverPythonModulesOptions { + /** Working directory for project-level modules. Default: process.cwd() */ + cwd?: string; + /** Home directory for user-level modules. Default: os.homedir() */ + homeDir?: string; +} + +interface ModuleCandidate { + name: string; + path: string; + source: PythonModuleSource; +} + +async function listModuleCandidates(dir: string, source: PythonModuleSource): Promise { + try { + const entries = await readdir(dir, { withFileTypes: true }); + return entries + .filter((entry) => entry.isFile() && entry.name.endsWith(".py")) + .map((entry) => ({ + name: entry.name, + path: resolve(dir, entry.name), + source, + })); + } catch { + return []; + } +} + +async function readModuleContent(candidate: ModuleCandidate): Promise { + try { + const content = await Bun.file(candidate.path).text(); + return { path: candidate.path, content, source: candidate.source }; + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + throw new Error(`Failed to read Python module ${candidate.path}: ${message}`); + } +} + +/** + * Discover Python prelude extension modules from user and project directories. + */ +export async function discoverPythonModules(options: DiscoverPythonModulesOptions = {}): Promise { + const cwd = options.cwd ?? process.cwd(); + const homeDir = options.homeDir ?? homedir(); + + const userDirs = [join(homeDir, ".omp", "agent", "modules"), join(homeDir, ".pi", "agent", "modules")]; + const projectDirs = [resolve(cwd, ".omp", "modules"), resolve(cwd, ".pi", "modules")]; + + const userCandidates = (await Promise.all(userDirs.map((dir) => listModuleCandidates(dir, "user")))).flat(); + const projectCandidates = (await Promise.all(projectDirs.map((dir) => listModuleCandidates(dir, "project")))).flat(); + + const byName = new Map(); + for (const candidate of userCandidates) { + if (!byName.has(candidate.name)) { + byName.set(candidate.name, candidate); + } + } + for (const candidate of projectCandidates) { + const existing = byName.get(candidate.name); + if (!existing || existing.source === "user") { + byName.set(candidate.name, candidate); + } + } + + const sorted = Array.from(byName.values()).sort((a, b) => a.name.localeCompare(b.name)); + return Promise.all(sorted.map((candidate) => readModuleContent(candidate))); +} + +/** + * Load Python prelude extension modules into an active kernel. + */ +export async function loadPythonModules( + executor: PythonModuleExecutor, + options: DiscoverPythonModulesOptions = {}, +): Promise { + const modules = await discoverPythonModules(options); + for (const module of modules) { + const result = await executor.execute(module.content, { silent: true, storeHistory: false }); + if (result.cancelled || result.status === "error") { + const details = result.error ? `${result.error.name}: ${result.error.value}` : "unknown error"; + throw new Error(`Failed to load Python module ${module.path}: ${details}`); + } + } + return modules; +} diff --git a/packages/coding-agent/test/fixtures/python-modules/alpha.py b/packages/coding-agent/test/fixtures/python-modules/alpha.py new file mode 100644 index 000000000..88f238870 --- /dev/null +++ b/packages/coding-agent/test/fixtures/python-modules/alpha.py @@ -0,0 +1,2 @@ +def alpha(): + return "alpha" diff --git a/packages/coding-agent/test/fixtures/python-modules/beta.py b/packages/coding-agent/test/fixtures/python-modules/beta.py new file mode 100644 index 000000000..8571dc97f --- /dev/null +++ b/packages/coding-agent/test/fixtures/python-modules/beta.py @@ -0,0 +1,2 @@ +def beta(): + return "beta" diff --git a/packages/coding-agent/test/fixtures/python-modules/delta.py b/packages/coding-agent/test/fixtures/python-modules/delta.py new file mode 100644 index 000000000..a6d0265e2 --- /dev/null +++ b/packages/coding-agent/test/fixtures/python-modules/delta.py @@ -0,0 +1,2 @@ +def delta(): + return "delta" diff --git a/packages/coding-agent/test/fixtures/python-modules/gamma.py b/packages/coding-agent/test/fixtures/python-modules/gamma.py new file mode 100644 index 000000000..fd25b457a --- /dev/null +++ b/packages/coding-agent/test/fixtures/python-modules/gamma.py @@ -0,0 +1,2 @@ +def gamma(): + return "gamma"