@oh-my-pi/pi-coding-agent 6.9.0 → 7.0.0

package/src/core/bash-executor.ts

@@ -56,7 +56,7 @@ export async function executeBash(command: string, options?: BashExecutorOptions
 			cancelled: false,
 			...(await sink.dump()),
 		};
-	} catch (err) {
+	} catch (err: unknown) {
 		// Exception covers NonZeroExitError, AbortError, TimeoutError
 		if (err instanceof Exception) {
 			if (err.aborted) {

package/src/core/extensions/index.ts

@@ -98,6 +98,8 @@ export type {
 	TurnStartEvent,
 	UserBashEvent,
 	UserBashEventResult,
+	UserPythonEvent,
+	UserPythonEventResult,
 	WriteToolResultEvent,
 } from "./types";
 // Type guards

package/src/core/extensions/runner.ts

@@ -40,6 +40,8 @@ import type {
 	ToolResultEventResult,
 	UserBashEvent,
 	UserBashEventResult,
+	UserPythonEvent,
+	UserPythonEventResult,
 } from "./types";
 
 /** Combined result from all before_agent_start handlers */
@@ -461,6 +463,35 @@ export class ExtensionRunner {
 		return undefined;
 	}
 
+	async emitUserPython(event: UserPythonEvent): Promise<UserPythonEventResult | undefined> {
+		const ctx = this.createContext();
+
+		for (const ext of this.extensions) {
+			const handlers = ext.handlers.get("user_python");
+			if (!handlers || handlers.length === 0) continue;
+
+			for (const handler of handlers) {
+				try {
+					const handlerResult = await handler(event, ctx);
+					if (handlerResult) {
+						return handlerResult as UserPythonEventResult;
+					}
+				} catch (err) {
+					const message = err instanceof Error ? err.message : String(err);
+					const stack = err instanceof Error ? err.stack : undefined;
+					this.emitError({
+						extensionPath: ext.path,
+						event: "user_python",
+						error: message,
+						stack,
+					});
+				}
+			}
+		}
+
+		return undefined;
+	}
+
 	/** Emit input event. Transforms chain, "handled" short-circuits. */
 	async emitInput(
 		text: string,

package/src/core/extensions/types.ts

@@ -21,6 +21,7 @@ import type { ExecOptions, ExecResult } from "../exec";
 import type { KeybindingsManager } from "../keybindings";
 import type { CustomMessage } from "../messages";
 import type { ModelRegistry } from "../model-registry";
+import type { PythonResult } from "../python-executor";
 import type {
 	BranchSummaryEntry,
 	CompactionEntry,
@@ -405,6 +406,21 @@ export interface UserBashEvent {
 	cwd: string;
 }
 
+// ============================================================================
+// User Python Events
+// ============================================================================
+
+/** Fired when user executes Python code via $ or $$ prefix */
+export interface UserPythonEvent {
+	type: "user_python";
+	/** The Python code to execute */
+	code: string;
+	/** True if $$ prefix was used (excluded from LLM context) */
+	excludeFromContext: boolean;
+	/** Current working directory */
+	cwd: string;
+}
+
 // ============================================================================
 // Input Events
 // ============================================================================
@@ -521,6 +537,7 @@ export type ExtensionEvent =
 	| TurnStartEvent
 	| TurnEndEvent
 	| UserBashEvent
+	| UserPythonEvent
 	| InputEvent
 	| ToolCallEvent
 	| ToolResultEvent;
@@ -554,6 +571,12 @@ export interface UserBashEventResult {
 	result?: BashResult;
 }
 
+/** Result from user_python event handler */
+export interface UserPythonEventResult {
+	/** Full replacement: extension handled execution, use this result */
+	result?: PythonResult;
+}
+
 export interface ToolResultEventResult {
 	content?: (TextContent | ImageContent)[];
 	details?: unknown;
@@ -671,6 +694,7 @@ export interface ExtensionAPI {
 	on(event: "tool_call", handler: ExtensionHandler<ToolCallEvent, ToolCallEventResult>): void;
 	on(event: "tool_result", handler: ExtensionHandler<ToolResultEvent, ToolResultEventResult>): void;
 	on(event: "user_bash", handler: ExtensionHandler<UserBashEvent, UserBashEventResult>): void;
+	on(event: "user_python", handler: ExtensionHandler<UserPythonEvent, UserPythonEventResult>): void;
 
 	// =========================================================================
 	// Tool Registration

package/src/core/messages.ts

@@ -39,6 +39,23 @@ export interface BashExecutionMessage {
 	excludeFromContext?: boolean;
 }
 
+/**
+ * Message type for user-initiated Python executions via the $ command.
+ * Shares the same kernel session as the agent's Python tool.
+ */
+export interface PythonExecutionMessage {
+	role: "pythonExecution";
+	code: string;
+	output: string;
+	exitCode: number | undefined;
+	cancelled: boolean;
+	truncated: boolean;
+	fullOutputPath?: string;
+	timestamp: number;
+	/** If true, this message is excluded from LLM context ($$ prefix) */
+	excludeFromContext?: boolean;
+}
+
 /**
  * Message type for extension-injected messages via sendMessage().
  */
@@ -95,6 +112,7 @@ export interface FileMentionMessage {
 declare module "@oh-my-pi/pi-agent-core" {
 	interface CustomAgentMessages {
 		bashExecution: BashExecutionMessage;
+		pythonExecution: PythonExecutionMessage;
 		custom: CustomMessage;
 		hookMessage: HookMessage;
 		branchSummary: BranchSummaryMessage;
@@ -124,6 +142,27 @@ export function bashExecutionToText(msg: BashExecutionMessage): string {
 	return text;
 }
 
+/**
+ * Convert a PythonExecutionMessage to user message text for LLM context.
+ */
+export function pythonExecutionToText(msg: PythonExecutionMessage): string {
+	let text = `Ran Python:\n\`\`\`python\n${msg.code}\n\`\`\`\n`;
+	if (msg.output) {
+		text += `Output:\n\`\`\`\n${msg.output}\n\`\`\``;
+	} else {
+		text += "(no output)";
+	}
+	if (msg.cancelled) {
+		text += "\n\n(execution cancelled)";
+	} else if (msg.exitCode !== null && msg.exitCode !== undefined && msg.exitCode !== 0) {
+		text += `\n\nExecution failed with code ${msg.exitCode}`;
+	}
+	if (msg.truncated && msg.fullOutputPath) {
+		text += `\n\n[Output truncated. Full output: ${msg.fullOutputPath}]`;
+	}
+	return text;
+}
+
 export function createBranchSummaryMessage(summary: string, fromId: string, timestamp: string): BranchSummaryMessage {
 	return {
 		role: "branchSummary",
@@ -185,6 +224,15 @@ export function convertToLlm(messages: AgentMessage[]): Message[] {
 						content: [{ type: "text", text: bashExecutionToText(m) }],
 						timestamp: m.timestamp,
 					};
+				case "pythonExecution":
+					if (m.excludeFromContext) {
+						return undefined;
+					}
+					return {
+						role: "user",
+						content: [{ type: "text", text: pythonExecutionToText(m) }],
+						timestamp: m.timestamp,
+					};
 				case "custom":
 				case "hookMessage": {
 					const content = typeof m.content === "string" ? [{ type: "text" as const, text: m.content }] : m.content;

package/src/core/model-registry.ts

@@ -432,6 +432,13 @@ export class ModelRegistry {
 		return this.models.find((m) => m.provider === provider && m.id === modelId);
 	}
 
+	/**
+	 * Get the base URL associated with a provider, if any model defines one.
+	 */
+	getProviderBaseUrl(provider: string): string | undefined {
+		return this.models.find((m) => m.provider === provider && m.baseUrl)?.baseUrl;
+	}
+
 	/**
 	 * Get API key for a model.
 	 */

package/src/core/python-executor.ts

@@ -27,6 +27,8 @@ export interface PythonExecutorOptions {
 	reset?: boolean;
 	/** Use shared gateway across pi instances (default: true) */
 	useSharedGateway?: boolean;
+	/** Session file path for accessing task outputs */
+	sessionFile?: string;
 }
 
 export interface PythonKernelExecutor {
@@ -79,6 +81,7 @@ export async function warmPythonEnvironment(
 	cwd: string,
 	sessionId?: string,
 	useSharedGateway?: boolean,
+	sessionFile?: string,
 ): Promise<{ ok: boolean; reason?: string; docs: PreludeHelper[] }> {
 	try {
 		await ensureKernelAvailable(cwd);
@@ -97,6 +100,7 @@ export async function warmPythonEnvironment(
 			cwd,
 			async (kernel) => kernel.introspectPrelude(),
 			useSharedGateway,
+			sessionFile,
 		);
 		cachedPreludeDocs = docs;
 		return { ok: true, docs };
@@ -115,8 +119,14 @@ export function resetPreludeDocsCache(): void {
 	cachedPreludeDocs = null;
 }
 
-async function createKernelSession(sessionId: string, cwd: string, useSharedGateway?: boolean): Promise<KernelSession> {
-	const kernel = await PythonKernel.start({ cwd, useSharedGateway });
+async function createKernelSession(
+	sessionId: string,
+	cwd: string,
+	useSharedGateway?: boolean,
+	sessionFile?: string,
+): Promise<KernelSession> {
+	const env = sessionFile ? { OMP_SESSION_FILE: sessionFile } : undefined;
+	const kernel = await PythonKernel.start({ cwd, useSharedGateway, env });
 	const session: KernelSession = {
 		id: sessionId,
 		kernel,
@@ -137,7 +147,12 @@ async function createKernelSession(sessionId: string, cwd: string, useSharedGate
 	return session;
 }
 
-async function restartKernelSession(session: KernelSession, cwd: string, useSharedGateway?: boolean): Promise<void> {
+async function restartKernelSession(
+	session: KernelSession,
+	cwd: string,
+	useSharedGateway?: boolean,
+	sessionFile?: string,
+): Promise<void> {
 	session.restartCount += 1;
 	if (session.restartCount > 1) {
 		throw new Error("Python kernel restarted too many times in this session");
@@ -147,7 +162,8 @@ async function restartKernelSession(session: KernelSession, cwd: string, useShar
 	} catch (err) {
 		logger.warn("Failed to shutdown crashed kernel", { error: err instanceof Error ? err.message : String(err) });
 	}
-	const kernel = await PythonKernel.start({ cwd, useSharedGateway });
+	const env = sessionFile ? { OMP_SESSION_FILE: sessionFile } : undefined;
+	const kernel = await PythonKernel.start({ cwd, useSharedGateway, env });
 	session.kernel = kernel;
 	session.dead = false;
 	session.lastUsedAt = Date.now();
@@ -170,17 +186,18 @@ async function withKernelSession<T>(
 	cwd: string,
 	handler: (kernel: PythonKernel) => Promise<T>,
 	useSharedGateway?: boolean,
+	sessionFile?: string,
 ): Promise<T> {
 	let session = kernelSessions.get(sessionId);
 	if (!session) {
-		session = await createKernelSession(sessionId, cwd, useSharedGateway);
+		session = await createKernelSession(sessionId, cwd, useSharedGateway, sessionFile);
 		kernelSessions.set(sessionId, session);
 	}
 
 	const run = async (): Promise<T> => {
 		session!.lastUsedAt = Date.now();
 		if (session!.dead || !session!.kernel.isAlive()) {
-			await restartKernelSession(session!, cwd, useSharedGateway);
+			await restartKernelSession(session!, cwd, useSharedGateway, sessionFile);
 		}
 		try {
 			const result = await handler(session!.kernel);
@@ -190,7 +207,7 @@ async function withKernelSession<T>(
 			if (!session!.dead && session!.kernel.isAlive()) {
 				throw err;
 			}
-			await restartKernelSession(session!, cwd, useSharedGateway);
+			await restartKernelSession(session!, cwd, useSharedGateway, sessionFile);
 			const result = await handler(session!.kernel);
 			session!.restartCount = 0;
 			return result;
@@ -273,8 +290,11 @@ export async function executePython(code: string, options?: PythonExecutorOption
 
 	const kernelMode = options?.kernelMode ?? "session";
 	const useSharedGateway = options?.useSharedGateway;
+	const sessionFile = options?.sessionFile;
+
 	if (kernelMode === "per-call") {
-		const kernel = await PythonKernel.start({ cwd, useSharedGateway });
+		const env = sessionFile ? { OMP_SESSION_FILE: sessionFile } : undefined;
+		const kernel = await PythonKernel.start({ cwd, useSharedGateway, env });
 		try {
 			return await executeWithKernel(kernel, code, options);
 		} finally {
@@ -294,5 +314,6 @@ export async function executePython(code: string, options?: PythonExecutorOption
 		cwd,
 		async (kernel) => executeWithKernel(kernel, code, options),
 		useSharedGateway,
+		sessionFile,
 	);
 }

package/src/core/python-gateway-coordinator.ts

@@ -13,7 +13,7 @@ import {
 } from "node:fs";
 import { createServer } from "node:net";
 import { delimiter, join } from "node:path";
-import { logger } from "@oh-my-pi/pi-utils";
+import { logger, postmortem } from "@oh-my-pi/pi-utils";
 import type { Subprocess } from "bun";
 import { getAgentDir } from "../config";
 import { getShellConfig, killProcessTree } from "../utils/shell";
@@ -161,6 +161,58 @@ let localGatewayUrl: string | null = null;
 let idleShutdownTimer: ReturnType<typeof setTimeout> | null = null;
 let isCoordinatorInitialized = false;
 let localClientFile: string | null = null;
+let postmortemRegistered = false;
+
+/**
+ * Register cleanup handler for process exit. Called lazily on first gateway acquisition.
+ * Ensures the gateway process we spawned is killed when omp exits, preventing orphaned processes.
+ */
+function ensurePostmortemCleanup(): void {
+	if (postmortemRegistered) return;
+	postmortemRegistered = true;
+
+	postmortem.register("shared-gateway", async () => {
+		cancelIdleShutdown();
+
+		// Clean up our client file first so refcount is accurate
+		if (localClientFile) {
+			try {
+				unlinkSync(localClientFile);
+			} catch {
+				// Ignore cleanup errors
+			}
+			localClientFile = null;
+		}
+
+		// If we spawned the gateway, kill it only if no other clients remain
+		if (localGatewayProcess) {
+			const clients = pruneStaleClientInfos(listClientInfos());
+			const remainingRefs = clients.reduce((sum, c) => sum + c.info.refCount, 0);
+
+			if (remainingRefs === 0) {
+				logger.debug("Cleaning up shared gateway on process exit", { pid: localGatewayProcess.pid });
+				try {
+					await killProcessTree(localGatewayProcess.pid);
+				} catch (err) {
+					logger.warn("Failed to kill shared gateway on exit", {
+						error: err instanceof Error ? err.message : String(err),
+					});
+				}
+				clearGatewayInfo();
+			} else {
+				logger.debug("Leaving shared gateway running for other clients", {
+					pid: localGatewayProcess.pid,
+					remainingRefs,
+				});
+			}
+
+			localGatewayProcess = null;
+			localGatewayUrl = null;
+		}
+
+		isCoordinatorInitialized = false;
+	});
+}
 
 function filterEnv(env: Record<string, string | undefined>): Record<string, string | undefined> {
 	const filtered: Record<string, string | undefined> = {};
@@ -665,6 +717,8 @@ export async function acquireSharedGateway(cwd: string): Promise<AcquireResult |
 		return null;
 	}
 
+	ensurePostmortemCleanup();
+
 	try {
 		return await withGatewayLock(async () => {
 			const existingInfo = readGatewayInfo();

package/src/core/python-prelude.py

@@ -24,14 +24,6 @@ if "__omp_prelude_loaded__" not in globals():
         _emit_status("pwd", path=str(p))
         return p
 
-    @_category("Navigation")
-    def cd(path: str | Path) -> Path:
-        """Change directory."""
-        p = Path(path).expanduser().resolve()
-        os.chdir(p)
-        _emit_status("cd", path=str(p))
-        return p
-
     @_category("Shell")
     def env(key: str | None = None, value: str | None = None):
         """Get/set environment variables."""
@@ -914,6 +906,207 @@ if "__omp_prelude_loaded__" not in globals():
         _emit_status("git_has_changes", has_changes=has_changes)
         return has_changes
 
+    @_category("Agent")
+    def output(
+        *ids: str,
+        format: str = "raw",
+        query: str | None = None,
+        offset: int | None = None,
+        limit: int | None = None,
+    ) -> str | dict | list[dict]:
+        """Read task/agent output by ID. Returns text or JSON depending on format.
+        
+        Args:
+            *ids: Output IDs to read (e.g., 'explore_0', 'reviewer_1')
+            format: 'raw' (default), 'json' (dict with metadata), 'stripped' (no ANSI)
+            query: jq-like query for JSON outputs (e.g., '.endpoints[0].file')
+            offset: Line number to start reading from (1-indexed)
+            limit: Maximum number of lines to read
+        
+        Returns:
+            Single ID: str (format='raw'/'stripped') or dict (format='json')
+            Multiple IDs: list of dict with 'id' and 'content'/'data' keys
+        
+        Examples:
+            output('explore_0')  # Read as raw text
+            output('reviewer_0', format='json')  # Read with metadata
+            output('explore_0', query='.files[0]')  # Extract JSON field
+            output('explore_0', offset=10, limit=20)  # Lines 10-29
+            output('explore_0', 'reviewer_1')  # Read multiple outputs
+        """
+        session_file = os.environ.get("OMP_SESSION_FILE")
+        if not session_file:
+            _emit_status("output", error="No session file available")
+            raise RuntimeError("No session - output artifacts unavailable")
+        
+        artifacts_dir = session_file.rsplit(".", 1)[0]  # Strip .jsonl extension
+        if not Path(artifacts_dir).exists():
+            _emit_status("output", error="Artifacts directory not found", path=artifacts_dir)
+            raise RuntimeError(f"No artifacts directory found: {artifacts_dir}")
+        
+        if not ids:
+            _emit_status("output", error="No IDs provided")
+            raise ValueError("At least one output ID is required")
+        
+        if query and (offset is not None or limit is not None):
+            _emit_status("output", error="query cannot be combined with offset/limit")
+            raise ValueError("query cannot be combined with offset/limit")
+        
+        results: list[dict] = []
+        not_found: list[str] = []
+        
+        for output_id in ids:
+            output_path = Path(artifacts_dir) / f"{output_id}.md"
+            if not output_path.exists():
+                not_found.append(output_id)
+                continue
+            
+            raw_content = output_path.read_text(encoding="utf-8")
+            raw_lines = raw_content.splitlines()
+            total_lines = len(raw_lines)
+            
+            selected_content = raw_content
+            range_info: dict | None = None
+            
+            # Handle query
+            if query:
+                try:
+                    json_value = json.loads(raw_content)
+                except json.JSONDecodeError as e:
+                    _emit_status("output", id=output_id, error=f"Not valid JSON: {e}")
+                    raise ValueError(f"Output {output_id} is not valid JSON: {e}")
+                
+                # Apply jq-like query
+                result_value = _apply_query(json_value, query)
+                try:
+                    selected_content = json.dumps(result_value, indent=2) if result_value is not None else "null"
+                except (TypeError, ValueError):
+                    selected_content = str(result_value)
+            
+            # Handle offset/limit
+            elif offset is not None or limit is not None:
+                start_line = max(1, offset or 1)
+                if start_line > total_lines:
+                    _emit_status("output", id=output_id, error=f"Offset {start_line} beyond end ({total_lines} lines)")
+                    raise ValueError(f"Offset {start_line} is beyond end of output ({total_lines} lines) for {output_id}")
+                
+                effective_limit = limit if limit is not None else total_lines - start_line + 1
+                end_line = min(total_lines, start_line + effective_limit - 1)
+                selected_lines = raw_lines[start_line - 1 : end_line]
+                selected_content = "\n".join(selected_lines)
+                range_info = {"start_line": start_line, "end_line": end_line, "total_lines": total_lines}
+            
+            # Strip ANSI codes if requested
+            if format == "stripped":
+                import re
+                selected_content = re.sub(r"\x1b\[[0-9;]*m", "", selected_content)
+            
+            # Build result
+            if format == "json":
+                result_data = {
+                    "id": output_id,
+                    "path": str(output_path),
+                    "line_count": total_lines if not query else len(selected_content.splitlines()),
+                    "char_count": len(raw_content) if not query else len(selected_content),
+                    "content": selected_content,
+                }
+                if range_info:
+                    result_data["range"] = range_info
+                if query:
+                    result_data["query"] = query
+                results.append(result_data)
+            else:
+                results.append({"id": output_id, "content": selected_content})
+        
+        # Handle not found
+        if not_found:
+            available = sorted(
+                [f.stem for f in Path(artifacts_dir).glob("*.md")]
+            )
+            error_msg = f"Output not found: {', '.join(not_found)}"
+            if available:
+                error_msg += f"\n\nAvailable outputs: {', '.join(available[:20])}"
+                if len(available) > 20:
+                    error_msg += f" (and {len(available) - 20} more)"
+            _emit_status("output", not_found=not_found, available_count=len(available))
+            raise FileNotFoundError(error_msg)
+        
+        # Return format
+        if len(ids) == 1:
+            if format == "json":
+                _emit_status("output", id=ids[0], chars=results[0]["char_count"])
+                return results[0]
+            _emit_status("output", id=ids[0], chars=len(results[0]["content"]))
+            return results[0]["content"]
+        
+        # Multiple IDs
+        if format == "json":
+            total_chars = sum(r["char_count"] for r in results)
+            _emit_status("output", count=len(results), total_chars=total_chars)
+            return results
+        
+        combined_output: list[dict] = []
+        for r in results:
+            combined_output.append({"id": r["id"], "content": r["content"]})
+        total_chars = sum(len(r["content"]) for r in combined_output)
+        _emit_status("output", count=len(combined_output), total_chars=total_chars)
+        return combined_output
+
+    def _apply_query(data: any, query: str) -> any:
+        """Apply jq-like query to data. Supports .key, [index], and chaining."""
+        if not query:
+            return data
+        
+        query = query.strip()
+        if query.startswith("."):
+            query = query[1:]
+        if not query:
+            return data
+        
+        # Parse query into tokens
+        tokens = []
+        current_token = ""
+        i = 0
+        while i < len(query):
+            ch = query[i]
+            if ch == ".":
+                if current_token:
+                    tokens.append(("key", current_token))
+                    current_token = ""
+            elif ch == "[":
+                if current_token:
+                    tokens.append(("key", current_token))
+                    current_token = ""
+                # Find matching ]
+                j = i + 1
+                while j < len(query) and query[j] != "]":
+                    j += 1
+                bracket_content = query[i+1:j]
+                if bracket_content.startswith('"') and bracket_content.endswith('"'):
+                    tokens.append(("key", bracket_content[1:-1]))
+                else:
+                    tokens.append(("index", int(bracket_content)))
+                i = j
+            else:
+                current_token += ch
+            i += 1
+        if current_token:
+            tokens.append(("key", current_token))
+        
+        # Apply tokens
+        current = data
+        for token_type, value in tokens:
+            if token_type == "index":
+                if not isinstance(current, list) or value >= len(current):
+                    return None
+                current = current[value]
+            elif token_type == "key":
+                if not isinstance(current, dict) or value not in current:
+                    return None
+                current = current[value]
+        
+        return current
+
     def __omp_prelude_docs__() -> list[dict[str, str]]:
         """Return prelude helper docs for templating. Discovers functions by _omp_category attribute."""
         helpers: list[dict[str, str]] = []

package/src/core/session-manager.ts

@@ -13,6 +13,7 @@ import {
 	createCustomMessage,
 	type FileMentionMessage,
 	type HookMessage,
+	type PythonExecutionMessage,
 } from "./messages";
 import type { SessionStorage, SessionStorageWriter } from "./session-storage";
 import { FileSessionStorage, MemorySessionStorage } from "./session-storage";
@@ -1306,7 +1307,15 @@ export class SessionManager {
 	 * so it is easier to find them.
 	 * These need to be appended via appendCompaction() and appendBranchSummary() methods.
 	 */
-	appendMessage(message: Message | CustomMessage | HookMessage | BashExecutionMessage | FileMentionMessage): string {
+	appendMessage(
+		message:
+			| Message
+			| CustomMessage
+			| HookMessage
+			| BashExecutionMessage
+			| PythonExecutionMessage
+			| FileMentionMessage,
+	): string {
 		const entry: SessionMessageEntry = {
 			type: "message",
 			id: generateId(this.byId),