Skip to content
Source
Sign inSign up

mcpInjection.ts

324 lines9.9 KB

1import { tokenizeCliArgs } from "@t3tools/shared/cliArgs";
2import * as Effect from "effect/Effect";
3import * as FileSystem from "effect/FileSystem";
4​
5import type { McpProviderSessionConfig } from "@t3tools/provider-core/server/mcpSession";
6import {
7 PI_T3_MCP_EXTENSION_FILENAME,
8 PI_T3_MCP_EXTENSION_SOURCE,
9 T3_MCP_BEARER_ENV,
10 T3_MCP_URL_ENV,
11 T3_PI_RUNTIME_MODE_ENV,
12 T3_PI_MCP_EXTENSION_PATH_ENV,
13} from "./mcpExtensionSource.ts";
14​
15const RESERVED_PI_LAUNCH_ARGUMENTS = new Set([
16 "--continue",
17 "-c",
18 "--export",
19 "--fork",
20 "--help",
21 "-h",
22 "--list-models",
23 "--mode",
24 "--no-session",
25 "--print",
26 "-p",
27 "--resume",
28 "-r",
29 "--session",
30 "--session-id",
31 "--version",
32 "-v",
33]);
34​
35export const PI_ARGUMENTS_WITH_VALUES = new Set([
36 "--api-key",
37 "--append-system-prompt",
38 "--exclude-tools",
39 "-xt",
40 "--extension",
41 "-e",
42 "--model",
43 "--models",
44 "--name",
45 "-n",
46 "--prompt-template",
47 "--provider",
48 "--session-dir",
49 "--skill",
50 "--system-prompt",
51 "--theme",
52 "--thinking",
53 "--tools",
54 "-t",
55 "--tui-mode",
56 "--use-theme",
57]);
58​
59const PI_ARGUMENTS_WITHOUT_VALUES = new Set([
60 "--approve",
61 "-a",
62 "--no-approve",
63 "-na",
64 "--no-builtin-tools",
65 "-nbt",
66 "--no-context-files",
67 "-nc",
68 "--no-extensions",
69 "-ne",
70 "--no-prompt-templates",
71 "-np",
72 "--no-skills",
73 "-ns",
74 "--no-themes",
75 "--no-tools",
76 "-nt",
77 "--offline",
78 "--verbose",
79]);
80​
81export type PiLaunchArgsResolution =
82 | { readonly ok: true; readonly args: ReadonlyArray<string> }
83 | { readonly ok: false; readonly message: string };
84​
85function reservedPiArgument(arg: string): string | undefined {
86 for (const reserved of RESERVED_PI_LAUNCH_ARGUMENTS) {
87 if (arg === reserved || (reserved.startsWith("--") && arg.startsWith(`${reserved}=`))) {
88 return reserved;
89 }
90 }
91 return undefined;
93​
94function normalizePiBuiltInEqualsArguments(args: ReadonlyArray<string>): ReadonlyArray<string> {
95 return args.flatMap((arg) => {
96 const equalsIndex = arg.indexOf("=");
97 if (equalsIndex <= 0) return [arg];
98 const option = arg.slice(0, equalsIndex);
99 return PI_ARGUMENTS_WITH_VALUES.has(option) ? [option, arg.slice(equalsIndex + 1)] : [arg];
100 });
102​
103export function resolvePiLaunchArgs(launchArgs: string): PiLaunchArgsResolution {
104 // Pi treats equals-form tokens as extension flags; split only known built-ins.
105 const args = normalizePiBuiltInEqualsArguments(tokenizeCliArgs(launchArgs));
106 let hasProvider = false;
107 let hasModel = false;
108 for (let index = 0; index < args.length; index += 1) {
109 const arg = args[index];
110 if (arg === undefined) continue;
111 const reserved = reservedPiArgument(arg);
112 if (reserved !== undefined) {
113 return {
114 ok: false,
115 message: `Pi launch argument '${reserved}' is controlled by T3 Code and cannot be overridden.`,
116 };
117 }
118 if (arg === "--") {
119 return {
120 ok: false,
121 message: "Pi launch arguments cannot include positional prompts.",
122 };
123 }
124 if (PI_ARGUMENTS_WITH_VALUES.has(arg)) {
125 const value = args[index + 1];
126 if (value === undefined) {
127 return { ok: false, message: `Pi launch argument '${arg}' requires a value.` };
128 }
129 if (arg === "--provider") hasProvider = true;
130 if (arg === "--model") hasModel = true;
131 index += 1;
132 continue;
133 }
134 if (PI_ARGUMENTS_WITHOUT_VALUES.has(arg) || (arg.startsWith("--") && arg.includes("="))) {
135 continue;
136 }
137 if (arg.startsWith("--")) {
138 // Arbitrary extension flags may consume one following non-flag token.
139 if (
140 args[index + 1] !== undefined &&
141 !args[index + 1]!.startsWith("-") &&
142 !args[index + 1]!.startsWith("@")
143 ) {
144 index += 1;
145 }
146 continue;
147 }
148 if (arg.startsWith("-")) {
149 return { ok: false, message: `Pi launch argument '${arg}' is not supported by T3 Code.` };
150 }
151 return {
152 ok: false,
153 message: `Pi launch arguments cannot include positional prompt '${arg}'.`,
154 };
155 }
156 if (hasProvider && !hasModel) {
157 return { ok: false, message: "Pi launch argument '--provider' requires '--model'." };
158 }
159 return { ok: true, args };
161​
162function bearerTokenFromAuthorizationHeader(header: string): string {
163 return header.startsWith("Bearer ") ? header.slice("Bearer ".length) : header;
165​
166function normalizedPiPath(value: string): string {
167 return value.replace(/\\/g, "/").replace(/\/+$/, "");
169​
170function hasExplicitExtension(args: ReadonlyArray<string>, extensionPath: string): boolean {
171 const wanted = normalizedPiPath(extensionPath);
172 for (let index = 0; index < args.length; index += 1) {
173 const arg = args[index];
174 if (arg !== "--extension" && arg !== "-e") continue;
175 const configured = args[index + 1];
176 if (configured !== undefined && normalizedPiPath(configured) === wanted) return true;
177 index += 1;
178 }
179 return false;
181​
182function withoutExplicitExtensions(args: ReadonlyArray<string>): ReadonlyArray<string> {
183 const filtered: string[] = [];
184 for (let index = 0; index < args.length; index += 1) {
185 const arg = args[index];
186 if (arg === undefined) continue;
187 if (arg === "--extension" || arg === "-e") {
188 index += 1;
189 continue;
190 }
191 if (arg.startsWith("--extension=") || arg.startsWith("-e=")) continue;
192 filtered.push(arg);
193 }
194 return filtered;
196​
197function withoutToolSelectionArgs(args: ReadonlyArray<string>): ReadonlyArray<string> {
198 const filtered: string[] = [];
199 for (let index = 0; index < args.length; index += 1) {
200 const arg = args[index];
201 if (arg === undefined) continue;
202 if (arg === "--tools" || arg === "-t" || arg === "--exclude-tools" || arg === "-xt") {
203 index += 1;
204 continue;
205 }
206 if (
207 arg === "--no-tools" ||
208 arg === "-nt" ||
209 arg === "--no-builtin-tools" ||
210 arg === "-nbt" ||
211 arg.startsWith("--tools=") ||
212 arg.startsWith("-t=") ||
213 arg.startsWith("--exclude-tools=") ||
214 arg.startsWith("-xt=")
215 ) {
216 continue;
217 }
218 filtered.push(arg);
219 }
220 return filtered;
222​
223function piT3McpExtensionDestPath(cacheDir: string): string {
224 return `${cacheDir.replace(/\\/g, "/")}/${PI_T3_MCP_EXTENSION_FILENAME}`;
226​
227export const materializePiT3McpExtension = Effect.fn("materializePiT3McpExtension")(function* (
228 cacheDir: string,
229) {
230 const fs = yield* FileSystem.FileSystem;
231 yield* fs.makeDirectory(cacheDir, { recursive: true });
232 const dest = piT3McpExtensionDestPath(cacheDir);
233 const existing = yield* fs.readFileString(dest).pipe(Effect.orElseSucceed(() => ""));
234 if (existing !== PI_T3_MCP_EXTENSION_SOURCE) {
235 // Concurrent pi launches read this file: stage it in a per-call directory
236 // and rename so readers and other writers never see a partial file.
237 yield* Effect.scoped(
238 Effect.gen(function* () {
239 const stagingDir = yield* fs.makeTempDirectoryScoped({ directory: cacheDir });
240 const staging = `${stagingDir.replace(/\\/g, "/")}/${PI_T3_MCP_EXTENSION_FILENAME}`;
241 yield* fs.writeFileString(staging, PI_T3_MCP_EXTENSION_SOURCE);
242 yield* fs.rename(staging, dest).pipe(Effect.retry({ times: 3 }));
243 }),
244 );
245 }
246 return dest;
247});
248​
249export function buildPiRpcLaunch(input: {
250 readonly launchArgs: ReadonlyArray<string>;
251 readonly environment: NodeJS.ProcessEnv;
252 readonly mcpSession: McpProviderSessionConfig | undefined;
253 readonly extensionPath: string | undefined;
254 readonly sessionPath?: string | undefined;
255 readonly ephemeral?: boolean;
256 // T3 stops these processes as soon as they answer. They run --offline because Pi's startup
257 // catalog refresh renews expired OAuth tokens, and a stop mid-request loses the rotated one.
258 readonly shortLived?: boolean;
259 readonly disableExtensions?: boolean;
260 readonly disableTools?: boolean;
261 readonly runtimeMode?: "approval-required" | "auto-accept-edits" | "auto" | "full-access";
262}): {
263 readonly args: ReadonlyArray<string>;
264 readonly env: NodeJS.ProcessEnv;
265} {
266 const hasT3Extension = input.disableExtensions !== true && input.extensionPath !== undefined;
267 const hasT3Mcp = hasT3Extension && input.mcpSession !== undefined;
268 const extensionSafeArgs =
269 input.disableExtensions === true
270 ? withoutExplicitExtensions(input.launchArgs)
271 : input.launchArgs;
272 const launchArgs =
273 input.disableTools === true ? withoutToolSelectionArgs(extensionSafeArgs) : extensionSafeArgs;
274 const args = [
275 "--mode",
276 "rpc",
277 ...(input.ephemeral === true
278 ? ["--no-session"]
279 : input.sessionPath !== undefined
280 ? ["--session", input.sessionPath]
281 : []),
282 ...launchArgs,
283 ...(input.shortLived === true && !launchArgs.includes("--offline") ? ["--offline"] : []),
284 // Restrictions must follow launch args so user flags cannot re-enable unattended tools/extensions.
285 ...(input.disableExtensions === true ? ["--no-extensions"] : []),
286 ...(input.disableTools === true ? ["--no-tools"] : []),
287 ];
288 if (
289 hasT3Extension &&
290 input.extensionPath !== undefined &&
291 !hasExplicitExtension(args, input.extensionPath)
292 ) {
293 args.push("--extension", input.extensionPath);
294 }
295 const environment = { ...input.environment };
296 // MCP credentials belong to this T3 session; discard inherited server/parent credentials.
297 delete environment[T3_MCP_URL_ENV];
298 delete environment[T3_MCP_BEARER_ENV];
299 delete environment[T3_PI_MCP_EXTENSION_PATH_ENV];
300​
301 return {
302 args,
303 env: {
304 ...environment,
305 ...(hasT3Extension && input.extensionPath !== undefined
306 ? { [T3_PI_MCP_EXTENSION_PATH_ENV]: input.extensionPath }
307 : {}),
308 ...(hasT3Extension && input.runtimeMode !== undefined
309 ? {
310 [T3_PI_RUNTIME_MODE_ENV]:
311 input.runtimeMode === "auto" ? "approval-required" : input.runtimeMode,
312 }
313 : {}),
314 ...(hasT3Mcp && input.mcpSession !== undefined
315 ? {
316 [T3_MCP_URL_ENV]: input.mcpSession.endpoint,
317 [T3_MCP_BEARER_ENV]: bearerTokenFromAuthorizationHeader(
318 input.mcpSession.authorizationHeader,
319 ),
320 }
321 : {}),
322 },
323 };