Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions e2e/mcp-late-dispatch.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,17 @@ function linearClient(opts: {
},
],
call: opts.call ?? (async () => "ISSUE-1"),
async callResult(toolName, args, signal) {
const text = await (opts.call ?? (async () => "ISSUE-1"))(
toolName,
args,
signal,
);
return {
blocks: [{ type: "text", text }],
isError: false,
};
},
async close() {
return undefined;
},
Expand Down
32 changes: 24 additions & 8 deletions e2e/subagent-permission.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -296,9 +296,13 @@ test.serial(
inputSchema: { type: "object", properties: {} },
},
],
call: async () => {
call: async () => "changed",
callResult: async () => {
calls++;
return "changed";
return {
blocks: [{ type: "text", text: "changed" }],
isError: false,
};
},
close: async () => undefined,
};
Expand Down Expand Up @@ -350,9 +354,13 @@ test.serial(
inputSchema: { type: "object", properties: {} },
},
],
call: async () => {
call: async () => "changed",
callResult: async () => {
calls++;
return "changed";
return {
blocks: [{ type: "text", text: "changed" }],
isError: false,
};
},
close: async () => undefined,
};
Expand Down Expand Up @@ -480,9 +488,13 @@ test.serial(
inputSchema: { type: "object", properties: {} },
},
],
call: async () => {
call: async () => "changed",
callResult: async () => {
calls++;
return "changed";
return {
blocks: [{ type: "text", text: "changed" }],
isError: false,
};
},
close: async () => undefined,
};
Expand Down Expand Up @@ -561,9 +573,13 @@ test("storing parent-gated MCP tools then wrapping again denies without requestA
inputSchema: { type: "object", properties: {} },
},
],
call: async () => {
call: async () => "changed",
callResult: async () => {
calls++;
return "changed";
return {
blocks: [{ type: "text", text: "changed" }],
isError: false,
};
},
close: async () => undefined,
};
Expand Down
134 changes: 36 additions & 98 deletions src/mcp/add-server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,44 +14,34 @@ import {
import { deleteAuthState } from "./auth-store.js";
import { EXA_MCP_SERVER_NAME } from "./exa.js";

export interface GlobalSettingsWriter {
readonly enqueue: <T>(job: () => Promise<T>) => Promise<T>;
readonly update: (
apply: (base: Settings) => Settings | null,
) => Promise<Settings | null>;
interface SettingsWriter<T> {
readonly enqueue: <R>(job: () => Promise<R>) => Promise<R>;
readonly update: (apply: (base: T) => T | null) => Promise<T | null>;
readonly updateAt: (
path: string,
apply: (base: Settings) => Settings | null,
) => Promise<Settings | null>;
readonly mutate: (
apply: (base: Settings) => Settings | null,
) => Promise<"ok" | "skipped">;
apply: (base: T) => T | null,
) => Promise<T | null>;
readonly mutate: (apply: (base: T) => T | null) => Promise<"ok" | "skipped">;
readonly mutateAt: (
path: string,
apply: (base: Settings) => Settings | null,
apply: (base: T) => T | null,
) => Promise<"ok" | "skipped">;
}

interface GlobalSettingsWriterDeps {
readonly load: (path: string) => Promise<Settings | null>;
readonly save: (path: string, settings: Settings) => Promise<void>;
}
export type GlobalSettingsWriter = SettingsWriter<Settings>;
export type LocalSettingsWriter = SettingsWriter<LocalSettings>;

const defaultWriterDeps: GlobalSettingsWriterDeps = {
load: loadGlobalSettingsWriteBase,
save: saveGlobalSettings,
};
interface SettingsWriterDeps<T> {
readonly load: (path: string) => Promise<T | null>;
readonly save: (path: string, settings: T) => Promise<void>;
}

export function createGlobalSettingsWriter(
function createSettingsWriter<T>(
path: string,
deps: Partial<GlobalSettingsWriterDeps> = {},
): GlobalSettingsWriter {
const writerDeps: GlobalSettingsWriterDeps = {
...defaultWriterDeps,
...deps,
};
deps: SettingsWriterDeps<T>,
): SettingsWriter<T> {
let tail = Promise.resolve();
const enqueue = <T>(job: () => Promise<T>): Promise<T> => {
const enqueue = <R>(job: () => Promise<R>): Promise<R> => {
const run = tail.then(job);
tail = run.then(
() => undefined,
Expand All @@ -62,20 +52,20 @@ export function createGlobalSettingsWriter(

const updateAt = (
settingsPath: string,
apply: (base: Settings) => Settings | null,
): Promise<Settings | null> =>
apply: (base: T) => T | null,
): Promise<T | null> =>
enqueue(async () => {
const base = await writerDeps.load(settingsPath);
const base = await deps.load(settingsPath);
if (base === null) return null;
const next = apply(base);
if (next === null) return base;
await writerDeps.save(settingsPath, next);
await deps.save(settingsPath, next);
return next;
});

const mutateAt = async (
settingsPath: string,
apply: (base: Settings) => Settings | null,
apply: (base: T) => T | null,
): Promise<"ok" | "skipped"> =>
(await updateAt(settingsPath, apply)) === null ? "skipped" : "ok";

Expand All @@ -88,78 +78,26 @@ export function createGlobalSettingsWriter(
};
}

export interface LocalSettingsWriter {
readonly enqueue: <T>(job: () => Promise<T>) => Promise<T>;
readonly update: (
apply: (base: LocalSettings) => LocalSettings | null,
) => Promise<LocalSettings | null>;
readonly updateAt: (
path: string,
apply: (base: LocalSettings) => LocalSettings | null,
) => Promise<LocalSettings | null>;
readonly mutate: (
apply: (base: LocalSettings) => LocalSettings | null,
) => Promise<"ok" | "skipped">;
readonly mutateAt: (
path: string,
apply: (base: LocalSettings) => LocalSettings | null,
) => Promise<"ok" | "skipped">;
}

interface LocalSettingsWriterDeps {
readonly load: (path: string) => Promise<LocalSettings | null>;
readonly save: (path: string, settings: LocalSettings) => Promise<void>;
export function createGlobalSettingsWriter(
path: string,
deps: Partial<SettingsWriterDeps<Settings>> = {},
): GlobalSettingsWriter {
return createSettingsWriter(path, {
load: loadGlobalSettingsWriteBase,
save: saveGlobalSettings,
...deps,
});
}

const defaultLocalWriterDeps: LocalSettingsWriterDeps = {
load: loadLocalSettingsWriteBase,
save: saveLocalSettings,
};

export function createLocalSettingsWriter(
path: string,
deps: Partial<LocalSettingsWriterDeps> = {},
deps: Partial<SettingsWriterDeps<LocalSettings>> = {},
): LocalSettingsWriter {
const writerDeps: LocalSettingsWriterDeps = {
...defaultLocalWriterDeps,
return createSettingsWriter(path, {
load: loadLocalSettingsWriteBase,
save: saveLocalSettings,
...deps,
};
let tail = Promise.resolve();
const enqueue = <T>(job: () => Promise<T>): Promise<T> => {
const run = tail.then(job);
tail = run.then(
() => undefined,
() => undefined,
);
return run;
};

const updateAt = (
settingsPath: string,
apply: (base: LocalSettings) => LocalSettings | null,
): Promise<LocalSettings | null> =>
enqueue(async () => {
const base = await writerDeps.load(settingsPath);
if (base === null) return null;
const next = apply(base);
if (next === null) return base;
await writerDeps.save(settingsPath, next);
return next;
});

const mutateAt = async (
settingsPath: string,
apply: (base: LocalSettings) => LocalSettings | null,
): Promise<"ok" | "skipped"> =>
(await updateAt(settingsPath, apply)) === null ? "skipped" : "ok";

return {
enqueue,
update: (apply) => updateAt(path, apply),
updateAt,
mutate: (apply) => mutateAt(path, apply),
mutateAt,
};
});
}

export type PersistMCPServerResult =
Expand Down
63 changes: 16 additions & 47 deletions src/mcp/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import {
} from "../telemetry/product-events.js";
import { liveTelemetry } from "../telemetry/singleton.js";
import { MCP_CLIENT_NAME } from "../branding.js";
import { closeQuietly, isAbortError } from "./quiet.js";

export interface MCPTool {
name: string;
Expand Down Expand Up @@ -60,16 +61,9 @@ export interface MCPClient {
args: Record<string, unknown>,
signal: AbortSignal,
): Promise<string>;
/** Validated content blocks before flattening — for post-policy
* archive capture. */
callBlocks?(
toolName: string,
args: Record<string, unknown>,
signal: AbortSignal,
): Promise<MCPContentBlock[]>;
/** Full tool-result envelope: blocks plus tool-level
* isError/structuredContent. */
callResult?(
callResult(
toolName: string,
args: Record<string, unknown>,
signal: AbortSignal,
Expand Down Expand Up @@ -176,15 +170,6 @@ function isRecoverableAuthError(
return err instanceof UnauthorizedError || err instanceof OAuthError;
}

function isAbortError(err: unknown): boolean {
return (
typeof err === "object" &&
err !== null &&
"name" in err &&
err.name === "AbortError"
);
}

export const MAX_BROWSER_AUTH_ATTEMPTS = 1;
export const BROWSER_AUTH_COOLDOWN_MS = 5 * 60_000;
export const BROWSER_AUTH_WAIT_MS = 2 * 60_000;
Expand Down Expand Up @@ -612,12 +597,13 @@ async function finishClient(
serverName,
tools,
async callResult(toolName, args, signal) {
const context =
authContext === undefined ? undefined : { ...authContext, signal };
const result = await withHTTPAuthorizationRecovery(context, () =>
client.callTool({ name: toolName, arguments: args }, undefined, {
signal,
}),
const result = await withHTTPAuthorizationRecovery(
authContext,
() =>
client.callTool({ name: toolName, arguments: args }, undefined, {
signal,
}),
signal,
);
const envelope: MCPToolResultEnvelope = {
blocks: validateMcpContentBlocks(result.content),
Expand All @@ -634,34 +620,17 @@ async function finishClient(
}
return envelope;
},
async callBlocks(toolName, args, signal) {
const result = await withHTTPAuthorizationRecovery(
authContext,
() =>
client.callTool({ name: toolName, arguments: args }, undefined, {
signal,
}),
signal,
);
return validateMcpContentBlocks(result.content);
},
async call(toolName, args, signal) {
const result = await withHTTPAuthorizationRecovery(
authContext,
() =>
client.callTool({ name: toolName, arguments: args }, undefined, {
signal,
}),
signal,
async call(this: MCPClient, toolName, args, signal) {
return unwrapToolContent(
(await this.callResult(toolName, args, signal)).blocks,
);
return unwrapToolContent(validateMcpContentBlocks(result.content));
},
async close() {
// Disarm first: the owner's hook reports unintentional death only.
if (liveTransport !== undefined) liveTransport.onclose = () => undefined;
closeLifecycle?.();
authContext?.callback.close();
await client.close().catch(() => undefined);
await closeQuietly(() => client.close());
},
};
}
Expand Down Expand Up @@ -730,7 +699,7 @@ async function connectStdio(
),
};
} catch (err) {
await client.close().catch(() => undefined);
await closeQuietly(() => client.close());
return {
ok: false,
serverName: config.name,
Expand Down Expand Up @@ -798,7 +767,7 @@ async function connectHttp(
signal: lifecycle.signal,
});
} finally {
await probeClient.close().catch(() => undefined);
await closeQuietly(() => probeClient.close());
}
},
};
Expand Down Expand Up @@ -855,7 +824,7 @@ async function connectHttp(
};
} catch (err) {
closeLifecycle();
await client?.close().catch(() => undefined);
await closeQuietly(() => client?.close() ?? Promise.resolve());
try {
callback?.close();
} catch {
Expand Down
Loading
Loading