File size: 5,028 Bytes
f500658 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | /**
* Print mode (single-shot): Send prompts, output result, exit.
*
* Used for:
* - `pi -p "prompt"` - text output
* - `pi --mode json "prompt"` - JSON event stream
*/
import type { AssistantMessage, ImageContent } from "@earendil-works/pi-ai";
import type { AgentSessionRuntime } from "../core/agent-session-runtime.ts";
import { flushRawStdout, waitForRawStdoutBackpressure, writeRawStdout } from "../core/output-guard.ts";
import { killTrackedDetachedChildren } from "../utils/shell.ts";
import { toJsonEvent } from "./json-event.ts";
/**
* Options for print mode.
*/
export interface PrintModeOptions {
/** Output mode: "text" for final response only, "json" for all events */
mode: "text" | "json";
/** Array of additional prompts to send after initialMessage */
messages?: string[];
/** First message to send (may contain @file content) */
initialMessage?: string;
/** Images to attach to the initial message */
initialImages?: ImageContent[];
}
/**
* Run in print (single-shot) mode.
* Sends prompts to the agent and outputs the result.
*/
export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: PrintModeOptions): Promise<number> {
const { mode, messages = [], initialMessage, initialImages } = options;
let exitCode = 0;
let session = runtimeHost.session;
let unsubscribe: (() => void) | undefined;
let unsubscribeBackpressure: (() => void) | undefined;
let disposed = false;
const signalCleanupHandlers: Array<() => void> = [];
const disposeRuntime = async (): Promise<void> => {
if (disposed) return;
disposed = true;
unsubscribe?.();
unsubscribeBackpressure?.();
await runtimeHost.dispose();
};
const registerSignalHandlers = (): void => {
const signals: NodeJS.Signals[] = ["SIGTERM"];
if (process.platform !== "win32") {
signals.push("SIGHUP");
}
for (const signal of signals) {
const handler = () => {
killTrackedDetachedChildren();
void disposeRuntime().finally(() => {
process.exit(signal === "SIGHUP" ? 129 : 143);
});
};
process.on(signal, handler);
signalCleanupHandlers.push(() => process.off(signal, handler));
}
};
registerSignalHandlers();
runtimeHost.setRebindSession(async () => {
await rebindSession();
});
const rebindSession = async (): Promise<void> => {
session = runtimeHost.session;
await session.bindExtensions({
mode: mode === "json" ? "json" : "print",
commandContextActions: {
waitForIdle: () => session.waitForIdle(),
newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions),
fork: async (entryId, forkOptions) => {
const result = await runtimeHost.fork(entryId, forkOptions);
return { cancelled: result.cancelled };
},
navigateTree: async (targetId, navigateOptions) => {
const result = await session.navigateTree(targetId, {
summarize: navigateOptions?.summarize,
customInstructions: navigateOptions?.customInstructions,
replaceInstructions: navigateOptions?.replaceInstructions,
label: navigateOptions?.label,
});
return { cancelled: result.cancelled };
},
switchSession: async (sessionPath, switchOptions) => {
return runtimeHost.switchSession(sessionPath, switchOptions);
},
reload: async () => {
await session.reload();
},
},
onError: (err) => {
console.error(`Extension error (${err.extensionPath}): ${err.error}`);
},
});
unsubscribe?.();
unsubscribeBackpressure?.();
unsubscribe = session.subscribe((event) => {
if (mode === "json") {
writeRawStdout(`${JSON.stringify(toJsonEvent(event))}\n`);
}
});
unsubscribeBackpressure =
mode === "json"
? session.agent.subscribe(async () => {
await waitForRawStdoutBackpressure();
})
: undefined;
};
try {
if (mode === "json") {
const header = session.sessionManager.getHeader();
if (header) {
writeRawStdout(`${JSON.stringify(header)}\n`);
}
}
await rebindSession();
if (initialMessage) {
await session.prompt(initialMessage, { images: initialImages });
}
for (const message of messages) {
await session.prompt(message);
}
if (mode === "text") {
const state = session.state;
const lastMessage = state.messages[state.messages.length - 1];
if (lastMessage?.role === "assistant") {
const assistantMsg = lastMessage as AssistantMessage;
if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") {
console.error(assistantMsg.errorMessage || `Request ${assistantMsg.stopReason}`);
exitCode = 1;
} else {
for (const content of assistantMsg.content) {
if (content.type === "text") {
writeRawStdout(`${content.text}\n`);
}
}
}
}
}
return exitCode;
} catch (error: unknown) {
console.error(error instanceof Error ? error.message : String(error));
return 1;
} finally {
for (const cleanup of signalCleanupHandlers) {
cleanup();
}
await disposeRuntime();
await flushRawStdout();
}
}
|