#!/usr/bin/env python3 from __future__ import annotations import json import sys from pathlib import Path PI_PKG = Path("/home/sebas/.npm-global/lib/node_modules/@mariozechner/pi-coding-agent") PACKAGE_JSON = PI_PKG / "package.json" EXPECTED_VERSION = "0.67.6" def replace_once(text: str, old: str, new: str, path: Path) -> tuple[str, bool]: if new in text: return text, False if old not in text: raise SystemExit(f"patch target not found in {path}: {old[:80]!r}") return text.replace(old, new, 1), True def patch_file(path: Path, replacements: list[tuple[str, str]]) -> bool: text = path.read_text() changed = False for old, new in replacements: text, did = replace_once(text, old, new, path) changed = changed or did if changed: path.write_text(text) return changed def main() -> int: if not PACKAGE_JSON.exists(): print(f"missing pi package: {PACKAGE_JSON}", file=sys.stderr) return 1 version = json.loads(PACKAGE_JSON.read_text()).get("version") if version != EXPECTED_VERSION: print( f"warning: expected pi version {EXPECTED_VERSION}, found {version}. patch may need refresh.", file=sys.stderr, ) patches: dict[Path, list[tuple[str, str]]] = { PI_PKG / "dist/core/extensions/types.d.ts": [ ( """/** Fired before forking a session (can be cancelled) */\nexport interface SessionBeforeForkEvent {\n type: \"session_before_fork\";\n entryId: string;\n}\n/** Fired before context compaction (can be cancelled or customized) */\nexport interface SessionBeforeCompactEvent {\n""", """/** Fired before forking a session (can be cancelled) */\nexport interface SessionBeforeForkEvent {\n type: \"session_before_fork\";\n entryId: string;\n}\n/** Fired before graceful shutdown paths such as Ctrl+C, Ctrl+D, /quit, or extension-triggered exit (can be cancelled). */\nexport interface SessionBeforeShutdownEvent {\n type: \"session_before_shutdown\";\n reason: \"ctrl_c\" | \"ctrl_d\" | \"command\" | \"extension\" | \"signal\" | \"rpc_eof\" | \"rpc_shutdown\" | \"unknown\";\n signal?: \"SIGHUP\" | \"SIGTERM\";\n}\n/** Fired before context compaction (can be cancelled or customized) */\nexport interface SessionBeforeCompactEvent {\n""", ), ( "export type SessionEvent = SessionStartEvent | SessionBeforeSwitchEvent | SessionBeforeForkEvent | SessionBeforeCompactEvent | SessionCompactEvent | SessionShutdownEvent | SessionBeforeTreeEvent | SessionTreeEvent;\n", "export type SessionEvent = SessionStartEvent | SessionBeforeSwitchEvent | SessionBeforeForkEvent | SessionBeforeShutdownEvent | SessionBeforeCompactEvent | SessionCompactEvent | SessionShutdownEvent | SessionBeforeTreeEvent | SessionTreeEvent;\n", ), ( """export interface SessionBeforeForkResult {\n cancel?: boolean;\n skipConversationRestore?: boolean;\n}\nexport interface SessionBeforeCompactResult {\n""", """export interface SessionBeforeForkResult {\n cancel?: boolean;\n skipConversationRestore?: boolean;\n}\nexport interface SessionBeforeShutdownResult {\n cancel?: boolean;\n}\nexport interface SessionBeforeCompactResult {\n""", ), ( """ on(event: \"session_before_switch\", handler: ExtensionHandler): void;\n on(event: \"session_before_fork\", handler: ExtensionHandler): void;\n on(event: \"session_before_compact\", handler: ExtensionHandler): void;\n""", """ on(event: \"session_before_switch\", handler: ExtensionHandler): void;\n on(event: \"session_before_fork\", handler: ExtensionHandler): void;\n on(event: \"session_before_shutdown\", handler: ExtensionHandler): void;\n on(event: \"session_before_compact\", handler: ExtensionHandler): void;\n""", ), ], PI_PKG / "dist/core/extensions/index.d.ts": [ ( "SessionBeforeCompactEvent, SessionBeforeCompactResult, SessionBeforeForkEvent, SessionBeforeForkResult, SessionBeforeSwitchEvent, SessionBeforeSwitchResult, SessionBeforeTreeEvent, SessionBeforeTreeResult", "SessionBeforeCompactEvent, SessionBeforeCompactResult, SessionBeforeForkEvent, SessionBeforeForkResult, SessionBeforeShutdownEvent, SessionBeforeShutdownResult, SessionBeforeSwitchEvent, SessionBeforeSwitchResult, SessionBeforeTreeEvent, SessionBeforeTreeResult", ), ], PI_PKG / "dist/core/extensions/runner.d.ts": [ ( "SessionBeforeCompactResult, SessionBeforeForkResult, SessionBeforeSwitchResult, SessionBeforeTreeResult", "SessionBeforeCompactResult, SessionBeforeForkResult, SessionBeforeShutdownResult, SessionBeforeSwitchResult, SessionBeforeTreeResult", ), ( """} ? SessionBeforeSwitchResult | undefined : TEvent extends {\n type: \"session_before_fork\";\n} ? SessionBeforeForkResult | undefined : TEvent extends {\n type: \"session_before_compact\";\n} ? SessionBeforeCompactResult | undefined : TEvent extends {\n""", """} ? SessionBeforeSwitchResult | undefined : TEvent extends {\n type: \"session_before_fork\";\n} ? SessionBeforeForkResult | undefined : TEvent extends {\n type: \"session_before_shutdown\";\n} ? SessionBeforeShutdownResult | undefined : TEvent extends {\n type: \"session_before_compact\";\n} ? SessionBeforeCompactResult | undefined : TEvent extends {\n""", ), ], PI_PKG / "dist/core/extensions/runner.js": [ ( """ isSessionBeforeEvent(event) {\n return (event.type === \"session_before_switch\" ||\n event.type === \"session_before_fork\" ||\n event.type === \"session_before_compact\" ||\n event.type === \"session_before_tree\");\n }\n""", """ isSessionBeforeEvent(event) {\n return (event.type === \"session_before_switch\" ||\n event.type === \"session_before_fork\" ||\n event.type === \"session_before_shutdown\" ||\n event.type === \"session_before_compact\" ||\n event.type === \"session_before_tree\");\n }\n""", ), ], PI_PKG / "dist/modes/interactive/interactive-mode.js": [ ( """ // Shutdown state\n shutdownRequested = false;\n // Extension UI state\n""", """ // Shutdown state\n shutdownRequested = false;\n shutdownRequestReason = \"extension\";\n shutdownRequestSignal = undefined;\n // Extension UI state\n""", ), ( """ shutdownHandler: () => {\n this.shutdownRequested = true;\n if (!this.session.isStreaming) {\n void this.shutdown();\n }\n },\n""", """ shutdownHandler: () => {\n this.shutdownRequested = true;\n this.shutdownRequestReason = \"extension\";\n this.shutdownRequestSignal = undefined;\n if (!this.session.isStreaming) {\n void this.shutdown(\"extension\");\n }\n },\n""", ), ( """ shutdown: () => {\n this.shutdownRequested = true;\n },\n""", """ shutdown: () => {\n this.shutdownRequested = true;\n this.shutdownRequestReason = \"extension\";\n this.shutdownRequestSignal = undefined;\n },\n""", ), ( """ handleCtrlC() {\n const now = Date.now();\n if (now - this.lastSigintTime < 500) {\n void this.shutdown();\n }\n else {\n this.clearEditor();\n this.lastSigintTime = now;\n }\n }\n handleCtrlD() {\n // Only called when editor is empty (enforced by CustomEditor)\n void this.shutdown();\n }\n""", """ handleCtrlC() {\n const now = Date.now();\n if (now - this.lastSigintTime < 500) {\n void this.shutdown(\"ctrl_c\");\n }\n else {\n this.clearEditor();\n this.lastSigintTime = now;\n }\n }\n handleCtrlD() {\n // Only called when editor is empty (enforced by CustomEditor)\n void this.shutdown(\"ctrl_d\");\n }\n""", ), ( """ /**\n * Gracefully shutdown the agent.\n * Emits shutdown event to extensions, then exits.\n */\n isShuttingDown = false;\n async shutdown() {\n if (this.isShuttingDown)\n return;\n this.isShuttingDown = true;\n this.unregisterSignalHandlers();\n await this.runtimeHost.dispose();\n""", """ async emitBeforeShutdown(reason = \"unknown\", signal) {\n const runner = this.runtimeHost.session.extensionRunner;\n if (!runner?.hasHandlers(\"session_before_shutdown\")) {\n return { cancelled: false };\n }\n const result = await runner.emit({\n type: \"session_before_shutdown\",\n reason,\n signal,\n });\n return { cancelled: result?.cancel === true };\n }\n /**\n * Gracefully shutdown the agent.\n * Emits shutdown event to extensions, then exits.\n */\n isShuttingDown = false;\n async shutdown(reason = \"unknown\", signal) {\n if (this.isShuttingDown)\n return;\n this.isShuttingDown = true;\n const beforeResult = await this.emitBeforeShutdown(reason, signal);\n if (beforeResult.cancelled) {\n this.isShuttingDown = false;\n this.shutdownRequested = false;\n this.shutdownRequestReason = \"extension\";\n this.shutdownRequestSignal = undefined;\n return;\n }\n this.unregisterSignalHandlers();\n await this.runtimeHost.dispose();\n""", ), ( """ async checkShutdownRequested() {\n if (!this.shutdownRequested)\n return;\n await this.shutdown();\n }\n""", """ async checkShutdownRequested() {\n if (!this.shutdownRequested)\n return;\n await this.shutdown(this.shutdownRequestReason, this.shutdownRequestSignal);\n }\n""", ), ( """ const handler = () => {\n killTrackedDetachedChildren();\n void this.shutdown();\n };\n""", """ const handler = () => {\n killTrackedDetachedChildren();\n void this.shutdown(\"signal\", signal);\n };\n""", ), ( """ if (text === \"/quit\") {\n this.editor.setText(\"\");\n await this.shutdown();\n return;\n }\n""", """ if (text === \"/quit\") {\n this.editor.setText(\"\");\n await this.shutdown(\"command\");\n return;\n }\n""", ), ( """ }, () => {\n void this.shutdown();\n }, () => this.ui.requestRender(), {\n""", """ }, () => {\n void this.shutdown(\"command\");\n }, () => this.ui.requestRender(), {\n""", ), ], PI_PKG / "dist/modes/rpc/rpc-mode.js": [ ( """ // Shutdown request flag\n let shutdownRequested = false;\n let shuttingDown = false;\n const signalCleanupHandlers = [];\n""", """ // Shutdown request flag\n let shutdownRequested = false;\n let shutdownRequestReason = \"extension\";\n let shutdownRequestSignal = undefined;\n let shuttingDown = false;\n const signalCleanupHandlers = [];\n""", ), ( """ shutdownHandler: () => {\n shutdownRequested = true;\n },\n""", """ shutdownHandler: () => {\n shutdownRequested = true;\n shutdownRequestReason = \"extension\";\n shutdownRequestSignal = undefined;\n },\n""", ), ( """ const handler = () => {\n killTrackedDetachedChildren();\n void shutdown(signal === \"SIGHUP\" ? 129 : 143);\n };\n""", """ const handler = () => {\n killTrackedDetachedChildren();\n void shutdown(signal === \"SIGHUP\" ? 129 : 143, \"signal\", signal);\n };\n""", ), ( """ let detachInput = () => { };\n async function shutdown(exitCode = 0) {\n if (shuttingDown) {\n process.exit(exitCode);\n }\n shuttingDown = true;\n for (const cleanup of signalCleanupHandlers) {\n cleanup();\n }\n unsubscribe?.();\n await runtimeHost.dispose();\n detachInput();\n process.stdin.pause();\n process.exit(exitCode);\n }\n async function checkShutdownRequested() {\n if (!shutdownRequested)\n return;\n await shutdown();\n }\n""", """ let detachInput = () => { };\n async function emitBeforeShutdown(reason = \"rpc_shutdown\", signal) {\n const runner = runtimeHost.session.extensionRunner;\n if (!runner?.hasHandlers(\"session_before_shutdown\")) {\n return { cancelled: false };\n }\n const result = await runner.emit({\n type: \"session_before_shutdown\",\n reason,\n signal,\n });\n return { cancelled: result?.cancel === true };\n }\n async function shutdown(exitCode = 0, reason = \"rpc_shutdown\", signal) {\n if (shuttingDown) {\n process.exit(exitCode);\n }\n shuttingDown = true;\n const beforeResult = await emitBeforeShutdown(reason, signal);\n if (beforeResult.cancelled) {\n shuttingDown = false;\n shutdownRequested = false;\n shutdownRequestReason = \"extension\";\n shutdownRequestSignal = undefined;\n return;\n }\n for (const cleanup of signalCleanupHandlers) {\n cleanup();\n }\n unsubscribe?.();\n await runtimeHost.dispose();\n detachInput();\n process.stdin.pause();\n process.exit(exitCode);\n }\n async function checkShutdownRequested() {\n if (!shutdownRequested)\n return;\n await shutdown(0, shutdownRequestReason, shutdownRequestSignal);\n }\n""", ), ( """ const onInputEnd = () => {\n void shutdown();\n };\n""", """ const onInputEnd = () => {\n void shutdown(0, \"rpc_eof\");\n };\n""", ), ], } changed_any = False for path, replacements in patches.items(): if not path.exists(): raise SystemExit(f"missing patch target: {path}") changed = patch_file(path, replacements) changed_any = changed_any or changed print(f"patched {path}: {'changed' if changed else 'already'}") print("done") return 0 if __name__ == "__main__": raise SystemExit(main())