import z from "zod"import { Tool } from "../tool"import { spawn, type ChildProcess } from "child_process"import path from "path"import os from "os"import { Instance } from "@/project/instance"import { OpenScience } from "@/openscience" const KERNEL_SCRIPT = `import sys, json, io, traceback, os _out = sys.stdout_err = sys.stderr ns = {"__name__": "__main__", "__builtins__": __builtins__} # Pre-import common scientific packagesfor pkg, alias in [("numpy", "np"), ("pandas", "pd"), ("scipy", "scipy")]: try: mod = __import__(pkg) ns[alias] = mod ns[pkg] = mod except ImportError: pass _out.write("__OPENSCIENCE_KERNEL_READY__\\n")_out.flush() while True: lines = [] try: for line in sys.stdin: if line.rstrip("\\n") == "__OPENSCIENCE_CODE_END__": break lines.append(line) except EOFError: break if not lines: continue code = "".join(lines) stdout_buf = io.StringIO() stderr_buf = io.StringIO() sys.stdout = stdout_buf sys.stderr = stderr_buf ok = True try: # Try eval first for expression auto-display (like Jupyter) try: compiled = compile(code, "<cell>", "eval") result = eval(compiled, ns) if result is not None: print(repr(result)) except SyntaxError: exec(compile(code, "<cell>", "exec"), ns) except SystemExit: stderr_buf.write("SystemExit caught (kernel stays alive)\\n") ok = False except Exception: traceback.print_exc() ok = False finally: sys.stdout = _out sys.stderr = _err r = json.dumps({"ok": ok, "stdout": stdout_buf.getvalue(), "stderr": stderr_buf.getvalue()}) _out.write("__OPENSCIENCE_RESULT_START__\\n" + r + "\\n__OPENSCIENCE_RESULT_END__\\n") _out.flush()`.trim() interface Kernel { process: ChildProcess scriptPath: string lastUsed: number} const kernels = new Map<string, Kernel>() // Clean up all kernels on process exitfunction cleanupAll() { for (const [id, kernel] of kernels) { try { kernel.process.kill() } catch {} try { require("fs").unlinkSync(kernel.scriptPath) } catch {} kernels.delete(id) }} process.on("exit", cleanupAll)process.on("SIGTERM", cleanupAll)process.on("SIGINT", cleanupAll) function cleanupIdle() { const now = Date.now() const idle = 30 * 60 * 1000 // 30 min for (const [id, kernel] of kernels) { if (now - kernel.lastUsed > idle) { try { kernel.process.kill() } catch {} try { require("fs").unlinkSync(kernel.scriptPath) } catch {} kernels.delete(id) } }} async function getKernel(sessionID: string): Promise<Kernel> { // Clean up idle kernels while we're here cleanupIdle() const existing = kernels.get(sessionID) if (existing && !existing.process.killed && existing.process.exitCode === null) { existing.lastUsed = Date.now() return existing } // Dead kernel — clean up if (existing) { try { require("fs").unlinkSync(existing.scriptPath) } catch {} kernels.delete(sessionID) } // Start new kernel const scriptPath = path.join(os.tmpdir(), `openscience-kernel-${sessionID.slice(0, 8)}-${Date.now()}.py`) await Bun.write(scriptPath, KERNEL_SCRIPT) const pythonBin = await findPython() const proc = spawn(pythonBin, ["-u", scriptPath], { cwd: Instance.directory, env: { ...(await OpenScience.subprocessEnv(process.env)), PYTHONUNBUFFERED: "1" }, stdio: ["pipe", "pipe", "pipe"], }) // Collect kernel stderr (startup warnings, etc.) let kernelStderr = "" proc.stderr?.on("data", (data: Buffer) => { kernelStderr += data.toString() // Cap stderr buffer if (kernelStderr.length > 10_000) kernelStderr = kernelStderr.slice(-5000) }) // Wait for ready signal await new Promise<void>((resolve, reject) => { const timeout = setTimeout(() => { proc.kill() reject(new Error(`Kernel startup timed out. stderr: ${kernelStderr}`)) }, 15_000) let buf = "" const handler = (data: Buffer) => { buf += data.toString() if (buf.includes("__OPENSCIENCE_KERNEL_READY__")) { clearTimeout(timeout) proc.stdout?.off("data", handler) resolve() } } proc.stdout?.on("data", handler) proc.once("error", (err) => { clearTimeout(timeout) reject(err) }) proc.once("exit", (code) => { clearTimeout(timeout) reject(new Error(`Kernel exited during startup (code ${code}). stderr: ${kernelStderr}`)) }) }) const kernel: Kernel = { process: proc, scriptPath, lastUsed: Date.now() } kernels.set(sessionID, kernel) return kernel} function executeInKernel( kernel: Kernel, code: string, timeout: number,): Promise<{ ok: boolean; stdout: string; stderr: string }> { return new Promise((resolve, reject) => { const timer = setTimeout(() => { // Kill the timed-out kernel — it will be restarted on next call try { kernel.process.kill() } catch {} reject(new Error(`Cell execution timed out after ${Math.round(timeout / 1000)}s`)) }, timeout) let buffer = "" const handler = (data: Buffer) => { buffer += data.toString() const startMarker = "__OPENSCIENCE_RESULT_START__\n" const endMarker = "\n__OPENSCIENCE_RESULT_END__" const startIdx = buffer.indexOf(startMarker) const endIdx = buffer.indexOf(endMarker) if (startIdx !== -1 && endIdx !== -1 && endIdx > startIdx) { clearTimeout(timer) kernel.process.stdout?.off("data", handler) const json = buffer.slice(startIdx + startMarker.length, endIdx) try { resolve(JSON.parse(json)) } catch { resolve({ ok: false, stdout: "", stderr: `Kernel response parse error: ${json.slice(0, 500)}` }) } } } const exitHandler = (code: number | null) => { clearTimeout(timer) kernel.process.stdout?.off("data", handler) reject(new Error(`Kernel died during execution (exit code ${code})`)) } kernel.process.stdout?.on("data", handler) kernel.process.once("exit", exitHandler) kernel.process.stdin?.write(code + "\n__OPENSCIENCE_CODE_END__\n") })} async function findPython(): Promise<string> { for (const bin of ["python3", "python"]) { try { const proc = Bun.spawn([bin, "--version"], { stdout: "pipe", stderr: "pipe" }) await proc.exited if (proc.exitCode === 0) return bin } catch {} } throw new Error("Python not found. Install Python 3.10+ to use the notebook tool.")} export const NotebookTool = Tool.define("notebook", { description: [ "Execute Python code in a persistent kernel. Variables, imports, and state persist across calls.", "Use instead of `bash python` for analysis — no need to re-import or re-load data between cells.", "numpy (np), pandas (pd), scipy are pre-imported. Expression results auto-display like Jupyter.", ].join("\n"), parameters: z.object({ code: z.string().describe("Python code to execute in the persistent kernel"), timeout: z.number().default(120_000).describe("Execution timeout in ms (default: 120s, max: 600s)"), }), async execute(params, ctx) { const timeout = Math.min(Math.max(params.timeout, 5_000), 600_000) // Same permission as bash — this executes arbitrary code await ctx.ask({ permission: "bash", patterns: ["python (notebook)"], always: ["python*"], metadata: {}, }) const kernel = await getKernel(ctx.sessionID) const result = await executeInKernel(kernel, params.code, timeout) // Stream metadata updates for the UI ctx.metadata({ metadata: { output: result.stdout || result.stderr || "(no output)", ok: result.ok, }, }) const parts: string[] = [] if (result.stdout) parts.push(result.stdout) if (result.stderr) { parts.push(result.ok ? `[stderr]\n${result.stderr}` : `[ERROR]\n${result.stderr}`) } if (!parts.length) parts.push("(no output)") const output = parts.join("\n") return { title: result.ok ? "Python cell" : "Python cell (error)", output, metadata: { ok: result.ok, output: output.length > 30_000 ? output.slice(0, 30_000) + "\n\n..." : output, }, } },})
import z from "zod"
import { Tool } from "../tool"
import { spawn, type ChildProcess } from "child_process"
import path from "path"
import os from "os"
import { Instance } from "@/project/instance"
import { OpenScience } from "@/openscience"
const KERNEL_SCRIPT = `
import sys, json, io, traceback, os
_out = sys.stdout
_err = sys.stderr
ns = {"__name__": "__main__", "__builtins__": __builtins__}
# Pre-import common scientific packages
for pkg, alias in [("numpy", "np"), ("pandas", "pd"), ("scipy", "scipy")]:
try:
mod = __import__(pkg)
ns[alias] = mod
ns[pkg] = mod
except ImportError:
pass
_out.write("__OPENSCIENCE_KERNEL_READY__\\n")
_out.flush()
while True:
lines = []
try:
for line in sys.stdin:
if line.rstrip("\\n") == "__OPENSCIENCE_CODE_END__":
break
lines.append(line)
except EOFError:
break
if not lines:
continue
code = "".join(lines)
stdout_buf = io.StringIO()
stderr_buf = io.StringIO()
sys.stdout = stdout_buf
sys.stderr = stderr_buf
ok = True
try:
# Try eval first for expression auto-display (like Jupyter)
try:
compiled = compile(code, "<cell>", "eval")
result = eval(compiled, ns)
if result is not None:
print(repr(result))
except SyntaxError:
exec(compile(code, "<cell>", "exec"), ns)
except SystemExit:
stderr_buf.write("SystemExit caught (kernel stays alive)\\n")
ok = False
except Exception:
traceback.print_exc()
ok = False
finally:
sys.stdout = _out
sys.stderr = _err
r = json.dumps({"ok": ok, "stdout": stdout_buf.getvalue(), "stderr": stderr_buf.getvalue()})
_out.write("__OPENSCIENCE_RESULT_START__\\n" + r + "\\n__OPENSCIENCE_RESULT_END__\\n")
_out.flush()
`.trim()
interface Kernel {
process: ChildProcess
scriptPath: string
lastUsed: number
}
const kernels = new Map<string, Kernel>()
// Clean up all kernels on process exit
function cleanupAll() {
for (const [id, kernel] of kernels) {
try {
kernel.process.kill()
} catch {}
try {
require("fs").unlinkSync(kernel.scriptPath)
} catch {}
kernels.delete(id)
}
}
process.on("exit", cleanupAll)
process.on("SIGTERM", cleanupAll)
process.on("SIGINT", cleanupAll)
function cleanupIdle() {
const now = Date.now()
const idle = 30 * 60 * 1000 // 30 min
for (const [id, kernel] of kernels) {
if (now - kernel.lastUsed > idle) {
try {
kernel.process.kill()
} catch {}
try {
require("fs").unlinkSync(kernel.scriptPath)
} catch {}
kernels.delete(id)
}
}
}
async function getKernel(sessionID: string): Promise<Kernel> {
// Clean up idle kernels while we're here
cleanupIdle()
const existing = kernels.get(sessionID)
if (existing && !existing.process.killed && existing.process.exitCode === null) {
existing.lastUsed = Date.now()
return existing
}
// Dead kernel — clean up
if (existing) {
try {
require("fs").unlinkSync(existing.scriptPath)
} catch {}
kernels.delete(sessionID)
}
// Start new kernel
const scriptPath = path.join(os.tmpdir(), `openscience-kernel-${sessionID.slice(0, 8)}-${Date.now()}.py`)
await Bun.write(scriptPath, KERNEL_SCRIPT)
const pythonBin = await findPython()
const proc = spawn(pythonBin, ["-u", scriptPath], {
cwd: Instance.directory,
env: { ...(await OpenScience.subprocessEnv(process.env)), PYTHONUNBUFFERED: "1" },
stdio: ["pipe", "pipe", "pipe"],
})
// Collect kernel stderr (startup warnings, etc.)
let kernelStderr = ""
proc.stderr?.on("data", (data: Buffer) => {
kernelStderr += data.toString()
// Cap stderr buffer
if (kernelStderr.length > 10_000) kernelStderr = kernelStderr.slice(-5000)
})
// Wait for ready signal
await new Promise<void>((resolve, reject) => {
const timeout = setTimeout(() => {
proc.kill()
reject(new Error(`Kernel startup timed out. stderr: ${kernelStderr}`))
}, 15_000)
let buf = ""
const handler = (data: Buffer) => {
buf += data.toString()
if (buf.includes("__OPENSCIENCE_KERNEL_READY__")) {
clearTimeout(timeout)
proc.stdout?.off("data", handler)
resolve()
}
}
proc.stdout?.on("data", handler)
proc.once("error", (err) => {
clearTimeout(timeout)
reject(err)
})
proc.once("exit", (code) => {
clearTimeout(timeout)
reject(new Error(`Kernel exited during startup (code ${code}). stderr: ${kernelStderr}`))
})
})
const kernel: Kernel = { process: proc, scriptPath, lastUsed: Date.now() }
kernels.set(sessionID, kernel)
return kernel
}
function executeInKernel(
kernel: Kernel,
code: string,
timeout: number,
): Promise<{ ok: boolean; stdout: string; stderr: string }> {
return new Promise((resolve, reject) => {
const timer = setTimeout(() => {
// Kill the timed-out kernel — it will be restarted on next call
try {
kernel.process.kill()
} catch {}
reject(new Error(`Cell execution timed out after ${Math.round(timeout / 1000)}s`))
}, timeout)
let buffer = ""
const handler = (data: Buffer) => {
buffer += data.toString()
const startMarker = "__OPENSCIENCE_RESULT_START__\n"
const endMarker = "\n__OPENSCIENCE_RESULT_END__"
const startIdx = buffer.indexOf(startMarker)
const endIdx = buffer.indexOf(endMarker)
if (startIdx !== -1 && endIdx !== -1 && endIdx > startIdx) {
clearTimeout(timer)
kernel.process.stdout?.off("data", handler)
const json = buffer.slice(startIdx + startMarker.length, endIdx)
try {
resolve(JSON.parse(json))
} catch {
resolve({ ok: false, stdout: "", stderr: `Kernel response parse error: ${json.slice(0, 500)}` })
}
}
}
const exitHandler = (code: number | null) => {
clearTimeout(timer)
kernel.process.stdout?.off("data", handler)
reject(new Error(`Kernel died during execution (exit code ${code})`))
}
kernel.process.stdout?.on("data", handler)
kernel.process.once("exit", exitHandler)
kernel.process.stdin?.write(code + "\n__OPENSCIENCE_CODE_END__\n")
})
}
async function findPython(): Promise<string> {
for (const bin of ["python3", "python"]) {
try {
const proc = Bun.spawn([bin, "--version"], { stdout: "pipe", stderr: "pipe" })
await proc.exited
if (proc.exitCode === 0) return bin
} catch {}
}
throw new Error("Python not found. Install Python 3.10+ to use the notebook tool.")
}
export const NotebookTool = Tool.define("notebook", {
description: [
"Execute Python code in a persistent kernel. Variables, imports, and state persist across calls.",
"Use instead of `bash python` for analysis — no need to re-import or re-load data between cells.",
"numpy (np), pandas (pd), scipy are pre-imported. Expression results auto-display like Jupyter.",
].join("\n"),
parameters: z.object({
code: z.string().describe("Python code to execute in the persistent kernel"),
timeout: z.number().default(120_000).describe("Execution timeout in ms (default: 120s, max: 600s)"),
}),
async execute(params, ctx) {
const timeout = Math.min(Math.max(params.timeout, 5_000), 600_000)
// Same permission as bash — this executes arbitrary code
await ctx.ask({
permission: "bash",
patterns: ["python (notebook)"],
always: ["python*"],
metadata: {},
})
const kernel = await getKernel(ctx.sessionID)
const result = await executeInKernel(kernel, params.code, timeout)
// Stream metadata updates for the UI
ctx.metadata({
metadata: {
output: result.stdout || result.stderr || "(no output)",
ok: result.ok,
},
})
const parts: string[] = []
if (result.stdout) parts.push(result.stdout)
if (result.stderr) {
parts.push(result.ok ? `[stderr]\n${result.stderr}` : `[ERROR]\n${result.stderr}`)
}
if (!parts.length) parts.push("(no output)")
const output = parts.join("\n")
return {
title: result.ok ? "Python cell" : "Python cell (error)",
output,
metadata: {
ok: result.ok,
output: output.length > 30_000 ? output.slice(0, 30_000) + "\n\n..." : output,
},
}
},
})