1import * as Effect from "effect/Effect";
2import * as Option from "effect/Option";
3import * as Queue from "effect/Queue";
4import * as Schema from "effect/Schema";
5import * as ChildProcessSpawner from "effect/process/ChildProcessSpawner";
6
7import { TextGenerationError } from "@t3tools/contracts";
8import type { PiSettings } from "../settings.ts";
9
10import { makePiRpcConnection, piRecordField, piRecordString, resolvePiModelSlug } from "./rpc.ts";
11import { buildPiRpcLaunch, resolvePiLaunchArgs } from "./mcpInjection.ts";
12import * as TextGenerationOperations from "@t3tools/provider-core/server/textGenerationOperations";
13
14const PI_TIMEOUT_MS = 180_000;
15
16const isTextGenerationError = Schema.is(TextGenerationError);
17
18// Pi's `errorMessage` is the provider's raw response text and can echo
19// credentials, so only a leading HTTP status reaches the error detail.
20const piTurnErrorDetail = (errorMessage: string | undefined) => {
21 const status = errorMessage?.match(/^([45]\d\d)\b/)?.[1];
22 return status === undefined
23 ? "Pi could not generate text: the model provider reported an error."
24 : `Pi could not generate text: the model provider returned HTTP ${status}.`;
25};
26
27export const makePiTextGeneration = Effect.fn("makePiTextGeneration")(function* (
28 piSettings: PiSettings,
29 environment: NodeJS.ProcessEnv = process.env,
30) {
31 const spawner = yield* ChildProcessSpawner.ChildProcessSpawner;
32
33 const runPiJson: TextGenerationOperations.Runner = Effect.fn("PiTextGeneration.runPiJson")(
34 function* (request) {
35 const { operation, cwd, prompt, modelSelection } = request;
36
37 const resolvedLaunchArgs = resolvePiLaunchArgs(piSettings.launchArgs);
38 if (!resolvedLaunchArgs.ok) {
39 return yield* new TextGenerationError({
40 operation,
41 detail: resolvedLaunchArgs.message,
42 });
43 }
44 const launch = buildPiRpcLaunch({
45 launchArgs: resolvedLaunchArgs.args,
46 environment,
47 mcpSession: undefined,
48 extensionPath: undefined,
49 ephemeral: true,
50 shortLived: true,
51 disableExtensions: true,
52 // Background naming/content helpers must never mutate the workspace.
53 disableTools: true,
54 });
55 const connection = yield* makePiRpcConnection({
56 command: piSettings.binaryPath || "pi",
57 args: launch.args,
58 cwd,
59 env: launch.env,
60 }).pipe(Effect.provideService(ChildProcessSpawner.ChildProcessSpawner, spawner));
61
62 if (modelSelection.model !== "default") {
63 // Reject unusable custom slugs; silently using Pi's default would misreport the requested model.
64 const inventory = yield* connection
65 .request({ type: "get_available_models" })
66 .pipe(Effect.orElseSucceed(() => undefined));
67 const parsed = resolvePiModelSlug(modelSelection.model, piRecordField(inventory, "models"));
68 if (parsed === null) {
69 return yield* new TextGenerationError({
70 operation,
71 detail: `Pi model '${modelSelection.model}' must use provider/model format.`,
72 });
73 }
74 yield* connection
75 .request({ type: "set_model", provider: parsed.provider, modelId: parsed.modelId })
76 .pipe(
77 // Model discovery loads extensions that this background helper disables.
78 Effect.mapError((error) =>
79 error._tag === "PiRpcError" && error.operation === "set_model"
80 ? new TextGenerationError({
81 operation,
82 detail:
83 `Pi rejected model '${modelSelection.model}': ${error.detail ?? "unknown error"}. ` +
84 "If the model comes from a Pi extension or package, text generation cannot use it " +
85 "because it runs Pi with --no-extensions. Pick a built-in Pi model or one defined " +
86 "in Pi's models.json.",
87 cause: error,
88 })
89 : error,
90 ),
91 );
92 }
93
94 yield* connection.request({ type: "prompt", message: prompt });
95 // Provider failures (bad key, rate limit, 5xx) arrive on the assistant
96 // `message_end`. Keep only the latest, since Pi can retry and still answer.
97 let turnError: string | undefined;
98 yield* Effect.gen(function* () {
99 while (true) {
100 const event = yield* Queue.take(connection.events);
101 if (event["type"] === "agent_settled") return;
102 const message = event["message"];
103 if (event["type"] === "message_end" && piRecordString(message, "role") === "assistant") {
104 turnError =
105 piRecordString(message, "stopReason") === "error"
106 ? piTurnErrorDetail(piRecordString(message, "errorMessage"))
107 : undefined;
108 }
109 }
110 });
111 const data = yield* connection.request({ type: "get_last_assistant_text" });
112 const text =
113 typeof data === "object" &&
114 data !== null &&
115 typeof (data as { text?: unknown }).text === "string"
116 ? (data as { text: string }).text.trim()
117 : "";
118 if (!text) {
119 return yield* new TextGenerationError({
120 operation,
121 detail: turnError ?? "Pi returned empty output.",
122 });
123 }
124 return yield* TextGenerationOperations.decodeJsonReply(request, "Pi", text);
125 },
126 (effect, request) => {
127 const { operation } = request;
128 return effect.pipe(
129 Effect.timeoutOption(PI_TIMEOUT_MS),
130 Effect.flatMap(
131 Option.match({
132 onNone: () =>
133 Effect.fail(new TextGenerationError({ operation, detail: "Pi request timed out." })),
134 onSome: (value) => Effect.succeed(value),
135 }),
136 ),
137 Effect.mapError((cause) =>
138 isTextGenerationError(cause)
139 ? cause
140 : new TextGenerationError({
141 operation,
142 detail: "Pi text generation failed.",
143 cause,
144 }),
145 ),
146 Effect.scoped,
147 );
148 },
149 );
150
151 return TextGenerationOperations.fromRunner(
152 "PiTextGeneration",
153 TextGenerationOperations.withoutInvestigation("Pi", runPiJson),
154 );
155});