diff --git a/bun.lock b/bun.lock index 9f50f7cd..80439732 100644 --- a/bun.lock +++ b/bun.lock @@ -18,6 +18,7 @@ "semver": "7.8.5", "smol-toml": "^1.9.0", "undici": "7.30.0", + "ws": "^8.22.0", "yaml": "^2.9.1", "zod": "^4.6.5", }, @@ -25,6 +26,7 @@ "@biomejs/biome": "2.5.15", "@types/bun": "1.4.2", "@types/semver": "^7.8.0", + "@types/ws": "^8.18.2", "braintrust": "3.37.0", "bunup": "^0.16.32", "husky": "^9.1.7", @@ -48,6 +50,7 @@ "version": "0.32.0", "dependencies": { "@modelcontextprotocol/sdk": "^1.29.0", + "ws": "^8.22.0", "zod": "^4.4.3", }, }, @@ -331,6 +334,8 @@ "@types/semver": ["@types/semver@7.8.0", "", {}, "sha512-1mAINjtQCXXeLkJ9ehXkwOcBpqtLxiVtKhpUf83DdRNdQKV0iXZpaHYqRr7nj+wvxuJzoAmAwXI+sCNMv1CzLQ=="], + "@types/ws": ["@types/ws@8.18.2", "", { "dependencies": { "@types/node": "*" } }, "sha512-67MQl+fpWKVTT1NYdnmo3U4sc/xPo/zQBncVnI74qmQa0z/b+1g6iYqNmGCPbxO+zz2aklb08a0oHfegiVd0/w=="], + "@typescript/typescript-aix-ppc64": ["@typescript/typescript-aix-ppc64@7.0.2", "", { "os": "aix", "cpu": "ppc64" }, "sha512-MTKKkWB7p/0E9xi1d1tHtZ5PiLkGEMIq88pK2CubZjOsLtYTLqhgIgi6zepFa+9GHZ6h05NMCkQxGKiPXMxXtQ=="], "@typescript/typescript-darwin-arm64": ["@typescript/typescript-darwin-arm64@7.0.2", "", { "os": "darwin", "cpu": "arm64" }, "sha512-gowzar9MwS/aRWp6f3a4KUqzRjAZjOsmGNCM6LcTgXum+dBfgsBVMN+AgvOCCbguXyick6LJhpBszxMebJ8syA=="], @@ -721,6 +726,8 @@ "wrappy": ["wrappy@1.0.2", "", {}, "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ=="], + "ws": ["ws@8.22.0", "", { "peerDependencies": { "bufferutil": "^4.0.1", "utf-8-validate": ">=5.0.2" }, "optionalPeers": ["bufferutil", "utf-8-validate"] }, "sha512-Ydggc987+RO0AnWtZ/7Wq9FtNvcrL1b/RO0ud9mWjUPgDrsAAwQSF51sm2hm1XofbU/4jkpGEsLFsZZxU+1DOg=="], + "wsl-utils": ["wsl-utils@1.0.1", "", { "dependencies": { "is-wsl": "^3.1.0", "powershell-utils": "^0.2.0" } }, "sha512-c16PKLIRxBYpR3jSWsdS7SFp0COm5SOQdiFjY818dYFHJbMGuLBszpF07LBWmOS1uzrHa9VGdynsEgoeKy0oOQ=="], "yaml": ["yaml@2.9.1", "", { "bin": { "yaml": "bin.mjs" } }, "sha512-3NxN8+78OdzbT7C/WjGsyfPAtJaN3FNDsWxv7Y7mcDsT/oOmgW8BpyQQFFBnvZE3j9Y2Sdz1ULFLezL7Eb2yFw=="], diff --git a/changes/mcp-tool-progress.added.md b/changes/mcp-tool-progress.added.md new file mode 100644 index 00000000..8f78cca1 --- /dev/null +++ b/changes/mcp-tool-progress.added.md @@ -0,0 +1,6 @@ +--- +"githits": minor +"@githits/mcp": minor +--- + +- **Live progress for `research` and `get_example`** - When an MCP host requests progress, the local MCP server now forwards GitHits status updates as `notifications/progress` while these long-running tools work; hosts that display them (such as Claude Code) show the current stage under the running call. Progress is best effort and authorized by the caller's own token: if it is unavailable, calls run exactly as before and the final result is unchanged. `@githits/mcp` adds an optional `progress` server option and a credential-free `RelayReceiver`, and now depends on `ws`. diff --git a/changes/project-key-token-refresh.fixed.md b/changes/project-key-token-refresh.fixed.md new file mode 100644 index 00000000..ea0bd57e --- /dev/null +++ b/changes/project-key-token-refresh.fixed.md @@ -0,0 +1,6 @@ +--- +"githits": patch +"@githits/mcp": patch +--- + +- **No refresh attempt for project API keys** - Authentication failures with a `gh-proj-` project API key no longer trigger a credential refresh, matching personal `ghi-` API tokens; only signed-in sessions refresh. diff --git a/package.json b/package.json index 702904c7..98723532 100644 --- a/package.json +++ b/package.json @@ -116,6 +116,7 @@ "semver": "7.8.5", "smol-toml": "^1.9.0", "undici": "7.30.0", + "ws": "^8.22.0", "yaml": "^2.9.1", "zod": "^4.6.5" }, @@ -123,6 +124,7 @@ "@biomejs/biome": "2.5.15", "@types/bun": "1.4.2", "@types/semver": "^7.8.0", + "@types/ws": "^8.18.2", "braintrust": "3.37.0", "bunup": "^0.16.32", "husky": "^9.1.7", diff --git a/packages/core-internal/src/index.ts b/packages/core-internal/src/index.ts index 7af96b6c..7ad2bc83 100644 --- a/packages/core-internal/src/index.ts +++ b/packages/core-internal/src/index.ts @@ -9,6 +9,7 @@ export * from "./services/grep-service.js"; export * from "./services/indexing-estimates.js"; export * from "./services/list-service.js"; export * from "./services/package-intelligence-service.js"; +export * from "./services/progress-request.js"; export * from "./services/promote-version-not-found.js"; export * from "./services/read-service.js"; export * from "./services/read-target.js"; diff --git a/packages/core-internal/src/services/agentic-ask-service.ts b/packages/core-internal/src/services/agentic-ask-service.ts index fe604df6..e9a7741c 100644 --- a/packages/core-internal/src/services/agentic-ask-service.ts +++ b/packages/core-internal/src/services/agentic-ask-service.ts @@ -7,6 +7,11 @@ import { isTokenRefreshableError, parseRetryAfterSeconds, } from "./githits-service.js"; +import { + discardResponseBody, + type ProgressRequest, + readProgressRejection, +} from "./progress-request.js"; import { type ServiceDiagnostics, withServiceDiagnostics, @@ -50,6 +55,8 @@ export type AgenticAskRequest = AgenticAskQuestion & export interface AgenticAskRequestOptions { signal?: AbortSignal; + /** Optional best-effort progress correlation; never changes the result. */ + progress?: ProgressRequest; } /** Expected backend envelope; successful JSON is deliberately not runtime-validated. */ @@ -181,7 +188,7 @@ export class AgenticAskServiceImpl implements AgenticAskService { error.code === "AUTH_REQUIRED") || isTokenRefreshableError(error), executeWithToken: (token) => - this.executeAsk(token, request, signal), + this.executeAsk(token, request, signal, options.progress), }), options.signal, this.runtime.timeoutMs ?? AGENTIC_ASK_REQUEST_TIMEOUT_MS, @@ -193,6 +200,7 @@ export class AgenticAskServiceImpl implements AgenticAskService { token: string, request: AgenticAskRequest, signal: AbortSignal, + progress?: ProgressRequest, ): Promise { const apiUrl = validateServiceUrl(this.apiUrl, "GITHITS_API_URL"); let response: Response; @@ -201,6 +209,7 @@ export class AgenticAskServiceImpl implements AgenticAskService { method: "POST", headers: { ...this.runtime.clientHeaders?.(), + ...progress?.headers, Authorization: `Bearer ${token}`, "Content-Type": "application/json", "User-Agent": this.runtime.userAgent ?? "githits-cli", @@ -222,6 +231,14 @@ export class AgenticAskServiceImpl implements AgenticAskService { throw cause; } + const progressRejection = progress && readProgressRejection(response); + if (progressRejection) { + // Rejected before admission: nothing ran, so resend without progress. + progress.rejected(progressRejection); + await discardResponseBody(response); + return this.executeAsk(token, request, signal); + } + const toolCallId = readAgenticAskResponseId( response.headers.get("X-GitHits-Tool-Call-Id"), ); diff --git a/packages/core-internal/src/services/execute-with-token-refresh.test.ts b/packages/core-internal/src/services/execute-with-token-refresh.test.ts index c210c677..c4a544e2 100644 --- a/packages/core-internal/src/services/execute-with-token-refresh.test.ts +++ b/packages/core-internal/src/services/execute-with-token-refresh.test.ts @@ -6,24 +6,27 @@ import { } from "./githits-service.js"; describe("executeWithTokenRefresh", () => { - it("never calls the refresh hook for an opaque ghi-* token", async () => { - const forceRefresh = mock(() => Promise.resolve("unexpected-token")); - const executeWithToken = mock(() => - Promise.reject(new TermsAcceptanceRequiredError()), - ); - - await expect( - executeWithTokenRefresh({ - getToken: mock(() => Promise.resolve("ghi-static-token")), - forceRefresh, - executeWithToken, - shouldRefresh: () => true, - }), - ).rejects.toBeInstanceOf(TermsAcceptanceRequiredError); - - expect(forceRefresh).not.toHaveBeenCalled(); - expect(executeWithToken).toHaveBeenCalledTimes(1); - }); + it.each(["ghi-static-token", "gh-proj-static-token"])( + "never calls the refresh hook for opaque API token %s", + async (token) => { + const forceRefresh = mock(() => Promise.resolve("unexpected-token")); + const executeWithToken = mock(() => + Promise.reject(new TermsAcceptanceRequiredError()), + ); + + await expect( + executeWithTokenRefresh({ + getToken: mock(() => Promise.resolve(token)), + forceRefresh, + executeWithToken, + shouldRefresh: () => true, + }), + ).rejects.toBeInstanceOf(TermsAcceptanceRequiredError); + + expect(forceRefresh).not.toHaveBeenCalled(); + expect(executeWithToken).toHaveBeenCalledTimes(1); + }, + ); it("rejects with an already-aborted caller reason before reading a token", async () => { const controller = new AbortController(); const reason = new Error("already cancelled"); diff --git a/packages/core-internal/src/services/execute-with-token-refresh.ts b/packages/core-internal/src/services/execute-with-token-refresh.ts index 8058d2c4..8c5b927b 100644 --- a/packages/core-internal/src/services/execute-with-token-refresh.ts +++ b/packages/core-internal/src/services/execute-with-token-refresh.ts @@ -39,9 +39,9 @@ export async function executeWithTokenRefresh( // A cancelled caller no longer wants the result; surface the caller's // reason instead of refreshing credentials for a retry. signal?.throwIfAborted(); - // Opaque ghi-* credentials are re-evaluated server-side and have no local - // refresh flow. Do not invoke the provider's refresh hook for them. - if (token.startsWith("ghi-") || !options.shouldRefresh(error)) { + // Opaque API tokens (personal ghi-*, project gh-proj-*) are re-evaluated + // server-side and have no local refresh flow; only Supabase JWTs refresh. + if (isOpaqueApiToken(token) || !options.shouldRefresh(error)) { throw error; } @@ -72,3 +72,7 @@ async function callerFirst( throw error; } } + +function isOpaqueApiToken(token: string): boolean { + return token.startsWith("ghi-") || token.startsWith("gh-proj-"); +} diff --git a/packages/core-internal/src/services/githits-service.ts b/packages/core-internal/src/services/githits-service.ts index b99b0864..bb410bab 100644 --- a/packages/core-internal/src/services/githits-service.ts +++ b/packages/core-internal/src/services/githits-service.ts @@ -16,6 +16,10 @@ import { AuthenticationError, SERVER_AUTHENTICATION_REJECTED_MESSAGE, } from "./githits-service-errors.js"; +import { + type ProgressRequest, + readProgressRejection, +} from "./progress-request.js"; import { type ServiceDiagnostics, withServiceDiagnostics, @@ -107,6 +111,8 @@ export interface SearchParams { /** Optional browser-standard controls for a GitHits service request. */ export interface GitHitsServiceRequestOptions { signal?: AbortSignal; + /** Optional best-effort progress correlation; never changes the result. */ + progress?: ProgressRequest; } export interface GitHitsServiceRuntimeOptions { @@ -148,23 +154,32 @@ export class GitHitsServiceImpl implements GitHitsService { this.runtime.diagnostics, "githits.search.request", async () => { - const { response, body } = await this.request( - "/search", - { - method: "POST", - headers: this.headers(), - body: JSON.stringify({ - query: params.query, - source_format: params.sourceFormat, - language: params.language, - license_mode: params.licenseMode ?? "strict", - include_explanation: params.includeExplanation ?? false, - }), - ...(options?.signal ? { signal: options.signal } : {}), - }, - this.runtime.exampleRequestTimeoutMs ?? - DEFAULT_EXAMPLE_REQUEST_TIMEOUT_MS, - ); + const send = (progress?: ProgressRequest) => + this.request( + "/search", + { + method: "POST", + headers: this.headers(progress?.headers), + body: JSON.stringify({ + query: params.query, + source_format: params.sourceFormat, + language: params.language, + license_mode: params.licenseMode ?? "strict", + include_explanation: params.includeExplanation ?? false, + }), + ...(options?.signal ? { signal: options.signal } : {}), + }, + this.runtime.exampleRequestTimeoutMs ?? + DEFAULT_EXAMPLE_REQUEST_TIMEOUT_MS, + ); + let { response, body } = await send(options?.progress); + const progressRejection = + options?.progress && readProgressRejection(response); + if (progressRejection) { + // Rejected before admission: nothing ran, so resend without progress. + options.progress?.rejected(progressRejection); + ({ response, body } = await send()); + } if (!response.ok) { throw this.createError(response, body); @@ -175,9 +190,12 @@ export class GitHitsServiceImpl implements GitHitsService { ); } - private headers(): Record { + private headers( + extra?: Readonly>, + ): Record { return { ...this.runtime.clientHeaders?.(), + ...extra, Authorization: `Bearer ${this.token}`, "Content-Type": "application/json", "User-Agent": this.runtime.userAgent ?? "githits-cli", diff --git a/packages/core-internal/src/services/progress-request.test.ts b/packages/core-internal/src/services/progress-request.test.ts new file mode 100644 index 00000000..4a97ee6a --- /dev/null +++ b/packages/core-internal/src/services/progress-request.test.ts @@ -0,0 +1,192 @@ +import { describe, expect, it, mock } from "bun:test"; +import { AgenticAskServiceImpl } from "./agentic-ask-service.js"; +import displayContract from "./fixtures/ask-display-contract.json"; +import { GitHitsServiceImpl } from "./githits-service.js"; +import type { ProgressRequest } from "./progress-request.js"; +import { createMockTokenProvider } from "./test-helpers.js"; + +const PROGRESS_HEADERS = { + "X-GitHits-Progress-Id": "8f9a1c2e-0d4b-4e5f-9a6b-7c8d9e0f1a2b", + "X-GitHits-Relay-Id": "1b2c3d4e-5f60-4718-8a9b-0c1d2e3f4a5b", +}; + +interface Captured { + headers: Headers; + body: string; +} + +function recordingFetch(responses: Response[]): { + fetchFn: typeof fetch; + calls: Captured[]; +} { + const calls: Captured[] = []; + const fetchFn = mock( + async (_url: string | URL | Request, init?: RequestInit) => { + calls.push({ + headers: new Headers(init?.headers), + body: String(init?.body ?? ""), + }); + const next = responses.shift(); + if (!next) throw new Error("unexpected request"); + return next; + }, + ); + return { fetchFn: fetchFn as unknown as typeof fetch, calls }; +} + +function progressRequest(): ProgressRequest & { + rejected: ReturnType; +} { + return { headers: PROGRESS_HEADERS, rejected: mock(() => {}) }; +} + +function rejection(status: number, code: string): Response { + return new Response(JSON.stringify({ detail: "Progress request rejected" }), { + status, + headers: { "X-GitHits-Progress-Error": code }, + }); +} + +function askService(fetchFn: typeof fetch): AgenticAskServiceImpl { + return new AgenticAskServiceImpl( + "https://api.githits.test", + createMockTokenProvider(), + fetchFn, + { + clientHeaders: () => ({ + "x-githits-session-id": "session-id", + // A colliding header must never override auth. + Authorization: "Bearer client-header", + }), + }, + ); +} + +function searchService(fetchFn: typeof fetch): GitHitsServiceImpl { + return new GitHitsServiceImpl( + "https://api.githits.test", + "token", + fetchFn, + undefined, + { + clientHeaders: () => ({ "x-githits-session-id": "session-id" }), + }, + ); +} + +const question = { + target: "pypi:httpx", + question: "default timeout?", +} as const; + +describe("progress request headers", () => { + it("adds correlation headers to /ask without touching session or auth", async () => { + const { fetchFn, calls } = recordingFetch([ + Response.json(displayContract.cli), + ]); + const progress = progressRequest(); + await askService(fetchFn).ask(question, { progress }); + expect(calls).toHaveLength(1); + const headers = calls[0]?.headers; + expect(headers?.get("X-GitHits-Progress-Id")).toBe( + PROGRESS_HEADERS["X-GitHits-Progress-Id"], + ); + expect(headers?.get("X-GitHits-Relay-Id")).toBe( + PROGRESS_HEADERS["X-GitHits-Relay-Id"], + ); + expect(headers?.get("x-githits-session-id")).toBe("session-id"); + expect(headers?.get("Authorization")).not.toBe("Bearer client-header"); + expect(progress.rejected).not.toHaveBeenCalled(); + }); + + it("adds correlation headers to /search", async () => { + const { fetchFn, calls } = recordingFetch([new Response("# Example")]); + const progress = progressRequest(); + const result = await searchService(fetchFn).search( + { query: "timeouts", sourceFormat: "mcp" }, + { progress }, + ); + expect(result).toBe("# Example"); + expect(calls[0]?.headers.get("X-GitHits-Progress-Id")).toBe( + PROGRESS_HEADERS["X-GitHits-Progress-Id"], + ); + expect(calls[0]?.headers.get("x-githits-session-id")).toBe("session-id"); + }); + + it("omits progress headers when no progress is requested", async () => { + const { fetchFn, calls } = recordingFetch([ + Response.json(displayContract.cli), + ]); + await askService(fetchFn).ask(question); + expect(calls[0]?.headers.has("X-GitHits-Progress-Id")).toBe(false); + }); +}); + +describe("pre-execution progress rejection", () => { + it.each([ + [400, "invalid_progress_request"], + [401, "service_auth_required"], + [503, "unavailable"], + ])("/ask resends once without progress after %i %s", async (status, code) => { + const { fetchFn, calls } = recordingFetch([ + rejection(status, code), + Response.json(displayContract.cli), + ]); + const progress = progressRequest(); + const response = await askService(fetchFn).ask(question, { progress }); + expect(response.display_markdown).toBe( + displayContract.cli.display_markdown, + ); + expect(calls).toHaveLength(2); + expect(calls[1]?.headers.has("X-GitHits-Progress-Id")).toBe(false); + expect(calls[1]?.headers.has("X-GitHits-Relay-Id")).toBe(false); + expect(calls[1]?.body).toBe(calls[0]?.body); + expect(progress.rejected).toHaveBeenCalledWith(code); + }); + + it("/search resends once without progress", async () => { + const { fetchFn, calls } = recordingFetch([ + rejection(503, "unavailable"), + new Response("# Example"), + ]); + const progress = progressRequest(); + const result = await searchService(fetchFn).search( + { query: "timeouts", sourceFormat: "mcp" }, + { progress }, + ); + expect(result).toBe("# Example"); + expect(calls).toHaveLength(2); + expect(calls[1]?.headers.has("X-GitHits-Progress-Id")).toBe(false); + expect(calls[1]?.body).toBe(calls[0]?.body); + expect(progress.rejected).toHaveBeenCalledWith("unavailable"); + }); + + it("does not resend ordinary errors that lack the progress marker", async () => { + const { fetchFn, calls } = recordingFetch([ + new Response("busy", { status: 503 }), + ]); + const progress = progressRequest(); + await expect( + searchService(fetchFn).search( + { query: "q", sourceFormat: "mcp" }, + { progress }, + ), + ).rejects.toThrow(); + expect(calls).toHaveLength(1); + expect(progress.rejected).not.toHaveBeenCalled(); + }); + + it("does not resend a second time when the retry also fails", async () => { + const { fetchFn, calls } = recordingFetch([ + rejection(503, "unavailable"), + rejection(503, "unavailable"), + ]); + await expect( + searchService(fetchFn).search( + { query: "q", sourceFormat: "mcp" }, + { progress: progressRequest() }, + ), + ).rejects.toThrow(); + expect(calls).toHaveLength(2); + }); +}); diff --git a/packages/core-internal/src/services/progress-request.ts b/packages/core-internal/src/services/progress-request.ts new file mode 100644 index 00000000..af39104d --- /dev/null +++ b/packages/core-internal/src/services/progress-request.ts @@ -0,0 +1,32 @@ +/** + * Optional per-call progress correlation for backend requests. + * + * Progress is best effort: the HTTP result stays authoritative. The backend + * validates progress headers before admission, accounting and execution, so a + * rejection carrying `X-GitHits-Progress-Error` means nothing ran and the same + * request is safe to resend once without the progress headers. + */ +export interface ProgressRequest { + /** Correlation headers merged after client headers and before auth. */ + readonly headers: Readonly>; + /** Reports a pre-execution progress rejection; the call is then resent without progress. */ + rejected(code: string): void; +} + +export const PROGRESS_ERROR_HEADER = "X-GitHits-Progress-Error"; + +/** Returns the backend's progress rejection code, if this response is one. */ +export function readProgressRejection(response: Response): string | undefined { + if (response.ok) return undefined; + const code = response.headers.get(PROGRESS_ERROR_HEADER)?.trim(); + return code ? code : undefined; +} + +/** Releases an unread response body before resending; failures are irrelevant. */ +export async function discardResponseBody(response: Response): Promise { + try { + await response.body?.cancel(); + } catch { + // The connection is being abandoned either way. + } +} diff --git a/packages/mcp/README.md b/packages/mcp/README.md index 2fdfd91c..a48aa808 100644 --- a/packages/mcp/README.md +++ b/packages/mcp/README.md @@ -157,6 +157,8 @@ boundary is authenticated and reached. The package expects callers to provide service implementations through `McpToolServices` or a request-scoped `McpToolServicesProvider`. `McpToolServices.readService` is required for the advertised compact `read` tool; remote hosts should construct the `ReadServiceImpl` instance per request with that request's endpoint, token, headers, fetch function, and diagnostics. GitHits does not populate MCP initialize instructions because hosts expose them inconsistently; `quick_start` owns shared guidance instead. Callers may still pass their own `instructions` explicitly. Use `quickStartOptions` to configure the guide. Servers can pass `traceTool` to `createMcpServer()` or `registerMcpTools()` to wrap public tool execution for instrumentation without receiving arguments or auth data. +Servers can pass an optional `progress` hook to `createMcpServer()` or `registerMcpTools()` to forward live status for `research` and `get_example`. It is used only when the client sends a progress token (including `0`); updates become `notifications/progress` and stop before the tool returns. `RelayReceiver` is the bundled hook: it receives GitHits operation progress over one shared WebSocket, authorized either by the caller's own GitHits token (`{ kind: "caller", getToken }`) or by a service credential supplied at runtime (`{ kind: "service", serviceId, key }`). The package ships no credentials. Progress is best effort: if it is unavailable or rejected, the tool runs unchanged, resending a rejected request once without progress, and the final result is always authoritative. Call `close()` on shutdown. + Network-backed tools forward the MCP request's abort signal to their service call as an optional trailing `ServiceRequestOptions` argument (`{ signal }`), so a cancelled call or closed connection aborts the backend request. The bundled service implementations behind these tools honor it; custom service implementations should forward `options?.signal` to their own requests. Calls without a signal keep their single-argument shape. Only imports from `@githits/mcp`, `@githits/mcp/client`, `@githits/mcp/smoke-test`, `@githits/mcp/tools`, and `@githits/mcp/package.json` are public. The workspace alias `@githits/mcp/internal` is not exported, is not supported for external consumers, and must not be used by remote MCP server implementations. diff --git a/packages/mcp/bunup.config.ts b/packages/mcp/bunup.config.ts index 0cc11da8..30551d8e 100644 --- a/packages/mcp/bunup.config.ts +++ b/packages/mcp/bunup.config.ts @@ -12,7 +12,11 @@ export default defineConfig({ }, entry: ["src/index.ts", "src/client.ts", "src/smoke-test.ts", "src/tools.ts"], exports: true, - external: [/^@modelcontextprotocol\/sdk(?:\/.*)?$/, /^zod(?:\/.*)?$/], + external: [ + /^@modelcontextprotocol\/sdk(?:\/.*)?$/, + /^ws(?:\/.*)?$/, + /^zod(?:\/.*)?$/, + ], minifyWhitespace: true, packages: "bundle", preferredTsconfig: "../../tsconfig.json", diff --git a/packages/mcp/package.json b/packages/mcp/package.json index 641fec8e..5df2ff1a 100644 --- a/packages/mcp/package.json +++ b/packages/mcp/package.json @@ -65,6 +65,7 @@ }, "dependencies": { "@modelcontextprotocol/sdk": "^1.29.0", + "ws": "^8.22.0", "zod": "^4.4.3" } } diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index c2d04517..28d5af48 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -23,4 +23,13 @@ export { getMcpToolDescriptors, registerMcpTools, } from "./mcp/server.js"; +export type { + McpProgressHook, + McpProgressOperation, + McpProgressUpdate, + RelayAuth, + RelayReceiverOptions, + RelaySocket, +} from "./progress/relay-receiver.js"; +export { RelayReceiver } from "./progress/relay-receiver.js"; export type { McpToolServices } from "./tools/tool-services.js"; diff --git a/packages/mcp/src/mcp/local-server.ts b/packages/mcp/src/mcp/local-server.ts index 31aae03e..cc6abe79 100644 --- a/packages/mcp/src/mcp/local-server.ts +++ b/packages/mcp/src/mcp/local-server.ts @@ -1,4 +1,5 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import type { McpProgressHook } from "../progress/relay-receiver.js"; import type { McpToolServices } from "../tools/tool-services.js"; import type { ToolTermsRemediation } from "../tools/types.js"; import { @@ -54,6 +55,8 @@ export interface CreateLocalMcpServerOptions { authAction?: McpAuthAction; termsRemediation?: ToolTermsRemediation; traceTool?: McpToolExecutionHook; + /** Optional progress source; used only when the client sends a progress token. */ + progress?: McpProgressHook; /** Experimental definitions; defaults to `LOCAL_EXPERIMENTAL_TOOL_DEFINITIONS`. */ experimentalTools?: readonly LocalExperimentalToolDefinition[]; } @@ -81,5 +84,6 @@ export function createLocalMcpServer( authAction: options.authAction, termsRemediation: options.termsRemediation, traceTool: options.traceTool, + progress: options.progress, }); } diff --git a/packages/mcp/src/mcp/progress.test.ts b/packages/mcp/src/mcp/progress.test.ts new file mode 100644 index 00000000..6ae05991 --- /dev/null +++ b/packages/mcp/src/mcp/progress.test.ts @@ -0,0 +1,217 @@ +import { afterEach, describe, expect, it } from "bun:test"; +import type { ProgressRequest } from "@githits/core-internal"; +import { Client } from "@modelcontextprotocol/sdk/client/index.js"; +import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js"; +import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import { + CallToolResultSchema, + ProgressNotificationSchema, +} from "@modelcontextprotocol/sdk/types.js"; +import type { + McpProgressHook, + McpProgressOperation, + McpProgressUpdate, +} from "../progress/relay-receiver.js"; +import type { McpToolServices } from "../tools/tool-services.js"; +import { createDescriptorServices, createMcpServer } from "./server.js"; + +interface FakeHook extends McpProgressHook { + begun: number; + ended: number; + notify?: (update: McpProgressUpdate) => void; +} + +/** Emits two updates while the backend call is in flight. */ +function fakeHook(available = true): FakeHook { + const hook: FakeHook = { + begun: 0, + ended: 0, + begin: async (notify): Promise => { + hook.begun += 1; + if (!available) return undefined; + hook.notify = notify; + return { + headers: { + "X-GitHits-Progress-Id": "op", + "X-GitHits-Relay-Id": "relay", + }, + rejected: () => {}, + end: () => { + hook.ended += 1; + hook.notify = undefined; + }, + }; + }, + }; + return hook; +} + +function services(hook: FakeHook, seen: (ProgressRequest | undefined)[]) { + const backendCall = async (progress?: ProgressRequest) => { + seen.push(progress); + hook.notify?.({ progress: 1, message: "Searching for sources." }); + await new Promise((resolve) => setTimeout(resolve, 5)); + hook.notify?.({ progress: 2, message: "Preparing the answer." }); + }; + return { + ...createDescriptorServices(), + githitsService: { + search: async (_params, options) => { + await backendCall(options?.progress); + return "# Example"; + }, + }, + agenticAskService: { + ask: async (_request, options) => { + await backendCall(options?.progress); + return { display_markdown: "Answer" }; + }, + }, + } as McpToolServices; +} + +const servers: McpServer[] = []; +afterEach(async () => { + for (const server of servers.splice(0)) await server.close(); +}); + +async function connect( + progress: McpProgressHook | undefined, + toolServices: McpToolServices, +) { + const server = createMcpServer({ + metadata: { name: "test", version: "0.0.0" }, + services: toolServices, + ...(progress ? { progress } : {}), + }); + servers.push(server); + const client = new Client({ name: "client", version: "0.0.0" }); + const [clientTransport, serverTransport] = + InMemoryTransport.createLinkedPair(); + await server.connect(serverTransport); + await client.connect(clientTransport); + return client; +} + +const CALLS = [ + ["get_example", { query: "httpx timeout" }], + ["research", { question: "default timeout?", target: "pypi:httpx" }], +] as const; + +describe("MCP progress over a real transport", () => { + it.each(CALLS)( + "%s forwards progress before the final result", + async (name, args) => { + const hook = fakeHook(); + const seen: (ProgressRequest | undefined)[] = []; + const client = await connect(hook, services(hook, seen)); + const updates: { progress: number; message?: string }[] = []; + const result = await client.callTool( + { name, arguments: args }, + CallToolResultSchema, + { + onprogress: (update) => + updates.push({ + progress: update.progress, + message: update.message, + }), + }, + ); + expect(result.isError).toBeFalsy(); + expect(updates).toEqual([ + { progress: 1, message: "Searching for sources." }, + { progress: 2, message: "Preparing the answer." }, + ]); + expect(seen[0]?.headers).toEqual({ + "X-GitHits-Progress-Id": "op", + "X-GitHits-Relay-Id": "relay", + }); + expect(hook.ended).toBe(1); + }, + ); + + it("supports progress token 0", async () => { + const hook = fakeHook(); + const client = await connect(hook, services(hook, [])); + const tokens: unknown[] = []; + client.setNotificationHandler( + ProgressNotificationSchema, + (notification) => { + tokens.push(notification.params.progressToken); + }, + ); + await client.request( + { + method: "tools/call", + params: { + name: "get_example", + arguments: { query: "httpx timeout" }, + _meta: { progressToken: 0 }, + }, + }, + CallToolResultSchema, + ); + expect(tokens).toEqual([0, 0]); + }); + + it("does not engage progress when the client sends no token", async () => { + const hook = fakeHook(); + const seen: (ProgressRequest | undefined)[] = []; + const client = await connect(hook, services(hook, seen)); + const result = await client.callTool({ + name: "get_example", + arguments: { query: "httpx timeout" }, + }); + expect(result.isError).toBeFalsy(); + expect(hook.begun).toBe(0); + expect(seen).toEqual([undefined]); + }); + + it("runs unchanged when progress is unavailable", async () => { + const hook = fakeHook(false); + const seen: (ProgressRequest | undefined)[] = []; + const client = await connect(hook, services(hook, seen)); + const updates: unknown[] = []; + const result = await client.callTool( + { name: "research", arguments: { question: "q?", target: "pypi:httpx" } }, + CallToolResultSchema, + { onprogress: (update) => updates.push(update) }, + ); + expect(result.isError).toBeFalsy(); + expect(hook.begun).toBe(1); + expect(seen).toEqual([undefined]); + expect(updates).toEqual([]); + }); + + it("ends progress when the backend call fails", async () => { + const hook = fakeHook(); + const toolServices = { + ...services(hook, []), + githitsService: { + search: async () => { + throw new Error("Server error (500)"); + }, + }, + } as McpToolServices; + const client = await connect(hook, toolServices); + const result = await client.callTool( + { name: "get_example", arguments: { query: "q" } }, + CallToolResultSchema, + { onprogress: () => {} }, + ); + expect(result.isError).toBe(true); + expect(hook.ended).toBe(1); + }); + + it("leaves servers without a hook unchanged", async () => { + const seen: (ProgressRequest | undefined)[] = []; + const client = await connect(undefined, services(fakeHook(), seen)); + const result = await client.callTool( + { name: "get_example", arguments: { query: "q" } }, + CallToolResultSchema, + { onprogress: () => {} }, + ); + expect(result.isError).toBeFalsy(); + expect(seen).toEqual([undefined]); + }); +}); diff --git a/packages/mcp/src/mcp/server.ts b/packages/mcp/src/mcp/server.ts index d7460410..8d08e9ea 100644 --- a/packages/mcp/src/mcp/server.ts +++ b/packages/mcp/src/mcp/server.ts @@ -1,4 +1,10 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; +import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js"; +import type { + ServerNotification, + ServerRequest, +} from "@modelcontextprotocol/sdk/types.js"; +import type { McpProgressHook } from "../progress/relay-receiver.js"; import { type CompleteToolAnnotations, createCodeDiffTool, @@ -25,6 +31,7 @@ import { type McpAuthAction, withErrorHandling } from "../tools/shared.js"; import type { McpToolServices } from "../tools/tool-services.js"; import type { ToolExecutionContext, + ToolProgress, ToolTermsRemediation, } from "../tools/types.js"; import { buildMcpQuickStart } from "./instructions.js"; @@ -73,6 +80,8 @@ export interface CreateMcpServerOptions { /** @deprecated Use `quickStartOptions`. */ instructionOptions?: Parameters[0]; traceTool?: McpToolExecutionHook; + /** Optional progress source; used only when the client sends a progress token. */ + progress?: McpProgressHook; } export interface McpToolDescriptor { @@ -191,6 +200,7 @@ export function registerMcpTools( termsRemediation?: ToolTermsRemediation; services: McpToolServicesProvider; traceTool?: McpToolExecutionHook; + progress?: McpProgressHook; }, ): void { registerMcpToolsWithFactories(server, STABLE_MCP_TOOL_FACTORIES, { @@ -210,6 +220,7 @@ export function registerMcpToolsWithFactories< termsRemediation?: ToolTermsRemediation; services: McpToolServicesProviderFor; traceTool?: McpToolExecutionHook; + progress?: McpProgressHook; descriptorServices: TServices; }, ): void { @@ -225,10 +236,12 @@ export function registerMcpToolsWithFactories< annotations: descriptor.annotations, }, async (args, extra) => { + const progress = bindProgress(options.progress, extra); const context: ToolExecutionContext = { authAction: options.authAction, termsRemediation: options.termsRemediation, signal: extra?.signal, + ...(progress ? { progress } : {}), }; const runHandler = async () => { const services = await withErrorHandling( @@ -279,6 +292,7 @@ export interface CreateMcpServerWithFactoriesOptions< termsRemediation?: ToolTermsRemediation; instructions?: string; traceTool?: McpToolExecutionHook; + progress?: McpProgressHook; } export function createMcpServerWithFactories< @@ -297,12 +311,43 @@ export function createMcpServerWithFactories< termsRemediation: options.termsRemediation, services: options.services, traceTool: options.traceTool, + progress: options.progress, descriptorServices: options.descriptorServices, }); return server; } +type ProgressRequestExtra = Pick< + RequestHandlerExtra, + "_meta" | "sendNotification" +>; + +/** + * Binds the host's progress source to one MCP request. Progress is offered + * only when the client supplied a token; `0` is a valid token. Notifications + * are best effort and stop when the tool call ends. + */ +function bindProgress( + hook: McpProgressHook | undefined, + extra: ProgressRequestExtra | undefined, +): ToolProgress | undefined { + const progressToken = extra?._meta?.progressToken; + if (!hook || progressToken === undefined || !extra?.sendNotification) { + return undefined; + } + const { sendNotification } = extra; + return { + begin: () => + hook.begin(({ progress, message }) => { + void sendNotification({ + method: "notifications/progress", + params: { progressToken, progress, message }, + }).catch(() => {}); + }), + }; +} + async function resolveMcpToolServices< TServices extends McpToolServices, TExtra, diff --git a/packages/mcp/src/progress/relay-receiver.test.ts b/packages/mcp/src/progress/relay-receiver.test.ts new file mode 100644 index 00000000..5ec60b68 --- /dev/null +++ b/packages/mcp/src/progress/relay-receiver.test.ts @@ -0,0 +1,409 @@ +import { afterEach, describe, expect, it, mock } from "bun:test"; +import { + type McpProgressUpdate, + type RelayAuth, + RelayReceiver, + type RelaySocket, +} from "./relay-receiver.js"; + +const RELAY_ID = "1b2c3d4e-5f60-4718-8a9b-0c1d2e3f4a5b"; + +class FakeSocket implements RelaySocket { + onopen: (() => void) | null = null; + onmessage: ((event: { data: unknown }) => void) | null = null; + onclose: (() => void) | null = null; + onerror: (() => void) | null = null; + closed = false; + + constructor(readonly url: string) {} + + open(): void { + this.onopen?.(); + } + + emit(event: Record): void { + this.onmessage?.({ data: JSON.stringify({ version: 1, ...event }) }); + } + + drop(): void { + this.onclose?.(); + } + + close(): void { + this.closed = true; + } +} + +interface Harness { + receiver: RelayReceiver; + sockets: FakeSocket[]; + setupCalls: { path: string; headers: Headers; body: string }[]; + responses: Response[]; + clock: { now: number }; +} + +const receivers: RelayReceiver[] = []; + +afterEach(() => { + for (const receiver of receivers.splice(0)) receiver.close(); +}); + +function access(seconds = 900, url = "wss://relay.test/client"): Response { + return Response.json({ + relay_id: RELAY_ID, + expires_at: Math.floor(Date.now() / 1000) + seconds, + connection_url: url, + }); +} + +function harness( + options: { + auth?: RelayAuth; + responses?: Response[]; + autoOpen?: boolean; + connectWaitMs?: number; + idleCloseMs?: number; + minRenewDelayMs?: number; + } = {}, +): Harness { + const sockets: FakeSocket[] = []; + const setupCalls: Harness["setupCalls"] = []; + const responses = options.responses ?? [access()]; + const clock = { now: Date.now() }; + let ids = 0; + const receiver = new RelayReceiver({ + apiUrl: "https://api.githits.test/", + auth: options.auth ?? { + kind: "caller", + getToken: async () => "gh-proj-token", + }, + fetchFn: (async (url: string | URL | Request, init?: RequestInit) => { + setupCalls.push({ + path: new URL(String(url)).pathname, + headers: new Headers(init?.headers), + body: String(init?.body ?? ""), + }); + const next = responses.shift(); + if (!next) throw new TypeError("network down"); + return next; + }) as unknown as typeof fetch, + createSocket: (url) => { + const socket = new FakeSocket(url); + sockets.push(socket); + if (options.autoOpen !== false) queueMicrotask(() => socket.open()); + return socket; + }, + connectWaitMs: options.connectWaitMs ?? 50, + firstConnectWaitMs: options.connectWaitMs ?? 50, + idleCloseMs: options.idleCloseMs ?? 60_000, + reconnectBaseMs: 5, + minRenewDelayMs: options.minRenewDelayMs ?? 60_000, + now: () => clock.now, + createId: () => `00000000-0000-4000-8000-00000000000${ids++}`, + }); + receivers.push(receiver); + return { receiver, sockets, setupCalls, responses, clock }; +} + +function collector(): { + notify: (update: McpProgressUpdate) => void; + updates: McpProgressUpdate[]; +} { + const updates: McpProgressUpdate[] = []; + return { notify: (update) => updates.push(update), updates }; +} + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +describe("RelayReceiver", () => { + it("mints with the caller token and routes events to the matching call", async () => { + const { receiver, sockets, setupCalls } = harness(); + const first = collector(); + const second = collector(); + const a = await receiver.begin(first.notify); + const b = await receiver.begin(second.notify); + expect(a?.headers).toEqual({ + "X-GitHits-Progress-Id": "00000000-0000-4000-8000-000000000000", + "X-GitHits-Relay-Id": RELAY_ID, + }); + expect(setupCalls).toHaveLength(1); + expect(setupCalls[0]?.path).toBe("/progress/relays"); + expect(setupCalls[0]?.headers.get("Authorization")).toBe( + "Bearer gh-proj-token", + ); + expect(sockets).toHaveLength(1); + + const socket = sockets[0] as FakeSocket; + socket.emit({ + operation_id: a?.headers["X-GitHits-Progress-Id"], + sequence: 1, + message: "Searching for sources.", + }); + socket.emit({ + operation_id: b?.headers["X-GitHits-Progress-Id"]?.toUpperCase(), + sequence: 2, + message: "Selecting relevant sources.", + }); + socket.emit({ operation_id: "unknown", sequence: 3, message: "x" }); + expect(first.updates).toEqual([ + { progress: 1, message: "Searching for sources." }, + ]); + expect(second.updates).toEqual([ + { progress: 2, message: "Selecting relevant sources." }, + ]); + }); + + it("drops duplicate, stale and malformed events and stops after end", async () => { + const { receiver, sockets } = harness(); + const { notify, updates } = collector(); + const operation = await receiver.begin(notify); + const socket = sockets[0] as FakeSocket; + const operation_id = operation?.headers["X-GitHits-Progress-Id"]; + socket.emit({ operation_id, sequence: 3, message: "Three" }); + socket.emit({ operation_id, sequence: 3, message: "Duplicate" }); + socket.emit({ operation_id, sequence: 2, message: "Stale" }); + socket.onmessage?.({ data: "not json" }); + socket.emit({ operation_id, sequence: 0, message: "Zero" }); + socket.onmessage?.({ + data: JSON.stringify({ + version: 2, + operation_id, + sequence: 9, + message: "v2", + }), + }); + operation?.end(); + socket.emit({ operation_id, sequence: 4, message: "After end" }); + expect(updates).toEqual([{ progress: 3, message: "Three" }]); + }); + + it("decodes binary frames and caps message length", async () => { + const { receiver, sockets } = harness(); + const { notify, updates } = collector(); + const operation = await receiver.begin(notify); + const payload = JSON.stringify({ + version: 1, + operation_id: operation?.headers["X-GitHits-Progress-Id"], + sequence: 1, + message: "x".repeat(500), + }); + sockets[0]?.onmessage?.({ data: new TextEncoder().encode(payload) }); + expect(updates[0]?.message).toHaveLength(200); + }); + + it("sends service credentials on setup and on each call", async () => { + const { receiver, setupCalls } = harness({ + auth: { + kind: "service", + serviceId: "example-service", + key: "k".repeat(48), + }, + }); + const operation = await receiver.begin(() => {}); + expect(setupCalls[0]?.headers.get("X-GitHits-Relay-Service")).toBe( + "example-service", + ); + expect(setupCalls[0]?.headers.has("Authorization")).toBe(false); + expect(operation?.headers["X-GitHits-Relay-Service"]).toBe( + "example-service", + ); + expect(operation?.headers["X-GitHits-Relay-Key"]).toBe("k".repeat(48)); + }); + + it.each([ + ["unavailable", new Response(null, { status: 503 }), 5 * 60_000], + [ + "rate limited", + new Response(null, { status: 429, headers: { "Retry-After": "7" } }), + 7_000, + ], + ])( + "backs off after a %s mint without retrying per call", + async (_name, failure, backoff) => { + const { receiver, setupCalls, responses, clock } = harness({ + responses: [failure], + }); + expect(await receiver.begin(() => {})).toBeUndefined(); + expect(await receiver.begin(() => {})).toBeUndefined(); + expect(setupCalls).toHaveLength(1); + responses.push(access()); + clock.now += backoff + 1; + expect(await receiver.begin(() => {})).toBeDefined(); + expect(setupCalls).toHaveLength(2); + }, + ); + + it("returns no progress when the caller has no token", async () => { + const { receiver, setupCalls } = harness({ + auth: { kind: "caller", getToken: async () => undefined }, + }); + expect(await receiver.begin(() => {})).toBeUndefined(); + expect(setupCalls).toHaveLength(0); + }); + + it("proceeds after the connect wait while the socket is still opening", async () => { + const { receiver, sockets } = harness({ + autoOpen: false, + connectWaitMs: 20, + }); + const started = Date.now(); + const operation = await receiver.begin(() => {}); + expect(Date.now() - started).toBeGreaterThanOrEqual(15); + expect(operation?.headers["X-GitHits-Relay-Id"]).toBe(RELAY_ID); + expect(sockets).toHaveLength(1); + }); + + it("backs off when the socket cannot open", async () => { + const { receiver, sockets, setupCalls } = harness({ autoOpen: false }); + const pending = receiver.begin(() => {}); + await sleep(0); + sockets[0]?.drop(); + await pending; + expect(await receiver.begin(() => {})).toBeUndefined(); + expect(setupCalls).toHaveLength(1); + }); + + it("reconnects through renew with the same relay after a drop mid-call", async () => { + const { receiver, sockets, setupCalls, responses } = harness(); + const { notify, updates } = collector(); + const operation = await receiver.begin(notify); + const operation_id = operation?.headers["X-GitHits-Progress-Id"]; + responses.push(access(900, "wss://relay.test/renewed")); + sockets[0]?.drop(); + await sleep(40); + expect(setupCalls.map((call) => call.path)).toEqual([ + "/progress/relays", + "/progress/relays/renew", + ]); + expect(JSON.parse(setupCalls[1]?.body ?? "{}")).toEqual({ + relay_id: RELAY_ID, + }); + expect(sockets).toHaveLength(2); + sockets[1]?.emit({ operation_id, sequence: 4, message: "Resumed" }); + expect(updates).toEqual([{ progress: 4, message: "Resumed" }]); + }); + + it("renews before expiry and swaps sockets only after the new one opens", async () => { + const { receiver, sockets, setupCalls, responses } = harness({ + minRenewDelayMs: 10, + responses: [access(30)], + }); + const { notify, updates } = collector(); + const operation = await receiver.begin(notify); + const operation_id = operation?.headers["X-GitHits-Progress-Id"]; + responses.push(access(900, "wss://relay.test/renewed")); + await sleep(60); + expect(setupCalls[1]?.path).toBe("/progress/relays/renew"); + expect(sockets).toHaveLength(2); + expect(sockets[0]?.closed).toBe(true); + sockets[1]?.emit({ operation_id, sequence: 1, message: "After renewal" }); + expect(updates).toEqual([{ progress: 1, message: "After renewal" }]); + }); + + it("keeps the working socket when a replacement fails to open", async () => { + const sockets: FakeSocket[] = []; + const responses = [access(30), access(900, "wss://relay.test/broken")]; + const receiver = new RelayReceiver({ + apiUrl: "https://api.githits.test", + auth: { kind: "caller", getToken: async () => "token" }, + fetchFn: (async () => + responses.shift() ?? + new Response(null, { status: 503 })) as unknown as typeof fetch, + createSocket: (url) => { + const socket = new FakeSocket(url); + sockets.push(socket); + queueMicrotask(() => + sockets.length === 1 ? socket.open() : socket.drop(), + ); + return socket; + }, + connectWaitMs: 50, + minRenewDelayMs: 10, + reconnectBaseMs: 60_000, + }); + receivers.push(receiver); + const { notify, updates } = collector(); + const operation = await receiver.begin(notify); + await sleep(60); + expect(sockets).toHaveLength(2); + expect(sockets[0]?.closed).toBe(false); + sockets[0]?.emit({ + operation_id: operation?.headers["X-GitHits-Progress-Id"], + sequence: 1, + message: "Still flowing", + }); + expect(updates).toHaveLength(1); + }); + + it("stops offering progress after a backend rejection", async () => { + const { receiver, clock } = harness(); + const operation = await receiver.begin(() => {}); + operation?.rejected("unavailable"); + operation?.end(); + expect(await receiver.begin(() => {})).toBeUndefined(); + clock.now += 5 * 60_000 + 1; + expect(await receiver.begin(() => {})).toBeDefined(); + }); + + it("closes an idle socket and reconnects on demand", async () => { + const { receiver, sockets, setupCalls, responses } = harness({ + idleCloseMs: 10, + }); + const operation = await receiver.begin(() => {}); + operation?.end(); + await sleep(30); + expect(sockets[0]?.closed).toBe(true); + responses.push(access()); + expect(await receiver.begin(() => {})).toBeDefined(); + expect(setupCalls.map((call) => call.path)).toEqual([ + "/progress/relays", + "/progress/relays/renew", + ]); + }); + + it("close() shuts the socket and disables progress", async () => { + const { receiver, sockets } = harness(); + await receiver.begin(() => {}); + receiver.close(); + expect(sockets[0]?.closed).toBe(true); + expect(await receiver.begin(() => {})).toBeUndefined(); + }); + + it("never lets a throwing notifier affect delivery", async () => { + const { receiver, sockets } = harness(); + const operation = await receiver.begin( + mock(() => { + throw new Error("client gone"); + }), + ); + expect(() => + sockets[0]?.emit({ + operation_id: operation?.headers["X-GitHits-Progress-Id"], + sequence: 1, + message: "x", + }), + ).not.toThrow(); + }); + it("waits longer for the first relay than for later calls", async () => { + const sockets: FakeSocket[] = []; + const receiver = new RelayReceiver({ + apiUrl: "https://api.githits.test", + auth: { kind: "caller", getToken: async () => "token" }, + fetchFn: (async () => { + await sleep(40); + return access(); + }) as unknown as typeof fetch, + createSocket: (url) => { + const socket = new FakeSocket(url); + sockets.push(socket); + queueMicrotask(() => socket.open()); + return socket; + }, + connectWaitMs: 5, + firstConnectWaitMs: 500, + }); + receivers.push(receiver); + const operation = await receiver.begin(() => {}); + expect(operation?.headers["X-GitHits-Relay-Id"]).toBe(RELAY_ID); + }); +}); diff --git a/packages/mcp/src/progress/relay-receiver.ts b/packages/mcp/src/progress/relay-receiver.ts new file mode 100644 index 00000000..600934a4 --- /dev/null +++ b/packages/mcp/src/progress/relay-receiver.ts @@ -0,0 +1,485 @@ +import type { ProgressRequest } from "@githits/core-internal"; +import WebSocket from "ws"; + +/** One MCP progress update derived from a backend progress event. */ +export interface McpProgressUpdate { + /** Strictly increasing per operation; suitable for MCP `progress`. */ + progress: number; + /** Fixed, server-allowlisted status text. */ + message: string; +} + +/** One correlated backend call; `end()` stops delivery before the result returns. */ +export interface McpProgressOperation extends ProgressRequest { + end(): void; +} + +/** + * Optional, credential-free progress source for long-running tools. + * + * `begin` returns `undefined` whenever progress is unavailable; tools then run + * exactly as without a hook. The HTTP result is always authoritative. + */ +export interface McpProgressHook { + begin( + notify: (update: McpProgressUpdate) => void, + ): Promise; +} + +/** + * How the receiver authenticates relay setup. + * + * - `caller`: the caller's own GitHits token (local MCP); progress reaches only + * that caller's receiver. + * - `service`: a trusted service credential for a server that multiplexes + * many callers over one socket. The credential is supplied at runtime and + * never ships in this package. + */ +export type RelayAuth = + | { kind: "caller"; getToken: () => Promise } + | { kind: "service"; serviceId: string; key: string }; + +/** The subset of a WebSocket the receiver uses; `ws` and WHATWG sockets fit. */ +export interface RelaySocket { + onopen: (() => void) | null; + onmessage: ((event: { data: unknown }) => void) | null; + onclose: (() => void) | null; + onerror: (() => void) | null; + close(): void; +} + +export interface RelayReceiverOptions { + apiUrl: string; + auth: RelayAuth; + fetchFn?: typeof fetch; + createSocket?: (url: string) => RelaySocket; + /** Extra headers for relay setup requests, such as client attribution. */ + clientHeaders?: () => Record; + /** Longest a call waits for an existing relay's socket before proceeding. */ + connectWaitMs?: number; + /** + * Longest a call waits while the first relay is minted and its socket opens. + * Minting takes about a second, so a shorter wait would skip progress for the + * first call of every process. + */ + firstConnectWaitMs?: number; + /** Close an unused socket after this long. */ + idleCloseMs?: number; + /** First reconnect delay; doubles per attempt up to 30 s. */ + reconnectBaseMs?: number; + /** Earliest renewal after a mint, even for short-lived receive URLs. */ + minRenewDelayMs?: number; + now?: () => number; + createId?: () => string; +} + +interface RelayAccess { + relayId: string; + url: string; + /** Epoch milliseconds. */ + expiresAt: number; +} + +interface Route { + notify: (update: McpProgressUpdate) => void; + lastSequence: number; +} + +class RelayUnavailableError extends Error { + constructor(readonly backoffMs: number) { + super("Progress relay unavailable"); + } +} + +const SETUP_TIMEOUT_MS = 5_000; +const RENEW_BEFORE_EXPIRY_MS = 60_000; +const MIN_RENEW_DELAY_MS = 10_000; +const RECONNECT_BASE_MS = 1_000; +const RECONNECT_MAX_MS = 30_000; +const NETWORK_BACKOFF_MS = 30_000; +const UNAVAILABLE_BACKOFF_MS = 5 * 60_000; +const MAX_MESSAGE_LENGTH = 200; + +/** + * Receives backend operation progress over one shared WebSocket and routes it + * to the MCP call that supplied each operation ID. + * + * Every failure degrades to "no progress": relay setup errors back off, socket + * drops reconnect in the background, and late or unknown events are dropped. + */ +export class RelayReceiver implements McpProgressHook { + private readonly fetchFn: typeof fetch; + private readonly createSocket: (url: string) => RelaySocket; + private readonly connectWaitMs: number; + private readonly firstConnectWaitMs: number; + private readonly idleCloseMs: number; + private readonly reconnectBaseMs: number; + private readonly minRenewDelayMs: number; + private readonly now: () => number; + private readonly createId: () => string; + private readonly routes = new Map(); + private relay: RelayAccess | undefined; + private socket: RelaySocket | undefined; + private connecting: Promise | undefined; + private unavailableUntil = 0; + private reconnectAttempts = 0; + private renewTimer: ReturnType | undefined; + private reconnectTimer: ReturnType | undefined; + private idleTimer: ReturnType | undefined; + private closed = false; + + constructor(private readonly options: RelayReceiverOptions) { + this.fetchFn = options.fetchFn ?? globalThis.fetch; + this.createSocket = + options.createSocket ?? + ((url) => new WebSocket(url) as unknown as RelaySocket); + this.connectWaitMs = options.connectWaitMs ?? 1_000; + this.firstConnectWaitMs = options.firstConnectWaitMs ?? 3_000; + this.idleCloseMs = options.idleCloseMs ?? 5 * 60_000; + this.reconnectBaseMs = options.reconnectBaseMs ?? RECONNECT_BASE_MS; + this.minRenewDelayMs = options.minRenewDelayMs ?? MIN_RENEW_DELAY_MS; + this.now = options.now ?? Date.now; + this.createId = options.createId ?? (() => globalThis.crypto.randomUUID()); + } + + async begin( + notify: (update: McpProgressUpdate) => void, + ): Promise { + if (this.closed || this.now() < this.unavailableUntil) return undefined; + this.clearTimer("idleTimer"); + const wait = this.relay ? this.connectWaitMs : this.firstConnectWaitMs; + const ready = this.ensureConnected(); + let waitTimer: ReturnType | undefined; + await Promise.race([ + ready, + new Promise((resolve) => { + waitTimer = setTimeout(resolve, wait); + }), + ]); + clearTimeout(waitTimer); + const relay = this.relay; + if (!relay || this.closed) { + this.scheduleIdleClose(); + return undefined; + } + const operationId = this.createId().toLowerCase(); + this.routes.set(operationId, { notify, lastSequence: 0 }); + let ended = false; + return { + headers: { + "X-GitHits-Progress-Id": operationId, + "X-GitHits-Relay-Id": relay.relayId, + ...this.serviceHeaders(), + }, + rejected: () => { + // Rejections happen before execution; stop offering progress for a while. + this.unavailableUntil = this.now() + UNAVAILABLE_BACKOFF_MS; + }, + end: () => { + if (ended) return; + ended = true; + this.routes.delete(operationId); + this.scheduleIdleClose(); + }, + }; + } + + /** Stops all timers and closes the socket; later calls get no progress. */ + close(): void { + this.closed = true; + this.routes.clear(); + this.clearTimer("renewTimer"); + this.clearTimer("reconnectTimer"); + this.clearTimer("idleTimer"); + this.closeSocket(); + } + + private ensureConnected(): Promise { + if (this.socket && this.relay && this.now() < this.relay.expiresAt) { + return Promise.resolve(); + } + this.connecting ??= this.connect().finally(() => { + this.connecting = undefined; + }); + return this.connecting; + } + + private async connect(): Promise { + let relay: RelayAccess; + try { + relay = await this.setup(this.relay?.relayId); + } catch (error) { + this.unavailableUntil = + this.now() + + (error instanceof RelayUnavailableError + ? error.backoffMs + : NETWORK_BACKOFF_MS); + return; + } + if (this.closed) return; + this.relay = relay; + this.scheduleRenewal(relay); + await this.openSocket(relay.url); + } + + /** Resolves once the socket opens or fails; an open socket replaces the old one. */ + private openSocket(url: string): Promise { + return new Promise((resolve) => { + let socket: RelaySocket; + try { + socket = this.createSocket(url); + } catch { + this.socketFailed(); + resolve(); + return; + } + let opened = false; + // Deliver from the replacement immediately; duplicates are dropped by sequence. + socket.onmessage = (event) => this.deliver(event.data); + socket.onerror = () => {}; + socket.onopen = () => { + opened = true; + this.reconnectAttempts = 0; + // Swap only after the new socket opens so renewal never leaves a gap + // and a failed replacement never orphans the working socket. + const previous = this.socket; + this.socket = socket; + if (previous) this.detach(previous); + if (this.closed) this.closeSocket(); + resolve(); + }; + socket.onclose = () => { + if (!opened) { + this.detach(socket); + if (!this.socket) this.socketFailed(); + resolve(); + return; + } + if (this.socket !== socket) return; + this.socket = undefined; + if (!this.closed && this.routes.size > 0) this.scheduleReconnect(); + }; + }); + } + + /** A socket that cannot open (e.g. a proxy blocking WebSockets) backs off. */ + private socketFailed(): void { + const delay = Math.min( + RECONNECT_MAX_MS, + this.reconnectBaseMs * 2 ** this.reconnectAttempts, + ); + this.unavailableUntil = Math.max(this.unavailableUntil, this.now() + delay); + if (this.routes.size > 0) this.scheduleReconnect(); + else this.reconnectAttempts += 1; + } + + private deliver(data: unknown): void { + const event = parseEvent(data); + if (!event) return; + const route = this.routes.get(event.operationId); + // Overlapping sockets during renewal can duplicate events; MCP progress + // must strictly increase, so drop anything not newer than the last one. + if (!route || event.sequence <= route.lastSequence) return; + route.lastSequence = event.sequence; + try { + route.notify({ progress: event.sequence, message: event.message }); + } catch { + // A failed notification must never affect the tool call. + } + } + + private async setup(relayId: string | undefined): Promise { + const path = relayId ? "/progress/relays/renew" : "/progress/relays"; + const headers: Record = { + ...this.options.clientHeaders?.(), + ...(await this.authHeaders()), + "Content-Type": "application/json", + }; + const response = await this.fetchFn( + `${this.options.apiUrl.replace(/\/+$/, "")}${path}`, + { + method: "POST", + headers, + body: relayId ? JSON.stringify({ relay_id: relayId }) : undefined, + signal: AbortSignal.timeout(SETUP_TIMEOUT_MS), + }, + ); + if (!response.ok) { + await response.body?.cancel().catch(() => {}); + throw new RelayUnavailableError(backoffFor(response)); + } + const access = parseRelayAccess(await response.json()); + if (!access) throw new RelayUnavailableError(UNAVAILABLE_BACKOFF_MS); + return access; + } + + private async authHeaders(): Promise> { + const { auth } = this.options; + if (auth.kind === "service") return this.serviceHeaders(); + const token = await auth.getToken(); + if (!token) throw new RelayUnavailableError(UNAVAILABLE_BACKOFF_MS); + return { Authorization: `Bearer ${token}` }; + } + + private serviceHeaders(): Record { + const { auth } = this.options; + return auth.kind === "service" + ? { + "X-GitHits-Relay-Service": auth.serviceId, + "X-GitHits-Relay-Key": auth.key, + } + : {}; + } + + private scheduleRenewal(relay: RelayAccess): void { + this.clearTimer("renewTimer"); + const delay = Math.max( + this.minRenewDelayMs, + relay.expiresAt - RENEW_BEFORE_EXPIRY_MS - this.now(), + ); + this.renewTimer = setTimeout(() => { + this.renewTimer = undefined; + if (this.routes.size === 0) { + // Nothing to deliver: let the next call reconnect on demand. + this.closeSocket(); + return; + } + void this.ensureFresh(); + }, delay); + unrefTimer(this.renewTimer); + } + + /** Mint a fresh receive URL for the same relay while the old socket stays up. */ + private ensureFresh(): Promise { + this.connecting ??= this.connect().finally(() => { + this.connecting = undefined; + }); + return this.connecting; + } + + private scheduleReconnect(): void { + if (this.reconnectTimer || this.closed) return; + const delay = Math.min( + RECONNECT_MAX_MS, + this.reconnectBaseMs * 2 ** this.reconnectAttempts, + ); + this.reconnectAttempts += 1; + this.reconnectTimer = setTimeout(() => { + this.reconnectTimer = undefined; + if (this.closed || this.routes.size === 0) return; + // Reconnects ignore the begin() backoff: in-flight calls still want events. + void this.ensureFresh().then(() => { + if (!this.socket && this.routes.size > 0) this.scheduleReconnect(); + }); + }, delay); + unrefTimer(this.reconnectTimer); + } + + private scheduleIdleClose(): void { + if (this.routes.size > 0 || this.closed) return; + this.clearTimer("idleTimer"); + this.idleTimer = setTimeout(() => { + this.idleTimer = undefined; + if (this.routes.size === 0) { + this.clearTimer("renewTimer"); + this.closeSocket(); + } + }, this.idleCloseMs); + unrefTimer(this.idleTimer); + } + + private closeSocket(): void { + const socket = this.socket; + this.socket = undefined; + if (socket) this.detach(socket); + } + + private detach(socket: RelaySocket): void { + socket.onopen = null; + socket.onmessage = null; + socket.onclose = null; + socket.onerror = () => {}; + try { + socket.close(); + } catch { + // Already closed. + } + } + + private clearTimer( + name: "renewTimer" | "reconnectTimer" | "idleTimer", + ): void { + const timer = this[name]; + if (timer) clearTimeout(timer); + this[name] = undefined; + } +} + +function backoffFor(response: Response): number { + if (response.status === 429) { + const seconds = Number(response.headers.get("Retry-After")); + if (Number.isFinite(seconds) && seconds > 0) return seconds * 1_000; + } + return UNAVAILABLE_BACKOFF_MS; +} + +function parseRelayAccess(value: unknown): RelayAccess | undefined { + if (!value || typeof value !== "object") return undefined; + const body = value as Record; + const { relay_id, connection_url, expires_at } = body; + if ( + typeof relay_id !== "string" || + typeof connection_url !== "string" || + !connection_url.startsWith("wss://") || + typeof expires_at !== "number" + ) { + return undefined; + } + return { + relayId: relay_id.toLowerCase(), + url: connection_url, + expiresAt: expires_at * 1_000, + }; +} + +interface RelayEvent { + operationId: string; + sequence: number; + message: string; +} + +function parseEvent(data: unknown): RelayEvent | undefined { + let text: string; + if (typeof data === "string") text = data; + else if (data instanceof ArrayBuffer || data instanceof Uint8Array) + text = new TextDecoder().decode(data); + else return undefined; + let value: unknown; + try { + value = JSON.parse(text); + } catch { + return undefined; + } + if (!value || typeof value !== "object") return undefined; + const event = value as Record; + if ( + event.version !== 1 || + typeof event.operation_id !== "string" || + typeof event.message !== "string" || + typeof event.sequence !== "number" || + !Number.isInteger(event.sequence) || + event.sequence < 1 + ) { + return undefined; + } + return { + operationId: event.operation_id.toLowerCase(), + sequence: event.sequence, + message: event.message.slice(0, MAX_MESSAGE_LENGTH), + }; +} + +function unrefTimer(timer: ReturnType): void { + // Background maintenance must not keep a CLI process alive. + (timer as { unref?: () => void }).unref?.(); +} diff --git a/packages/mcp/src/tools/get-example.ts b/packages/mcp/src/tools/get-example.ts index 99b77a01..28a0a984 100644 --- a/packages/mcp/src/tools/get-example.ts +++ b/packages/mcp/src/tools/get-example.ts @@ -1,3 +1,4 @@ +import type { ProgressRequest } from "@githits/core-internal"; import { z } from "zod"; import { extractSolutionId } from "../shared/extract-solution-id.js"; import { GET_EXAMPLE_GUARDRAIL } from "./guardrails.js"; @@ -29,6 +30,8 @@ export interface GetExampleSearchParams { /** Browser-standard request controls for example search. */ export interface GetExampleRequestOptions { signal?: AbortSignal; + /** Optional best-effort progress correlation; never changes the result. */ + progress?: ProgressRequest; } /** Minimal injected service contract used by the example tool. */ @@ -92,9 +95,19 @@ export function createGetExampleTool( licenseMode: args.license_mode, includeExplanation: false, }; - const markdown = context?.signal - ? await service.search(searchParams, { signal: context.signal }) - : await service.search(searchParams); + const progress = await context?.progress?.begin(); + let markdown: string; + try { + markdown = + context?.signal || progress + ? await service.search(searchParams, { + ...(context?.signal ? { signal: context.signal } : {}), + ...(progress ? { progress } : {}), + }) + : await service.search(searchParams); + } finally { + progress?.end(); + } const solutionId = extractSolutionId(markdown); const payload = solutionId ? { result: markdown, solution_id: solutionId } diff --git a/packages/mcp/src/tools/research.ts b/packages/mcp/src/tools/research.ts index 51e53f82..38dbcd2a 100644 --- a/packages/mcp/src/tools/research.ts +++ b/packages/mcp/src/tools/research.ts @@ -89,10 +89,15 @@ export function createResearchTool( args.format, ); } + const progress = await context?.progress?.begin(); try { - const requestOptions = context?.signal - ? { signal: context.signal } - : undefined; + const requestOptions = + context?.signal || progress + ? { + ...(context?.signal ? { signal: context.signal } : {}), + ...(progress ? { progress } : {}), + } + : undefined; const response = args.source_format === "url" ? await service.ask( @@ -145,6 +150,8 @@ export function createResearchTool( ...(failure.threadId ? { thread_id: failure.threadId } : {}), }), ); + } finally { + progress?.end(); } }, }; diff --git a/packages/mcp/src/tools/types.ts b/packages/mcp/src/tools/types.ts index 502a1f0f..129b6afa 100644 --- a/packages/mcp/src/tools/types.ts +++ b/packages/mcp/src/tools/types.ts @@ -1,3 +1,4 @@ +import type { ProgressRequest } from "@githits/core-internal"; import type { z } from "zod"; /** Annotation fields required by OpenAI's MCP marketplace validation. */ @@ -38,11 +39,24 @@ export interface ToolTermsRemediation { action: string; } +/** One correlated backend call; `end()` must run before the tool returns. */ +export interface ToolProgressOperation extends ProgressRequest { + end(): void; +} + +/** Progress already bound to the current MCP request's progress token. */ +export interface ToolProgress { + /** Returns `undefined` when progress is unavailable; run without it then. */ + begin(): Promise; +} + /** Host-provided execution state shared by MCP and direct tool callers. */ export interface ToolExecutionContext { authAction?: McpAuthAction; termsRemediation?: ToolTermsRemediation; signal?: AbortSignal; + /** Present only when the host enabled progress and the client asked for it. */ + progress?: ToolProgress; } /** diff --git a/src/commands/mcp.ts b/src/commands/mcp.ts index 99754de8..41740a8a 100644 --- a/src/commands/mcp.ts +++ b/src/commands/mcp.ts @@ -5,6 +5,7 @@ import { highlight, type LocalExperimentalMcpPolicy, type LocalMcpToolServices, + RelayReceiver, shouldUseColors, type ToolTermsRemediation, } from "@githits/mcp/internal"; @@ -38,6 +39,8 @@ export interface StartMcpServerOptions { onServerCreated?: (server: LocalMcpServer) => void; experimentalPolicy?: LocalExperimentalMcpPolicy; termsRemediation?: ToolTermsRemediation; + /** Best-effort progress for long tools; closed when the transport closes. */ + progress?: RelayReceiver; } export interface CreateMcpCommandStartupOptions { @@ -61,8 +64,11 @@ export async function startMcpServer( metadata: LOCAL_MCP_SERVER_METADATA, policy: options.experimentalPolicy ?? DISABLED_LOCAL_MCP_POLICY, termsRemediation: options.termsRemediation ?? LOCAL_TERMS_REMEDIATION, + progress: options.progress, }); const transport = new StdioServerTransport(); + // An open progress socket must not keep the process alive after the client leaves. + server.server.onclose = () => options.progress?.close(); options.onServerCreated?.(server); @@ -97,6 +103,7 @@ function readMcpClientVersion( export interface McpCommandStartup { services: LocalMcpToolServices; + progress?: RelayReceiver; experimentalPolicy: LocalExperimentalMcpPolicy; onServerCreated: (server: LocalMcpServer) => void; } @@ -144,8 +151,15 @@ export async function createMcpCommandStartup( } catch { warn("GitHits MCP skill maintenance failed; continuing startup."); } + // Caller relay: authorized by the user's own token, which is resolved lazily + // on the first tool call that asks for progress, never at startup. + const progress = new RelayReceiver({ + apiUrl: services.apiUrl, + auth: { kind: "caller", getToken: () => services.tokenProvider.getToken() }, + }); return { services, + progress, experimentalPolicy, onServerCreated: (created: LocalMcpServer) => { server = created; @@ -220,6 +234,7 @@ Authenticated tool calls require a valid GitHits token.`, await startServer(startup.services, { experimentalPolicy: startup.experimentalPolicy, onServerCreated: startup.onServerCreated, + progress: startup.progress, }); }); @@ -240,6 +255,7 @@ in MCP configuration files. Use 'githits mcp' for interactive setup.`, await startServer(startup.services, { experimentalPolicy: startup.experimentalPolicy, onServerCreated: startup.onServerCreated, + progress: startup.progress, }); }); }