Skip to content

Commit 52bef09

Browse files
hafiezulhafiezul
committed
feat(pi): harden recovery and continuation safety (#6)
1 parent 5aa2b8a commit 52bef09

6 files changed

Lines changed: 448 additions & 66 deletions

File tree

‎apps/server/src/provider/Drivers/PiDriver.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { ServerSettingsService } from "../../serverSettings.ts";
1717
import * as TextGeneration from "../../textGeneration/TextGeneration.ts";
1818
import { ProviderDriverError } from "../Errors.ts";
1919
import { makePiAdapter, makePiImageAttachmentLoader } from "../Layers/PiAdapter.ts";
20+
import { ProviderEventLoggers } from "../Layers/ProviderEventLoggers.ts";
2021
import { makePiSessionRuntime } from "./PiSessionRuntime.ts";
2122
import { makeManagedServerProvider } from "../makeManagedServerProvider.ts";
2223
import {
@@ -79,6 +80,7 @@ export type PiDriverEnv =
7980
| FileSystem.FileSystem
8081
| Path.Path
8182
| ProcessRunner
83+
| ProviderEventLoggers
8284
| ServerConfig
8385
| ServerSettingsService;
8486

@@ -115,6 +117,7 @@ export const PiDriver: ProviderDriver<PiSettings, PiDriverEnv> = {
115117
Effect.gen(function* () {
116118
const serverConfig = yield* ServerConfig;
117119
const serverSettings = yield* ServerSettingsService;
120+
const eventLoggers = yield* ProviderEventLoggers;
118121
const processRunner = yield* ProcessRunner;
119122
const fileSystem = yield* FileSystem.FileSystem;
120123
const path = yield* Path.Path;
@@ -140,6 +143,7 @@ export const PiDriver: ProviderDriver<PiSettings, PiDriverEnv> = {
140143
instanceId,
141144
sessionDirectory,
142145
environment: processEnv,
146+
...(eventLoggers.native ? { nativeEventLogger: eventLoggers.native } : {}),
143147
loadImageAttachment: makePiImageAttachmentLoader({
144148
attachmentsDir: serverConfig.attachmentsDir,
145149
fileSystem,

‎apps/server/src/provider/Drivers/PiSessionRuntime.test.ts‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,10 @@ function handle(command) {
7979
respond(command);
8080
return;
8181
case "prompt":
82+
if (process.env.PI_TEST_EXIT_ON_PROMPT === "true") {
83+
process.stderr.write("simulated Pi process loss\\n");
84+
process.exit(23);
85+
}
8286
if (sessionDirectory) {
8387
mkdirSync(sessionDirectory, { recursive: true });
8488
writeFileSync(sessionFile, JSON.stringify({ type: "session", id: sessionId, message: command.message }) + "\\n");
@@ -238,3 +242,38 @@ it.effect("starts Pi in the configured project and extension configuration conte
238242
NodeFS.rmSync(configDirectory, { recursive: true, force: true });
239243
}).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
240244
);
245+
246+
it.effect("reports a diagnosable transport failure when Pi exits during an accepted prompt", () =>
247+
Effect.gen(function* () {
248+
const binaryPath = makeMockPiBinary();
249+
const runtime = yield* makePiSessionRuntime({
250+
binaryPath,
251+
configDirectory: "",
252+
launchArgs: "",
253+
cwd: process.cwd(),
254+
environment: { PI_TEST_EXIT_ON_PROMPT: "true" },
255+
});
256+
yield* runtime.start();
257+
const eventFiber = yield* Stream.take(runtime.events, 1).pipe(
258+
Stream.runCollect,
259+
Effect.forkChild,
260+
);
261+
yield* Effect.yieldNow;
262+
263+
const promptFailure = yield* runtime
264+
.prompt({ message: "Persist this prompt before the transport fails." })
265+
.pipe(Effect.flip);
266+
const events = Array.from(yield* Fiber.join(eventFiber).pipe(Effect.timeout("1 second")));
267+
268+
expect(["read-stdout", "process-exit"]).toContain(promptFailure.operation);
269+
expect(events).toEqual([
270+
expect.objectContaining({
271+
type: "pi_rpc_transport_failure",
272+
operation: expect.stringMatching(/read-stdout|process-exit/),
273+
detail: expect.stringContaining("Pi RPC"),
274+
}),
275+
]);
276+
277+
NodeFS.rmSync(NodePath.dirname(binaryPath), { recursive: true, force: true });
278+
}).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
279+
);

‎apps/server/src/provider/Drivers/PiSessionRuntime.ts‎

Lines changed: 70 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
import { resolveSpawnCommand } from "@t3tools/shared/shell";
2+
import * as Cause from "effect/Cause";
23
import * as Deferred from "effect/Deferred";
34
import * as Effect from "effect/Effect";
5+
import * as Exit from "effect/Exit";
46
import * as Option from "effect/Option";
57
import * as Queue from "effect/Queue";
68
import * as Ref from "effect/Ref";
@@ -35,6 +37,23 @@ export class PiSessionRuntimeError extends Schema.TaggedErrorClass<PiSessionRunt
3537

3638
const isPiSessionRuntimeError = Schema.is(PiSessionRuntimeError);
3739

40+
/** Synthetic raw event emitted when Pi's process or RPC transport becomes unusable. */
41+
export const PI_RPC_TRANSPORT_FAILURE_EVENT_TYPE = "pi_rpc_transport_failure";
42+
43+
export interface PiRpcTransportFailureEvent {
44+
readonly type: typeof PI_RPC_TRANSPORT_FAILURE_EVENT_TYPE;
45+
readonly operation: string;
46+
readonly detail: string;
47+
}
48+
49+
export function isPiSessionRuntimeTransportError(error: PiSessionRuntimeError): boolean {
50+
return (
51+
error.operation === "read-stdout" ||
52+
error.operation === "process-exit" ||
53+
error.operation === "write-stdin"
54+
);
55+
}
56+
3857
export interface PiSessionRuntimeOptions {
3958
readonly binaryPath: string;
4059
readonly configDirectory: string;
@@ -160,6 +179,15 @@ function stringValue(value: unknown): string | undefined {
160179
return typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
161180
}
162181

182+
export function isPiRpcTransportFailureEvent(value: unknown): value is PiRpcTransportFailureEvent {
183+
return (
184+
isRecord(value) &&
185+
value.type === PI_RPC_TRANSPORT_FAILURE_EVENT_TYPE &&
186+
stringValue(value.operation) !== undefined &&
187+
stringValue(value.detail) !== undefined
188+
);
189+
}
190+
163191
function parseModel(value: unknown): PiRpcModel | undefined {
164192
if (!isRecord(value)) {
165193
return undefined;
@@ -288,7 +316,7 @@ export const makePiSessionRuntime = (
288316
const pendingRequests = yield* Ref.make(new Map<string, PendingRequest>());
289317
const requestNumber = yield* Ref.make(0);
290318
const writeLock = yield* Semaphore.make(1);
291-
const rawEvents = yield* Queue.unbounded<unknown>();
319+
const rawEvents = yield* Queue.unbounded<unknown, Cause.Done>();
292320
const closed = yield* Ref.make(false);
293321
const decoder = makePiJsonlDecoder();
294322

@@ -303,6 +331,26 @@ export const makePiSessionRuntime = (
303331
),
304332
);
305333

334+
const finishUnexpectedTransportFailure = Effect.fn(
335+
"PiSessionRuntime.finishUnexpectedTransportFailure",
336+
)(function* (error: PiSessionRuntimeError) {
337+
if (yield* Ref.getAndSet(closed, true)) {
338+
return;
339+
}
340+
341+
yield* Queue.offer(rawEvents, {
342+
type: PI_RPC_TRANSPORT_FAILURE_EVENT_TYPE,
343+
operation: error.operation,
344+
detail: error.detail,
345+
} satisfies PiRpcTransportFailureEvent);
346+
yield* failPendingRequests(error);
347+
yield* child.kill({ forceKillAfter: PI_RPC_FORCE_KILL_AFTER }).pipe(Effect.ignore);
348+
// `end` drains the failure event before completing the event stream.
349+
// `shutdown` would interrupt the adapter before it could mark an active
350+
// T3 turn interrupted and retain the diagnostic payload.
351+
yield* Queue.end(rawEvents);
352+
});
353+
306354
const removePendingRequest = (id: string) =>
307355
Ref.update(pendingRequests, (pending) => {
308356
if (!pending.has(id)) {
@@ -370,12 +418,19 @@ export const makePiSessionRuntime = (
370418
Effect.forEach(decoder.push(chunk), handleRecord, { discard: true }).pipe(Effect.asVoid),
371419
),
372420
Effect.ensuring(flushDecoder()),
373-
Effect.catch(() =>
374-
failPendingRequests(
375-
new PiSessionRuntimeError({
376-
operation: "read-stdout",
377-
detail: "Pi RPC stdout closed unexpectedly.",
378-
}),
421+
Effect.exit,
422+
Effect.flatMap((exit) =>
423+
finishUnexpectedTransportFailure(
424+
Exit.isSuccess(exit)
425+
? new PiSessionRuntimeError({
426+
operation: "read-stdout",
427+
detail: "Pi RPC stdout closed unexpectedly.",
428+
})
429+
: new PiSessionRuntimeError({
430+
operation: "read-stdout",
431+
detail: "Pi RPC stdout connection failed.",
432+
cause: exit.cause,
433+
}),
379434
),
380435
),
381436
Effect.forkIn(runtimeScope),
@@ -388,15 +443,15 @@ export const makePiSessionRuntime = (
388443
yield* child.exitCode.pipe(
389444
Effect.matchEffect({
390445
onFailure: (cause) =>
391-
failPendingRequests(
446+
finishUnexpectedTransportFailure(
392447
new PiSessionRuntimeError({
393448
operation: "process-exit",
394449
detail: "Could not read Pi process exit status.",
395450
cause,
396451
}),
397452
),
398453
onSuccess: (code) =>
399-
failPendingRequests(
454+
finishUnexpectedTransportFailure(
400455
new PiSessionRuntimeError({
401456
operation: "process-exit",
402457
detail:
@@ -406,7 +461,6 @@ export const makePiSessionRuntime = (
406461
}),
407462
),
408463
}),
409-
Effect.andThen(Queue.shutdown(rawEvents)),
410464
Effect.forkIn(runtimeScope),
411465
);
412466

@@ -434,6 +488,9 @@ export const makePiSessionRuntime = (
434488
cause,
435489
}),
436490
),
491+
Effect.tapError((error) =>
492+
error.operation === "write-stdin" ? finishUnexpectedTransportFailure(error) : Effect.void,
493+
),
437494
);
438495

439496
const request = Effect.fn("PiSessionRuntime.request")(function* (
@@ -480,7 +537,9 @@ export const makePiSessionRuntime = (
480537
return response;
481538
});
482539

483-
const getState = (timeout = PI_RPC_REQUEST_TIMEOUT) =>
540+
const getState = (
541+
timeout: typeof PI_RPC_REQUEST_TIMEOUT | typeof PI_RPC_START_TIMEOUT = PI_RPC_REQUEST_TIMEOUT,
542+
) =>
484543
request({ type: "get_state" }, timeout).pipe(
485544
Effect.flatMap((response) => {
486545
const state = parseState(response);

‎apps/server/src/provider/Layers/PiAdapter.test.ts‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1109,4 +1109,123 @@ describe("PiAdapter", () => {
11091109
expect(yield* second.hasSession(THREAD_ID)).toBe(false);
11101110
}).pipe(Effect.scoped),
11111111
);
1112+
1113+
it.effect(
1114+
"interrupts an accepted turn, closes Pi, and retains its transport failure in native logs",
1115+
() =>
1116+
Effect.gen(function* () {
1117+
const runtime = yield* makeRuntimeFactory();
1118+
const nativeEvents: Array<{
1119+
readonly event?: {
1120+
readonly method?: string;
1121+
readonly payload?: unknown;
1122+
readonly provider?: string;
1123+
readonly providerThreadId?: string;
1124+
readonly threadId?: string;
1125+
};
1126+
}> = [];
1127+
const nativeThreadIds: Array<ThreadId | null> = [];
1128+
const adapter = yield* makePiAdapter(decodePiSettings({}), {
1129+
instanceId: INSTANCE_A,
1130+
sessionDirectory: "/tmp/t3-pi-sessions/pi_personal",
1131+
nativeEventLogger: {
1132+
filePath: "memory://pi-native-events",
1133+
write: (event, threadId) =>
1134+
Effect.sync(() => {
1135+
nativeEvents.push(event as (typeof nativeEvents)[number]);
1136+
nativeThreadIds.push(threadId);
1137+
}),
1138+
close: () => Effect.void,
1139+
},
1140+
makeRuntime: runtime.factory,
1141+
});
1142+
const events: ProviderRuntimeEvent[] = [];
1143+
yield* adapter.streamEvents.pipe(
1144+
Stream.runForEach((event) =>
1145+
Effect.sync(() => {
1146+
events.push(event);
1147+
}),
1148+
),
1149+
Effect.forkScoped,
1150+
);
1151+
yield* adapter.startSession(sessionStart(INSTANCE_A));
1152+
yield* Effect.yieldNow;
1153+
const turn = yield* adapter.sendTurn({
1154+
threadId: THREAD_ID,
1155+
input: "Persist this accepted prompt exactly once.",
1156+
attachments: [],
1157+
});
1158+
1159+
yield* runtime.emit({
1160+
type: "pi_rpc_transport_failure",
1161+
operation: "process-exit",
1162+
detail: "Pi RPC process exited with code 23.",
1163+
});
1164+
yield* Effect.yieldNow;
1165+
yield* Effect.yieldNow;
1166+
yield* Effect.yieldNow;
1167+
1168+
expect(events).toEqual(
1169+
expect.arrayContaining([
1170+
expect.objectContaining({
1171+
type: "turn.completed",
1172+
turnId: turn.turnId,
1173+
payload: expect.objectContaining({
1174+
state: "interrupted",
1175+
errorMessage: "Pi RPC process exited with code 23.",
1176+
}),
1177+
}),
1178+
expect.objectContaining({
1179+
type: "runtime.error",
1180+
turnId: turn.turnId,
1181+
payload: expect.objectContaining({
1182+
class: "transport_error",
1183+
message: "Pi RPC process exited with code 23.",
1184+
}),
1185+
}),
1186+
expect.objectContaining({
1187+
type: "session.exited",
1188+
payload: {
1189+
reason: "Pi RPC process exited with code 23.",
1190+
recoverable: true,
1191+
exitKind: "error",
1192+
},
1193+
}),
1194+
]),
1195+
);
1196+
expect(runtime.getCloseCount()).toBe(1);
1197+
expect(yield* adapter.hasSession(THREAD_ID)).toBe(false);
1198+
expect(yield* adapter.listSessions()).toEqual([]);
1199+
1200+
const replayAttempt = yield* adapter
1201+
.sendTurn({
1202+
threadId: THREAD_ID,
1203+
input: "This must not be replayed automatically.",
1204+
attachments: [],
1205+
})
1206+
.pipe(Effect.flip);
1207+
expect(replayAttempt._tag).toBe("ProviderAdapterSessionNotFoundError");
1208+
expect(runtime.prompts).toEqual([
1209+
{ message: "Persist this accepted prompt exactly once." },
1210+
]);
1211+
expect(nativeEvents).toEqual(
1212+
expect.arrayContaining([
1213+
expect.objectContaining({
1214+
event: expect.objectContaining({
1215+
provider: "pi",
1216+
providerThreadId: THREAD_ID,
1217+
threadId: THREAD_ID,
1218+
method: "pi.rpc.transport.process-exit",
1219+
payload: {
1220+
type: "pi_rpc_transport_failure",
1221+
operation: "process-exit",
1222+
detail: "Pi RPC process exited with code 23.",
1223+
},
1224+
}),
1225+
}),
1226+
]),
1227+
);
1228+
expect(nativeThreadIds).toEqual([THREAD_ID]);
1229+
}).pipe(Effect.scoped),
1230+
);
11121231
});

0 commit comments

Comments
 (0)