Segfaulting code

Author: intentionally-left-nil

.gitignore

@@ -1,4 +1,6 @@
 .env
+.env.*
+.env.production
 runt
 echo-agent
 
@@ -18,6 +20,5 @@ __pycache__/
 *.pyc
 *.pyo
 
-.env
-.env.*
-.env.production
+# Javascript
+node_modules/

packages/python-runtime-agent/deno.json

@@ -15,7 +15,9 @@
   },
   "imports": {
     "@runt/lib": "jsr:@runt/lib@^0.6.2",
-    "@runt/schema": "jsr:@runt/schema@^0.6.2"
+    "@runt/schema": "jsr:@runt/schema@^0.6.2",
+    "@runtimed/jmp": "npm:@runtimed/jmp@^3.0.0",
+    "@std/path": "jsr:@std/path@^1.1.1"
   },
   "tasks": {
     "test": "deno test --allow-all",

packages/python-runtime-agent/src/python-worker.ts

@@ -1,5 +1,7 @@
 import { createLogger } from "@runt/lib";
 import * as path from "@std/path";
+import * as fs from "jsr:@std/fs";
+import jmp from "@runtimed/jmp";
 
 export class PythonWorker {
   private kernel: Deno.ChildProcess | null = null;
@@ -68,4 +70,111 @@ export class PythonWorker {
   getConnPath(): string | null {
     return this.connPath;
   }
+
+  /**
+   * Returns a JupyterKernelConnection instance after the kernel is started.
+   */
+  async getConnection(): Promise<JupyterKernelConnection> {
+    if (!this.connPath) throw new Error("Kernel not started");
+    return await JupyterKernelConnection.connect(this.connPath);
+  }
+}
+
+/**
+ * Connects to a Jupyter kernel using conn.json and @runtimed/jmp.
+ */
+export class JupyterKernelConnection {
+  private shellSocket: any;
+  private iopubSocket: any;
+  private controlSocket: any;
+  private stdinSocket: any;
+  private hbSocket: any;
+  private config: any;
+  private session: string;
+  private key: string;
+  private scheme: string;
+  private username: string;
+
+  private constructor(config: any, sockets: Record<string, any>) {
+    this.config = config;
+    this.shellSocket = sockets.shell;
+    this.iopubSocket = sockets.iopub;
+    this.controlSocket = sockets.control;
+    this.stdinSocket = sockets.stdin;
+    this.hbSocket = sockets.hb;
+    this.session = crypto.randomUUID();
+    this.key = config.key;
+    this.scheme = config.signature_scheme.replace("hmac-", "");
+    this.username = "runt";
+  }
+
+  static async connect(connPath: string): Promise<JupyterKernelConnection> {
+    const configRaw = await Deno.readTextFile(connPath);
+    const config = JSON.parse(configRaw);
+    const mkAddr = (port: number) => `${config.transport}://${config.ip}:${port}`;
+    const scheme = config.signature_scheme.replace("hmac-", "");
+    const key = config.key;
+    const shell = new jmp.Socket("dealer", scheme, key);
+    const iopub = new jmp.Socket("sub", scheme, key);
+    const control = new jmp.Socket("dealer", scheme, key);
+    const stdin = new jmp.Socket("dealer", scheme, key);
+    const hb = new jmp.Socket("req", scheme, key);
+    shell.connect(mkAddr(config.shell_port));
+    iopub.connect(mkAddr(config.iopub_port));
+    iopub.subscribe("");
+    control.connect(mkAddr(config.control_port));
+    stdin.connect(mkAddr(config.stdin_port));
+    hb.connect(mkAddr(config.hb_port));
+    return new JupyterKernelConnection(config, { shell, iopub, control, stdin, hb });
+  }
+
+  /**
+   * Execute Python code and return the result as a Promise.
+   */
+  async execute(code: string): Promise<{ result: string; outputs: unknown[] }> {
+    const msg_id = crypto.randomUUID();
+    const header = {
+      msg_id,
+      username: this.username,
+      session: this.session,
+      msg_type: "execute_request",
+      version: "5.3",
+    };
+    const content = {
+      code,
+      silent: false,
+      store_history: true,
+      user_expressions: {},
+      allow_stdin: false,
+      stop_on_error: true,
+    };
+    const msg = new jmp.Message();
+    msg.header = header;
+    msg.parent_header = {};
+    msg.metadata = {};
+    msg.content = content;
+    msg.idents = [];
+    // Send execute_request
+    this.shellSocket.send(msg);
+    // Listen for iopub messages for this msg_id
+    const outputs: unknown[] = [];
+    let result: string | undefined;
+    for await (const msg of this.iopubSocket) {
+      if (msg.parent_header?.msg_id !== msg_id) continue;
+      if (msg.header.msg_type === "execute_result" || msg.header.msg_type === "display_data") {
+        outputs.push(msg.content);
+        if (msg.content.data && msg.content.data["text/plain"]) {
+          result = msg.content.data["text/plain"];
+        }
+      } else if (msg.header.msg_type === "stream") {
+        outputs.push(msg.content);
+      } else if (msg.header.msg_type === "error") {
+        outputs.push(msg.content);
+        result = msg.content.ename + ": " + msg.content.evalue;
+      } else if (msg.header.msg_type === "status" && msg.content.execution_state === "idle") {
+        break;
+      }
+    }
+    return { result: result ?? "", outputs };
+  }
 }

packages/python-runtime-agent/test/python-worker-integration.test.ts

@@ -84,6 +84,16 @@ Deno.test({
       assertEquals(worker.getKernelPid(), null, "Kernel PID should be null after shutdown");
     });
 
+    await t.step("Can execute Python code and get result via JupyterKernelConnection", async () => {
+      worker = new PythonWorker(pythonPath);
+      await worker.start();
+      const conn = await worker.getConnection();
+      const { result, outputs } = await conn.execute("1+2");
+      assertEquals(result.trim(), "3");
+      assert(outputs.length > 0);
+      await worker.shutdown();
+    });
+
     // Teardown
     await t.step("[teardown] remove venv tempdir", async () => {
       if (venvDir) {