1/**
2 * PiRpc — stdio JSONL transport for the Pi coding agent's RPC mode.
3 *
4 * Spawns `pi --mode rpc` and speaks Pi's line-delimited JSON protocol:
5 * requests go to stdin as `{"type": "...", "id": "..."}` records, responses
6 * come back as `{"type": "response", "id": ..., "success": ...}` and are
7 * correlated by `id`; every other stdout record is a session event and is
8 * surfaced on the `events` queue in arrival order.
9 *
10 * Framing follows Pi's spec: LF-delimited only, with a trailing `\r`
11 * stripped. Lines are split manually (never with `readline`, which also
12 * splits on U+2028/U+2029 and would corrupt frames). Records that fail to
13 * parse as JSON are dropped with a debug log rather than failing the
14 * transport, so a chatty extension cannot take the session down.
15 *
16 * Used by `PiAdapterV2` for sessions and by `PiTextGeneration` /
17 * `PiProvider` for ephemeral one-shot processes.
18 */
19import * as Clock from "effect/Clock";
20import * as Deferred from "effect/Deferred";
21import * as Duration from "effect/Duration";
22import * as Effect from "effect/Effect";
23import * as Option from "effect/Option";
24import * as Predicate from "effect/Predicate";
25import * as Queue from "effect/Queue";
26import * as Scope from "effect/Scope";
27import * as Schema from "effect/Schema";
28import * as Stream from "effect/Stream";
29import * as ChildProcess from "effect/process/ChildProcess";
30import * as ChildProcessSpawner from "effect/process/ChildProcessSpawner";
31import * as HostProcess from "@t3tools/shared/HostProcess";
32import { resolveSpawnCommand } from "@t3tools/shared/shell";
33import {
34 MAX_TOOL_OUTPUT_IMAGES,
35 MAX_TOOL_OUTPUT_IMAGE_BASE64_LENGTH,
36} from "@t3tools/shared/toolOutput";
37
38import { signalProcessGroup } from "@t3tools/provider-core/server/processGroup";
39
40export class PiRpcError extends Schema.TaggedError<PiRpcError>()("PiRpcError", {
41 operation: Schema.String,
42 detail: Schema.optional(Schema.String),
43 cause: Schema.optional(Schema.Defect()),
44}) {
45 override get message(): string {
46 return `Pi RPC ${this.operation} failed${this.detail === undefined ? "" : `: ${this.detail}`}.`;
47 }
48}
49
50export class PiRpcTimeoutError extends Schema.TaggedError<PiRpcTimeoutError>()(
51 "PiRpcTimeoutError",
52 {
53 operation: Schema.String,
54 timeoutMs: Schema.Int.check(Schema.isBetween({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER })),
55 },
56) {
57 override get message(): string {
58 return `Pi RPC ${this.operation} failed: timed out after ${this.timeoutMs}ms.`;
59 }
60}
61
62export type PiRpcRecord = Record<string, unknown>;
63
64export function piRecordField(input: unknown, key: string): unknown {
65 return Predicate.isObject(input) ? input[key] : undefined;
66}
67
68export function piRecordString(input: unknown, key: string): string | undefined {
69 const value = piRecordField(input, key);
70 return Predicate.isString(value) ? value : undefined;
71}
72
73export function piRecordNumber(input: unknown, key: string): number | undefined {
74 const value = piRecordField(input, key);
75 return Predicate.isNumber(value) && Number.isFinite(value) ? value : undefined;
76}
77
78/**
79 * Splits a `provider/model` slug into the two fields `set_model` expects.
80 * Returns null for slugs without a usable separator so callers can reject the
81 * selection instead of silently leaving Pi on its configured default.
82 */
83export function parsePiModelSlug(slug: string): { provider: string; modelId: string } | null {
84 const separator = slug.indexOf("/");
85 if (separator <= 0 || separator === slug.length - 1) return null;
86 return { provider: slug.slice(0, separator), modelId: slug.slice(separator + 1) };
87}
88
89export interface PiRpcSpawnOptions {
90 readonly command: string;
91 readonly args: ReadonlyArray<string>;
92 readonly cwd: string | undefined;
93 readonly env: NodeJS.ProcessEnv;
94}
95
96export interface PiRpcConnection {
97 /** Fire-and-forget write (used for `extension_ui_response`). */
98 readonly send: (record: PiRpcRecord) => Effect.Effect<void, PiRpcError>;
99 /**
100 * Correlated request: assigns an `id`, waits for the matching response
101 * record, and returns its `data` (undefined when the command carries none).
102 * Fails on `success: false`, transport death, or timeout.
103 */
104 readonly request: (
105 record: PiRpcRecord,
106 timeoutMs?: number,
107 ) => Effect.Effect<unknown, PiRpcError | PiRpcTimeoutError>;
108 /**
109 * Session events (every non-response stdout record) in arrival order. The
110 * full queue is exposed so consumers can append order-preserving synthetic
111 * records of their own (see PiAdapterV2's settle probe).
112 */
113 readonly events: Queue.Queue<PiRpcRecord, PiRpcError>;
114 /** Resolves when the process has exited, with its exit code. */
115 readonly exited: Effect.Effect<number, PiRpcError>;
116 /**
117 * Kill the pi process group immediately (SIGTERM, grace, SIGKILL). Used by
118 * Stop-with-restart when the process may be wedged and `abort` cannot be
119 * trusted to land. The transport fails and the session manager respawns a
120 * fresh process on the next turn.
121 */
122 readonly terminate: Effect.Effect<void>;
123}
124
125const DEFAULT_REQUEST_TIMEOUT_MS = 30_000;
126const TERMINATION_GRACE = Duration.seconds(1);
127
128interface PendingPiRequest {
129 readonly deferred: Deferred.Deferred<unknown, PiRpcError>;
130}
131
132// MCP completion events carry images in both model content and the script
133// CallToolResult. Allow the shared image budget twice, plus bounded JSON/text
134// overhead, before dropping an extension's oversized record.
135const MAX_PI_RECORD_CHARS =
136 2 * MAX_TOOL_OUTPUT_IMAGES * MAX_TOOL_OUTPUT_IMAGE_BASE64_LENGTH + 8 * 1024 * 1024;
137
138function makeJsonlFramer() {
139 let buffer = "";
140 let dropping = false;
141 return (chunk: string): ReadonlyArray<string> => {
142 const lines: string[] = [];
143 let start = 0;
144 while (start < chunk.length) {
145 const newline = chunk.indexOf("\n", start);
146 const end = newline < 0 ? chunk.length : newline;
147 if (!dropping) {
148 if (buffer.length + end - start > MAX_PI_RECORD_CHARS) {
149 buffer = "";
150 dropping = true;
151 } else {
152 buffer += chunk.slice(start, end);
153 }
154 }
155 if (newline < 0) break;
156 if (!dropping && buffer.length > 0) {
157 lines.push(buffer.endsWith("\r") ? buffer.slice(0, -1) : buffer);
158 }
159 buffer = "";
160 dropping = false;
161 start = newline + 1;
162 }
163 return lines;
164 };
165}
166
167const UnknownFromJsonString = Schema.fromJsonString(Schema.Unknown);
168const decodeJsonLine = Schema.decodeSync(UnknownFromJsonString);
169const encodeJsonLine = Schema.encodeSync(UnknownFromJsonString);
170
171const PI_ERROR_DETAIL_MAX_CHARS = 200;
172
173/**
174 * Bounded, human-readable summary of a failed response's `error` payload.
175 * The untruncated value stays on the error's `cause`, so `message` never
176 * carries unbounded remote text while logs keep something diagnostic.
177 */
178function summarizePiError(error: unknown): string {
179 const text = typeof error === "string" ? error : JSON.stringify(error);
180 if (text === undefined) return "unknown error";
181 return text.length > PI_ERROR_DETAIL_MAX_CHARS
182 ? `${text.slice(0, PI_ERROR_DETAIL_MAX_CHARS)}…`
183 : text;
184}
185
186/**
187 * Stdout closure is not enough to diagnose an early crash: Pi 0.84+ can exit
188 * 1 on import before writing any protocol line. The numeric exit code is
189 * sanitized; stderr stays out of public details because it can carry
190 * credentials or prompt text.
191 */
192function describePiStdoutClosure(exitCode: number | undefined): string {
193 if (exitCode === undefined || !Number.isFinite(exitCode)) {
194 return "pi process closed stdout";
195 }
196 return `pi process exited with code ${exitCode}`;
197}
198
199function parsePiRecord(line: string): PiRpcRecord | undefined {
200 try {
201 const parsed: unknown = decodeJsonLine(line);
202 return Predicate.isObject(parsed) ? parsed : undefined;
203 } catch {
204 return undefined;
205 }
206}
207
208/**
209 * Kill the pi process group: SIGTERM, short grace, then SIGKILL.
210 *
211 * `hasExited` is consulted before each signal. Once the original child is
212 * gone its pid/pgid can be recycled by the OS, so escalating blindly could
213 * deliver SIGKILL to an unrelated process.
214 */
215const terminatePiProcess = (kill: (signal: NodeJS.Signals) => boolean, hasExited: () => boolean) =>
216 Effect.gen(function* () {
217 if (hasExited()) return;
218 if (!kill("SIGTERM")) return;
219 yield* Effect.sleep(TERMINATION_GRACE);
220 if (hasExited()) return;
221 kill("SIGKILL");
222 });
223
224export const makePiRpcConnection = Effect.fnUntraced(function* (options: PiRpcSpawnOptions) {
225 const spawner = yield* ChildProcessSpawner.ChildProcessSpawner;
226 const platform = yield* HostProcess.Platform;
227 const scope = yield* Effect.scope;
228
229 const spawnCommand = yield* resolveSpawnCommand(options.command, [...options.args], {
230 env: options.env,
231 }).pipe(Effect.mapError((cause) => new PiRpcError({ operation: "spawn", cause })));
232 const child = yield* spawner
233 .spawn(
234 ChildProcess.make(spawnCommand.command, spawnCommand.args, {
235 ...(options.cwd === undefined ? {} : { cwd: options.cwd }),
236 env: options.env,
237 extendEnv: false,
238 shell: spawnCommand.shell,
239 detached: platform !== "win32",
240 }),
241 )
242 .pipe(Effect.mapError((cause) => new PiRpcError({ operation: "spawn", cause })));
243
244 let childExited = false;
245 let diagnosingStdoutClose = false;
246
247 const killProcessGroup = (signal: NodeJS.Signals): boolean => {
248 try {
249 if (platform === "win32") {
250 process.kill(Number(child.pid), signal);
251 } else {
252 signalProcessGroup(Number(child.pid), signal);
253 }
254 return true;
255 } catch {
256 return false;
257 }
258 };
259
260 /** Signal 0 probes liveness without delivering anything. */
261 const hasExited = (): boolean => {
262 if (childExited) return true;
263 try {
264 if (platform === "win32") process.kill(Number(child.pid), 0);
265 else signalProcessGroup(Number(child.pid), 0);
266 return false;
267 } catch {
268 return true;
269 }
270 };
271
272 /**
273 * Windows has no process groups, so `process.kill` reaches only pi itself
274 * and leaves extension subprocesses running with inherited stdio handles.
275 * `taskkill /T` reaps the whole tree.
276 */
277 const terminateWindowsTree = Effect.gen(function* () {
278 if (hasExited()) return;
279 const taskkill = yield* spawner.spawn(
280 ChildProcess.make("taskkill", ["/PID", String(child.pid), "/T", "/F"]),
281 );
282 yield* taskkill.exitCode;
283 }).pipe(Effect.scoped, Effect.ignore);
284
285 const terminateProcess =
286 platform === "win32" ? terminateWindowsTree : terminatePiProcess(killProcessGroup, hasExited);
287
288 // Registered before any further setup: an interrupt or failure between the
289 // spawn and the rest of this constructor would otherwise leak a detached
290 // pi process with no finalizer to reap it.
291 yield* Scope.addFinalizer(scope, terminateProcess.pipe(Effect.ignore, Effect.uninterruptible));
292
293 const pendingRequests = new Map<string, PendingPiRequest>();
294 const events = yield* Queue.unbounded<PiRpcRecord, PiRpcError>();
295 const outgoing = yield* Queue.unbounded<Uint8Array, PiRpcError>();
296 const transportDown = yield* Deferred.make<never, PiRpcError>();
297 const exitDeferred = yield* Deferred.make<number, PiRpcError>();
298 let nextRequestId = 0;
299
300 const failTransport = (error: PiRpcError) =>
301 Effect.gen(function* () {
302 const claimed = yield* Deferred.fail(transportDown, error);
303 if (!claimed) return;
304 for (const [key, pending] of pendingRequests) {
305 pendingRequests.delete(key);
306 yield* Deferred.fail(pending.deferred, error);
307 }
308 // Closing `outgoing` is what makes `send` non-racy: once the writer is
309 // gone every later offer is refused rather than silently buffered.
310 yield* Queue.fail(outgoing, error);
311 yield* Queue.fail(events, error);
312 });
313
314 const routeRecord = (record: PiRpcRecord) =>
315 Effect.gen(function* () {
316 if (record["type"] === "response" && typeof record["id"] === "string") {
317 const pending = pendingRequests.get(record["id"]);
318 if (pending !== undefined) {
319 pendingRequests.delete(record["id"]);
320 if (record["success"] === true) {
321 yield* Deferred.succeed(pending.deferred, record["data"]);
322 } else {
323 yield* Deferred.fail(
324 pending.deferred,
325 new PiRpcError({
326 operation: String(record["command"] ?? "request"),
327 detail: summarizePiError(record["error"]),
328 ...(record["error"] === undefined ? {} : { cause: record["error"] }),
329 }),
330 );
331 }
332 return;
333 }
334 }
335 yield* Queue.offer(events, record);
336 });
337
338 // Watch exit before the reader so an immediate crash can populate
339 // `exitDeferred` before stdout-close diagnosis runs.
340 yield* child.exitCode.pipe(
341 Effect.matchEffect({
342 onFailure: (cause) =>
343 Deferred.fail(exitDeferred, new PiRpcError({ operation: "exit", cause })),
344 onSuccess: (code) =>
345 Effect.suspend(() => {
346 childExited = true;
347 return Deferred.succeed(exitDeferred, Number(code));
348 }),
349 }),
350 Effect.forkIn(scope),
351 );
352
353 // Reader: decode stdout into LF-delimited JSON records.
354 yield* Effect.gen(function* () {
355 const frame = makeJsonlFramer();
356 yield* child.stdout.pipe(
357 Stream.decodeText(),
358 Stream.runForEach((chunk) =>
359 Effect.gen(function* () {
360 const lines = frame(chunk);
361 for (const line of lines) {
362 const record = parsePiRecord(line);
363 if (record === undefined) {
364 yield* Effect.logDebug("Dropping non-JSON pi stdout line.", {
365 lineLength: line.length,
366 });
367 continue;
368 }
369 yield* routeRecord(record);
370 }
371 }),
372 ),
373 );
374 for (const line of frame("\n")) {
375 const trailing = parsePiRecord(line);
376 if (trailing !== undefined) yield* routeRecord(trailing);
377 }
378 }).pipe(
379 Effect.matchCauseEffect({
380 onFailure: (cause) => failTransport(new PiRpcError({ operation: "read", cause })),
381 onSuccess: () =>
382 Effect.gen(function* () {
383 // Claim the stdout-close path before waiting so a broken stdin
384 // writer cannot replace the exit diagnosis with a write error.
385 diagnosingStdoutClose = true;
386 const polled = yield* Deferred.poll(exitDeferred);
387 const maybeExitCode = Option.isSome(polled)
388 ? yield* polled.value.pipe(
389 Effect.map(Option.some),
390 Effect.orElseSucceed(() => Option.none<number>()),
391 )
392 : yield* Deferred.await(exitDeferred).pipe(
393 Effect.timeoutOption(Duration.millis(250)),
394 // Adapter tests run under TestClock. A planned stdout close
395 // without an exit, for example Stop-with-restart, must not wait
396 // on that clock or the transport never fails.
397 Effect.provideService(Clock.Clock, Clock.Clock.defaultValue()),
398 Effect.orElseSucceed(() => Option.none<number>()),
399 );
400 return yield* failTransport(
401 new PiRpcError({
402 operation: "read",
403 detail: describePiStdoutClosure(
404 Option.isSome(maybeExitCode) ? maybeExitCode.value : undefined,
405 ),
406 }),
407 );
408 }),
409 }),
410 Effect.forkIn(scope),
411 );
412
413 // Surface stderr as debug logs; pi reserves stdout for the protocol.
414 yield* child.stderr.pipe(
415 Stream.decodeText(),
416 Stream.runForEach((chunk) =>
417 chunk.trim().length === 0
418 ? Effect.void
419 : // Length only: pi's stderr is unbounded remote output and can carry
420 // credentials or prompt text, so it never enters a log annotation.
421 Effect.logDebug("pi stderr", { stderrLength: chunk.length }),
422 ),
423 Effect.ignore,
424 Effect.forkIn(scope),
425 );
426
427 // Writer starts after the reader so an already-closed stdout can mark
428 // diagnosis before a broken-pipe stdin claims the transport.
429 yield* Stream.fromQueue(outgoing).pipe(
430 Stream.run(child.stdin),
431 Effect.catchCause((cause) =>
432 diagnosingStdoutClose
433 ? Effect.void
434 : failTransport(new PiRpcError({ operation: "write", cause })),
435 ),
436 Effect.forkIn(scope),
437 );
438
439 const send = (record: PiRpcRecord): Effect.Effect<void, PiRpcError> =>
440 Effect.gen(function* () {
441 const accepted = yield* Queue.offer(
442 outgoing,
443 new TextEncoder().encode(`${encodeJsonLine(record)}\n`),
444 );
445 // A refused offer means `failTransport` already closed the queue, so the
446 // write can never land; surface the transport error instead of
447 // reporting a success the caller cannot rely on.
448 if (!accepted) {
449 return yield* Deferred.await(transportDown);
450 }
451 });
452
453 const request = (
454 record: PiRpcRecord,
455 timeoutMs = DEFAULT_REQUEST_TIMEOUT_MS,
456 ): Effect.Effect<unknown, PiRpcError | PiRpcTimeoutError> =>
457 Effect.gen(function* () {
458 const id = `t3-${nextRequestId++}`;
459 const deferred = yield* Deferred.make<unknown, PiRpcError>();
460 pendingRequests.set(id, { deferred });
461 yield* send({ ...record, id }).pipe(
462 Effect.tapError(() => Effect.sync(() => pendingRequests.delete(id))),
463 );
464 // Raced against the transport: a death that lands after this request was
465 // registered (or between `send`'s check and its enqueue) would otherwise
466 // leave the caller waiting out the full timeout for a reply that is
467 // never coming.
468 return yield* Effect.raceFirst(Deferred.await(deferred), Deferred.await(transportDown)).pipe(
469 Effect.timeoutOrElse({
470 duration: Duration.millis(timeoutMs),
471 orElse: () =>
472 Effect.fail(
473 new PiRpcTimeoutError({
474 operation: String(record["type"] ?? "request"),
475 timeoutMs,
476 }),
477 ),
478 }),
479 Effect.onInterrupt(() => Effect.sync(() => pendingRequests.delete(id))),
480 Effect.onError(() => Effect.sync(() => pendingRequests.delete(id))),
481 );
482 });
483
484 return {
485 send,
486 request,
487 events,
488 exited: Deferred.await(exitDeferred),
489 terminate: failTransport(
490 new PiRpcError({ operation: "terminate", detail: "pi process was stopped" }),
491 ).pipe(Effect.andThen(terminateProcess), Effect.ignore, Effect.uninterruptible),
492 } satisfies PiRpcConnection;
493});