Open History to see the last commit
1// Pi frames are LF-delimited with trailing CR stripped; readline also splits U+2028/U+2029 and corrupts them.
2import * as Clock from "effect/Clock";
3import * as Deferred from "effect/Deferred";
4import * as Duration from "effect/Duration";
5import * as Effect from "effect/Effect";
6import * as Option from "effect/Option";
7import * as Predicate from "effect/Predicate";
8import * as Queue from "effect/Queue";
9import * as Scope from "effect/Scope";
10import * as Schema from "effect/Schema";
11import * as Stream from "effect/Stream";
12import * as ChildProcess from "effect/process/ChildProcess";
13import * as ChildProcessSpawner from "effect/process/ChildProcessSpawner";
14import * as HostProcess from "@t3tools/shared/HostProcess";
15import { resolveSpawnCommand } from "@t3tools/shared/shell";
16import {
17 MAX_TOOL_OUTPUT_IMAGES,
18 MAX_TOOL_OUTPUT_IMAGE_BASE64_LENGTH,
19} from "@t3tools/shared/toolOutput";
20
21import { signalProcessGroup } from "@t3tools/provider-core/server/processGroup";
22
23export class PiRpcError extends Schema.TaggedError<PiRpcError>()("PiRpcError", {
24 operation: Schema.String,
25 detail: Schema.optional(Schema.String),
26 cause: Schema.optional(Schema.Defect()),
27}) {
28 override get message(): string {
29 return `Pi RPC ${this.operation} failed${this.detail === undefined ? "" : `: ${this.detail}`}.`;
30 }
31}
32
33export class PiRpcTimeoutError extends Schema.TaggedError<PiRpcTimeoutError>()(
34 "PiRpcTimeoutError",
35 {
36 operation: Schema.String,
37 timeoutMs: Schema.Int.check(Schema.isBetween({ minimum: 0, maximum: Number.MAX_SAFE_INTEGER })),
38 },
39) {
40 override get message(): string {
41 return `Pi RPC ${this.operation} failed: timed out after ${this.timeoutMs}ms.`;
42 }
43}
44
45export type PiRpcRecord = Record<string, unknown>;
46
47export function piRecordField(input: unknown, key: string): unknown {
48 return Predicate.isObject(input) ? input[key] : undefined;
49}
50
51export function piRecordString(input: unknown, key: string): string | undefined {
52 const value = piRecordField(input, key);
53 return Predicate.isString(value) ? value : undefined;
54}
55
56export function piRecordNumber(input: unknown, key: string): number | undefined {
57 const value = piRecordField(input, key);
58 return Predicate.isNumber(value) && Number.isFinite(value) ? value : undefined;
59}
60
61// Pi provider names and model ids may both contain "/", so a slug only splits
62// reliably through the `{provider, id}` pairs `get_available_models` reported.
63export function findReportedPiModel(
64 slug: string,
65 reportedModels: unknown,
66): { provider: string; modelId: string } | undefined {
67 if (!Array.isArray(reportedModels)) return undefined;
68 for (const model of reportedModels) {
69 const provider = piRecordString(model, "provider");
70 const modelId = piRecordString(model, "id");
71 if (provider !== undefined && modelId !== undefined && `${provider}/${modelId}` === slug)
72 return { provider, modelId };
73 }
74 return undefined;
75}
76
77export function resolvePiModelSlug(
78 slug: string,
79 reportedModels: unknown,
80): { provider: string; modelId: string } | null {
81 const reported = findReportedPiModel(slug, reportedModels);
82 if (reported !== undefined) return reported;
83 const separator = slug.indexOf("/");
84 if (separator <= 0 || separator === slug.length - 1) return null;
85 return { provider: slug.slice(0, separator), modelId: slug.slice(separator + 1) };
86}
87
88export interface PiRpcSpawnOptions {
89 readonly command: string;
90 readonly args: ReadonlyArray<string>;
91 readonly cwd: string | undefined;
92 readonly env: NodeJS.ProcessEnv;
93}
94
95export interface PiRpcConnection {
96 readonly send: (record: PiRpcRecord) => Effect.Effect<void, PiRpcError>;
97 readonly request: (
98 record: PiRpcRecord,
99 timeoutMs?: number,
100 ) => Effect.Effect<unknown, PiRpcError | PiRpcTimeoutError>;
101 readonly events: Queue.Queue<PiRpcRecord, PiRpcError>;
102 readonly exited: Effect.Effect<number, PiRpcError>;
103 readonly terminate: Effect.Effect<void>;
104}
105
106const DEFAULT_REQUEST_TIMEOUT_MS = 30_000;
107const TERMINATION_GRACE = Duration.seconds(1);
108
109interface PendingPiRequest {
110 readonly deferred: Deferred.Deferred<unknown, PiRpcError>;
111}
112
113// MCP duplicates images in model content and CallToolResult; budget both plus bounded overhead.
114const MAX_PI_RECORD_CHARS =
115 2 * MAX_TOOL_OUTPUT_IMAGES * MAX_TOOL_OUTPUT_IMAGE_BASE64_LENGTH + 8 * 1024 * 1024;
116
117function makeJsonlFramer() {
118 let buffer = "";
119 let dropping = false;
120 return (chunk: string) => {
121 const lines: string[] = [];
122 const droppedLengths: number[] = [];
123 let start = 0;
124 while (start < chunk.length) {
125 const newline = chunk.indexOf("\n", start);
126 const end = newline < 0 ? chunk.length : newline;
127 if (!dropping) {
128 if (buffer.length + end - start > MAX_PI_RECORD_CHARS) {
129 droppedLengths.push(buffer.length + end - start);
130 buffer = "";
131 dropping = true;
132 } else {
133 buffer += chunk.slice(start, end);
134 }
135 }
136 if (newline < 0) break;
137 if (!dropping && buffer.length > 0) {
138 lines.push(buffer.endsWith("\r") ? buffer.slice(0, -1) : buffer);
139 }
140 buffer = "";
141 dropping = false;
142 start = newline + 1;
143 }
144 return { lines, droppedLengths };
145 };
146}
147
148// Report the observed length before LF arrives; an unterminated response can time out first.
149const logDroppedPiRecords = (droppedLengths: ReadonlyArray<number>) =>
150 Effect.forEach(
151 droppedLengths,
152 (recordLength) =>
153 Effect.logWarning("Dropping oversized pi stdout record.", {
154 recordLength,
155 maxRecordLength: MAX_PI_RECORD_CHARS,
156 }),
157 { discard: true },
158 );
159
160const UnknownFromJsonString = Schema.fromJsonString(Schema.Unknown);
161const decodeJsonLine = Schema.decodeSync(UnknownFromJsonString);
162const encodeJsonLine = Schema.encodeSync(UnknownFromJsonString);
163
164const PI_ERROR_DETAIL_MAX_CHARS = 200;
165
166function summarizePiError(error: unknown): string {
167 const text = typeof error === "string" ? error : JSON.stringify(error);
168 if (text === undefined) return "unknown error";
169 return text.length > PI_ERROR_DETAIL_MAX_CHARS
170 ? `${text.slice(0, PI_ERROR_DETAIL_MAX_CHARS)}…`
171 : text;
172}
173
174// Pi 0.84+ may exit 1 during import without any protocol output.
175function describePiStdoutClosure(exitCode: number | undefined): string {
176 if (exitCode === undefined || !Number.isFinite(exitCode)) {
177 return "pi process closed stdout";
178 }
179 return `pi process exited with code ${exitCode}`;
180}
181
182function parsePiRecord(line: string): PiRpcRecord | undefined {
183 try {
184 const parsed: unknown = decodeJsonLine(line);
185 return Predicate.isObject(parsed) ? parsed : undefined;
186 } catch {
187 return undefined;
188 }
189}
190
191/**
192 * `hasExited` is consulted before each signal. Once the original child is
193 * gone its pid/pgid can be recycled by the OS, so escalating blindly could
194 * deliver SIGKILL to an unrelated process.
195 */
196const terminatePiProcess = Effect.fn("PiRpc.terminatePiProcess")(function* (
197 kill: (signal: NodeJS.Signals) => boolean,
198 hasExited: () => boolean,
199) {
200 if (hasExited()) return;
201 if (!kill("SIGTERM")) return;
202 yield* Effect.sleep(TERMINATION_GRACE);
203 if (hasExited()) return;
204 kill("SIGKILL");
205});
206
207export const makePiRpcConnection = Effect.fnUntraced(function* (options: PiRpcSpawnOptions) {
208 const spawner = yield* ChildProcessSpawner.ChildProcessSpawner;
209 const platform = yield* HostProcess.Platform;
210 const scope = yield* Effect.scope;
211
212 const spawnCommand = yield* resolveSpawnCommand(options.command, [...options.args], {
213 env: options.env,
214 }).pipe(Effect.mapError((cause) => new PiRpcError({ operation: "spawn", cause })));
215 const child = yield* spawner
216 .spawn(
217 ChildProcess.make(spawnCommand.command, spawnCommand.args, {
218 ...(options.cwd === undefined ? {} : { cwd: options.cwd }),
219 env: options.env,
220 extendEnv: false,
221 shell: spawnCommand.shell,
222 detached: platform !== "win32",
223 // On stdin EOF Pi quits without its SIGTERM handler, so the later SIGTERM kills it mid-cleanup.
224 stdin: { stream: "pipe", endOnDone: false },
225 }),
226 )
227 .pipe(Effect.mapError((cause) => new PiRpcError({ operation: "spawn", cause })));
228
229 let childExited = false;
230 let diagnosingStdoutClose = false;
231
232 const killProcessGroup = (signal: NodeJS.Signals): boolean => {
233 try {
234 if (platform === "win32") {
235 process.kill(Number(child.pid), signal);
236 } else {
237 signalProcessGroup(Number(child.pid), signal);
238 }
239 return true;
240 } catch {
241 return false;
242 }
243 };
244
245 const hasExited = (): boolean => {
246 if (childExited) return true;
247 try {
248 if (platform === "win32") process.kill(Number(child.pid), 0);
249 else signalProcessGroup(Number(child.pid), 0);
250 return false;
251 } catch {
252 return true;
253 }
254 };
255
256 // Windows process.kill leaves extension children and pipes alive; taskkill /T reaps the tree.
257 const terminateWindowsTree = Effect.gen(function* () {
258 if (hasExited()) return;
259 const taskkill = yield* spawner.spawn(
260 ChildProcess.make("taskkill", ["/PID", String(child.pid), "/T", "/F"]),
261 );
262 yield* taskkill.exitCode;
263 }).pipe(Effect.scoped, Effect.ignore);
264
265 const terminateProcess =
266 platform === "win32" ? terminateWindowsTree : terminatePiProcess(killProcessGroup, hasExited);
267
268 // Register teardown immediately after spawn so interrupted setup cannot leak a detached process.
269 yield* Scope.addFinalizer(scope, terminateProcess.pipe(Effect.ignore, Effect.uninterruptible));
270
271 const pendingRequests = new Map<string, PendingPiRequest>();
272 const events = yield* Queue.unbounded<PiRpcRecord, PiRpcError>();
273 const outgoing = yield* Queue.unbounded<Uint8Array, PiRpcError>();
274 const transportDown = yield* Deferred.make<never, PiRpcError>();
275 const exitDeferred = yield* Deferred.make<number, PiRpcError>();
276 let nextRequestId = 0;
277
278 const failTransport = Effect.fn("PiRpc.failTransport")(function* (error: PiRpcError) {
279 const claimed = yield* Deferred.fail(transportDown, error);
280 if (!claimed) return;
281 for (const [key, pending] of pendingRequests) {
282 pendingRequests.delete(key);
283 yield* Deferred.fail(pending.deferred, error);
284 }
285 yield* Queue.fail(outgoing, error);
286 yield* Queue.fail(events, error);
287 });
288
289 const routeRecord =
290 /* Protocol records and writes include streaming deltas; avoid a span per record. */
291 Effect.fnUntraced(function* (record: PiRpcRecord) {
292 if (record["type"] === "response" && typeof record["id"] === "string") {
293 const pending = pendingRequests.get(record["id"]);
294 if (pending !== undefined) {
295 pendingRequests.delete(record["id"]);
296 if (record["success"] === true) {
297 yield* Deferred.succeed(pending.deferred, record["data"]);
298 } else {
299 yield* Deferred.fail(
300 pending.deferred,
301 new PiRpcError({
302 operation: String(record["command"] ?? "request"),
303 detail: summarizePiError(record["error"]),
304 ...(record["error"] === undefined ? {} : { cause: record["error"] }),
305 }),
306 );
307 }
308 return;
309 }
310 }
311 yield* Queue.offer(events, record);
312 });
313
314 // Watch exit before reading stdout so immediate crashes populate exitDeferred for diagnosis.
315 yield* child.exitCode.pipe(
316 Effect.matchEffect({
317 onFailure: (cause) =>
318 Deferred.fail(exitDeferred, new PiRpcError({ operation: "exit", cause })),
319 onSuccess: (code) =>
320 Effect.suspend(() => {
321 childExited = true;
322 return Deferred.succeed(exitDeferred, Number(code));
323 }),
324 }),
325 Effect.forkIn(scope),
326 );
327
328 yield* Effect.gen(function* () {
329 const frame = makeJsonlFramer();
330 yield* child.stdout.pipe(
331 Stream.decodeText(),
332 Stream.runForEach((chunk) =>
333 Effect.gen(function* () {
334 const { lines, droppedLengths } = frame(chunk);
335 yield* logDroppedPiRecords(droppedLengths);
336 for (const line of lines) {
337 const record = parsePiRecord(line);
338 if (record === undefined) {
339 yield* Effect.logDebug("Dropping non-JSON pi stdout line.", {
340 lineLength: line.length,
341 });
342 continue;
343 }
344 yield* routeRecord(record);
345 }
346 }),
347 ),
348 );
349 const flushed = frame("\n");
350 yield* logDroppedPiRecords(flushed.droppedLengths);
351 for (const line of flushed.lines) {
352 const trailing = parsePiRecord(line);
353 if (trailing !== undefined) yield* routeRecord(trailing);
354 }
355 }).pipe(
356 Effect.matchCauseEffect({
357 onFailure: (cause) => failTransport(new PiRpcError({ operation: "read", cause })),
358 // The enclosing operation owns the span for this inline handler.
359 onSuccess: Effect.fnUntraced(function* () {
360 // Claim stdout-close diagnosis before waiting so broken stdin cannot replace the exit error.
361 diagnosingStdoutClose = true;
362 const polled = yield* Deferred.poll(exitDeferred);
363 const maybeExitCode = Option.isSome(polled)
364 ? yield* polled.value.pipe(
365 Effect.asSome,
366 Effect.orElseSucceed(() => Option.none<number>()),
367 )
368 : yield* Deferred.await(exitDeferred).pipe(
369 Effect.timeoutOption(Duration.millis(250)),
370 // Stdout-close waits need wall time; TestClock would strand Stop-with-restart without an exit.
371 Effect.provideService(Clock.Clock, Clock.Clock.defaultValue()),
372 Effect.orElseSucceed(() => Option.none<number>()),
373 );
374 return yield* failTransport(
375 new PiRpcError({
376 operation: "read",
377 detail: describePiStdoutClosure(
378 Option.isSome(maybeExitCode) ? maybeExitCode.value : undefined,
379 ),
380 }),
381 );
382 }),
383 }),
384 Effect.forkIn(scope),
385 );
386
387 yield* child.stderr.pipe(
388 Stream.decodeText(),
389 Stream.runForEach((chunk) =>
390 chunk.trim().length === 0
391 ? Effect.void
392 : // Log stderr length only; unbounded text may contain credentials or prompts.
393 Effect.logDebug("pi stderr", { stderrLength: chunk.length }),
394 ),
395 Effect.ignore,
396 Effect.forkIn(scope),
397 );
398
399 // Start the reader first so stdout-close diagnosis can win over broken-pipe stdin.
400 yield* Stream.fromQueue(outgoing).pipe(
401 Stream.run(child.stdin),
402 Effect.catchCause((cause) =>
403 diagnosingStdoutClose
404 ? Effect.void
405 : failTransport(new PiRpcError({ operation: "write", cause })),
406 ),
407 Effect.forkIn(scope),
408 );
409
410 const send =
411 /* Protocol records and writes include streaming deltas; avoid a span per record. */
412 Effect.fnUntraced(function* (record: PiRpcRecord): Effect.fn.Return<void, PiRpcError> {
413 const accepted = yield* Queue.offer(
414 outgoing,
415 new TextEncoder().encode(`${encodeJsonLine(record)}\n`),
416 );
417 if (!accepted) {
418 return yield* Deferred.await(transportDown);
419 }
420 });
421
422 const request = Effect.fn("PiRpc.request")(function* (
423 record: PiRpcRecord,
424 timeoutMs = DEFAULT_REQUEST_TIMEOUT_MS,
425 ): Effect.fn.Return<unknown, PiRpcError | PiRpcTimeoutError> {
426 const id = `t3-${nextRequestId++}`;
427 const deferred = yield* Deferred.make<unknown, PiRpcError>();
428 pendingRequests.set(id, { deferred });
429 yield* send({ ...record, id }).pipe(
430 Effect.tapError(() => Effect.sync(() => pendingRequests.delete(id))),
431 );
432 return yield* Effect.raceFirst(Deferred.await(deferred), Deferred.await(transportDown)).pipe(
433 Effect.timeoutOrElse({
434 duration: Duration.millis(timeoutMs),
435 orElse: () =>
436 Effect.fail(
437 new PiRpcTimeoutError({
438 operation: String(record["type"] ?? "request"),
439 timeoutMs,
440 }),
441 ),
442 }),
443 Effect.onInterrupt(() => Effect.sync(() => pendingRequests.delete(id))),
444 Effect.onError(() => Effect.sync(() => pendingRequests.delete(id))),
445 );
446 });
447
448 return {
449 send,
450 request,
451 events,
452 exited: Deferred.await(exitDeferred),
453 terminate: failTransport(
454 new PiRpcError({ operation: "terminate", detail: "pi process was stopped" }),
455 ).pipe(Effect.andThen(terminateProcess), Effect.ignore, Effect.uninterruptible),
456 } satisfies PiRpcConnection;
457});