copilot-relay 0.3.8 → 0.3.9

package/dist/main.js

@@ -1,6 +1,7 @@
 #!/usr/bin/env node
 import { defineCommand, runMain } from "citty";
 import fs from "node:fs/promises";
+import { createHash, randomUUID } from "node:crypto";
 import { inspect, promisify } from "node:util";
 import consola from "consola";
 import os from "node:os";
@@ -9,7 +10,6 @@ import { fileURLToPath } from "node:url";
 import { execFile } from "node:child_process";
 import { readFileSync } from "node:fs";
 import { createAdaptorServer } from "@hono/node-server";
-import { createHash, randomUUID } from "node:crypto";
 import { Agent, fetch as fetch$1 } from "undici";
 import { events } from "fetch-event-stream";
 import { Hono } from "hono";
@@ -407,7 +407,16 @@ const writeStoredCopilotToken = async (input) => {
 		refreshIn: input.refresh_in,
 		token: input.token
 	};
-	await fs.writeFile(paths.copilotTokenPath, `${JSON.stringify(payload, null, 2)}\n`, { mode: 384 });
+	const temporaryPath = `${paths.copilotTokenPath}.${randomUUID()}.tmp`;
+	try {
+		await fs.writeFile(temporaryPath, `${JSON.stringify(payload, null, 2)}\n`, {
+			mode: 384,
+			flag: "wx"
+		});
+		await fs.rename(temporaryPath, paths.copilotTokenPath);
+	} finally {
+		await fs.rm(temporaryPath, { force: true });
+	}
 };
 const getStoredTokenRemainingSeconds = (token) => {
 	const ageSeconds = Math.floor((Date.now() - token.refreshedAt) / 1e3);
@@ -460,8 +469,11 @@ const getGitHubUser = async (githubToken, vsCodeVersion) => {
 	if (!response.ok) throw new HTTPError("Failed to get GitHub user", response);
 	return await response.json();
 };
-const getCopilotToken = async (githubToken, vsCodeVersion) => {
-	const response = await fetch(`${githubApiBaseUrl}/copilot_internal/v2/token`, { headers: githubHeaders(githubToken, vsCodeVersion) });
+const getCopilotToken = async (githubToken, vsCodeVersion, signal) => {
+	const response = await fetch(`${githubApiBaseUrl}/copilot_internal/v2/token`, {
+		headers: githubHeaders(githubToken, vsCodeVersion),
+		signal
+	});
 	if (!response.ok) throw new HTTPError("Failed to get Copilot token", response);
 	return await response.json();
 };
@@ -489,11 +501,21 @@ const loadGitHubLogin = async (githubToken, vsCodeVersion) => {
 };
 const setupProxyAuth = async (config, options = {}) => {
 	let githubToken = await ensureGitHubToken(options);
-	const applyCopilotToken = async () => {
-		const tokenResponse = await getCopilotToken(githubToken, config.vsCodeVersion);
-		config.copilotToken = tokenResponse.token;
-		await writeStoredCopilotToken(tokenResponse);
-		return tokenResponse.refresh_in;
+	let refreshInFlight;
+	const applyCopilotToken = () => {
+		if (refreshInFlight) return refreshInFlight;
+		refreshInFlight = (async () => {
+			const tokenResponse = await getCopilotToken(githubToken, config.vsCodeVersion, AbortSignal.timeout(config.upstreamTimeoutMs > 0 ? config.upstreamTimeoutMs : 18e4));
+			if (typeof tokenResponse.token !== "string" || !tokenResponse.token.trim() || !Number.isFinite(tokenResponse.refresh_in) || tokenResponse.refresh_in <= 0) throw new Error("Copilot token exchange returned invalid credentials metadata.");
+			await writeStoredCopilotToken(tokenResponse);
+			config.copilotToken = tokenResponse.token;
+			config.copilotTokenGeneration = (config.copilotTokenGeneration ?? 0) + 1;
+			scheduleCopilotTokenRefresh(tokenResponse.refresh_in);
+			return tokenResponse.refresh_in;
+		})().finally(() => {
+			refreshInFlight = void 0;
+		});
+		return refreshInFlight;
 	};
 	const scheduleCopilotTokenRefresh = (refreshIn$1) => {
 		if (copilotTokenRefreshTimer) clearTimeout(copilotTokenRefreshTimer);
@@ -501,11 +523,10 @@ const setupProxyAuth = async (config, options = {}) => {
 		log.info(`Next Copilot token refresh in ${Math.round(refreshMs / 1e3)}s`);
 		copilotTokenRefreshTimer = setTimeout(async () => {
 			try {
-				const nextRefreshIn = await applyCopilotToken();
+				await applyCopilotToken();
 				log.info("Refreshed Copilot token");
-				scheduleCopilotTokenRefresh(nextRefreshIn);
-			} catch (error) {
-				log.error("Failed to refresh Copilot token:", error);
+			} catch {
+				log.error("Failed to refresh Copilot token; retrying in 60s");
 				scheduleCopilotTokenRefresh(60);
 			}
 		}, refreshMs);
@@ -517,6 +538,7 @@ const setupProxyAuth = async (config, options = {}) => {
 		config.copilotToken = storedToken.token;
 		refreshIn = getStoredTokenRemainingSeconds(storedToken);
 		log.info(`Using cached Copilot token at ${paths.copilotTokenPath}`);
+		scheduleCopilotTokenRefresh(refreshIn);
 	} else try {
 		refreshIn = await applyCopilotToken();
 	} catch (error) {
@@ -528,7 +550,14 @@ const setupProxyAuth = async (config, options = {}) => {
 		});
 		refreshIn = await applyCopilotToken();
 	}
-	scheduleCopilotTokenRefresh(refreshIn);
+	config.refreshCopilotToken = async (rejectedToken, generation) => {
+		if (config.copilotToken !== rejectedToken || (config.copilotTokenGeneration ?? 0) !== generation) return;
+		try {
+			await applyCopilotToken();
+		} catch {
+			throw new Error("Copilot token refresh failed; upstream request was not replayed.");
+		}
+	};
 	return {
 		githubLogin: await loadGitHubLogin(githubToken, config.vsCodeVersion),
 		source: "github-token"
@@ -1332,7 +1361,13 @@ const copilotDispatcher = new Agent({
 });
 const getCopilotProviderContext = (config) => ({
 	baseUrl: config.copilotBaseUrl,
-	token: config.copilotToken,
+	get token() {
+		return config.copilotToken;
+	},
+	get tokenGeneration() {
+		return config.copilotTokenGeneration ?? 0;
+	},
+	refreshToken: config.refreshCopilotToken,
 	vsCodeVersion: config.vsCodeVersion
 });
 const shouldRetryResponse = (response) => response.status >= 500 && response.status <= 599;
@@ -1398,27 +1433,68 @@ const logUpstreamLifecycle = (requestId, message) => {
 	}
 	log.debug(message);
 };
+const waitWithSignal = (pending, signal) => {
+	if (!signal) return pending;
+	if (signal.aborted) {
+		pending.catch(() => {});
+		return Promise.reject(signal.reason);
+	}
+	return new Promise((resolve$1, reject) => {
+		const abort = () => reject(signal.reason);
+		signal.addEventListener("abort", abort, { once: true });
+		pending.then(resolve$1, reject).finally(() => signal.removeEventListener("abort", abort));
+	});
+};
+const isAuthRejection = async (response, signal) => {
+	if (response.status === 401) return true;
+	if (response.status !== 403) return false;
+	const reader = response.clone().body?.getReader();
+	if (!reader) return false;
+	const decoder = new TextDecoder();
+	let text = "";
+	let size = 0;
+	try {
+		for (;;) {
+			const chunk = await waitWithSignal(reader.read(), signal);
+			if (chunk.done) return (text + decoder.decode()).trim().toLowerCase() === "forbidden";
+			size += chunk.value.byteLength;
+			if (size > 128) return false;
+			text += decoder.decode(chunk.value, { stream: true });
+		}
+	} catch (error) {
+		if (signal?.aborted) throw error;
+		return false;
+	} finally {
+		reader.cancel().catch(() => {});
+	}
+};
 const fetchCopilot = async (provider, path$1, init, options = {}) => {
 	if (!provider.token) throw new ProxyNotImplementedError("Copilot token is not configured for copilot-relay.");
 	const timeoutMs = options.timeoutMs ?? copilotRequestTimeoutMs;
 	const signal = createCopilotRequestSignal(options.signal, timeoutMs);
 	let lastError;
-	for (let attempt = 1; attempt <= maxFetchAttempts; attempt++) {
+	let transientRetries = 0;
+	let authRecoveryUsed = false;
+	for (let attempt = 1; attempt <= maxFetchAttempts + 1; attempt++) {
 		const upstreamRequestId = randomUUID();
+		const attemptedToken = provider.token;
+		const attemptedGeneration = provider.tokenGeneration ?? 0;
+		let response;
 		try {
 			const started = performance.now();
 			logUpstreamLifecycle(options.requestId, `send upstream method=${init.method ?? "GET"} path=${path$1} attempt=${attempt} upstream_request_id=${upstreamRequestId}`);
-			const response = await fetch$1(`${provider.baseUrl}${path$1}`, {
+			response = await fetch$1(`${provider.baseUrl}${path$1}`, {
 				...init,
-				headers: buildHeaders(provider, init, options, upstreamRequestId),
+				headers: buildHeaders({
+					...provider,
+					token: attemptedToken
+				}, init, options, upstreamRequestId),
 				dispatcher: copilotDispatcher,
 				signal
 			});
 			const ms = Math.round(performance.now() - started);
 			logUpstreamLifecycle(options.requestId, `return from upstream method=${init.method ?? "GET"} path=${path$1} status=${response.status} ms=${ms} attempt=${attempt} upstream_request_id=${upstreamRequestId}`);
 			log.debug(`${formatRequestId(options.requestId)}Copilot ${init.method ?? "GET"} ${path$1} -> ${response.status} ${ms}ms (attempt ${attempt}) upstream_request_id=${upstreamRequestId}`);
-			if (!shouldRetryResponse(response) || attempt === maxFetchAttempts) return response;
-			log.error(`${formatRequestId(options.requestId)}Copilot ${path$1} returned ${response.status}; retrying (${attempt}/${maxFetchAttempts}) upstream_request_id=${upstreamRequestId}`);
 		} catch (error) {
 			const abortError = toCopilotAbortHTTPError(error, signal, timeoutMs);
 			if (abortError) {
@@ -1427,12 +1503,37 @@ const fetchCopilot = async (provider, path$1, init, options = {}) => {
 			}
 			logUpstreamLifecycle(options.requestId, `upstream failed method=${init.method ?? "GET"} path=${path$1} attempt=${attempt} upstream_request_id=${upstreamRequestId}`);
 			lastError = error;
-			if (attempt === maxFetchAttempts) {
+			if (transientRetries >= maxFetchAttempts - 1) {
 				log.error(`${formatRequestId(options.requestId)}Copilot ${path$1} request failed after ${attempt} attempts upstream_request_id=${upstreamRequestId}`, error);
 				throw error;
 			}
-			log.error(`${formatRequestId(options.requestId)}Copilot ${path$1} request failed; retrying (${attempt}/${maxFetchAttempts}) upstream_request_id=${upstreamRequestId}`, error);
+			transientRetries++;
+			log.error(`${formatRequestId(options.requestId)}Copilot ${path$1} request failed; retrying (${transientRetries}/${maxFetchAttempts}) upstream_request_id=${upstreamRequestId}`, error);
+			continue;
+		}
+		try {
+			signal?.throwIfAborted();
+			if (!authRecoveryUsed && provider.refreshToken && await isAuthRejection(response, signal)) {
+				authRecoveryUsed = true;
+				log.info(`${formatRequestId(options.requestId)}Copilot ${path$1} authentication rejected status=${response.status}; refreshing token`);
+				await response.body?.cancel();
+				signal?.throwIfAborted();
+				await waitWithSignal(provider.refreshToken(attemptedToken, attemptedGeneration), signal);
+				signal?.throwIfAborted();
+				log.info(`${formatRequestId(options.requestId)}Copilot token refresh completed; retrying ${path$1}`);
+				continue;
+			}
+		} catch (error) {
+			response.body?.cancel().catch(() => {});
+			const abortError = toCopilotAbortHTTPError(error, signal, timeoutMs);
+			if (abortError) throw abortError;
+			log.error(`${formatRequestId(options.requestId)}Copilot token recovery failed; request not replayed`);
+			throw new Error("Copilot token recovery failed; request not replayed.");
 		}
+		if (!shouldRetryResponse(response) || transientRetries >= maxFetchAttempts - 1) return response;
+		transientRetries++;
+		log.error(`${formatRequestId(options.requestId)}Copilot ${path$1} returned ${response.status}; retrying (${transientRetries}/${maxFetchAttempts}) upstream_request_id=${upstreamRequestId}`);
+		await response.body?.cancel();
 	}
 	throw lastError;
 };
@@ -1981,7 +2082,8 @@ function translateTools(tools) {
 		type: "function",
 		name: tool.function.name,
 		description: tool.function.description,
-		parameters: normalizeResponsesToolSchema(tool.function.parameters)
+		parameters: normalizeResponsesToolSchema(tool.function.parameters),
+		strict: false
 	}));
 }
 function translateToolChoice(toolChoice) {
@@ -2861,6 +2963,36 @@ const createFailedSearchExecution = (payload, requestedQuery, model, message) =>
 	results: [],
 	text: message
 });
+const getSearchFailureCategory = (status$1) => {
+	if (status$1 === 503) return "service unavailable";
+	if (status$1 >= 500) return "server failure";
+	if (status$1 === 429) return "rate limited";
+	if (status$1 === 401) return "authentication rejected";
+	if (status$1 === 403) return "access denied";
+	return "request failed";
+};
+const getSearchFailureDetail = (body, token) => {
+	let detail = body.trim();
+	try {
+		const parsed = JSON.parse(detail);
+		const error = isRecord(parsed) && isRecord(parsed.error) ? parsed.error : void 0;
+		detail = typeof error?.message === "string" ? error.message : typeof error?.code === "string" ? error.code : "";
+	} catch {
+		if (/^[{[<]/.test(detail)) return "";
+	}
+	detail = scrubSensitiveUrls(sanitizeTerminalString(detail));
+	detail = detail.replace(/https?:\/\/[^\s'"`<>]+/gi, (raw) => {
+		try {
+			const url = new URL(raw);
+			return url.username || url.password || url.search || url.hash ? `${url.origin}/[redacted]` : raw;
+		} catch {
+			return "[redacted]";
+		}
+	});
+	if (token) detail = detail.replaceAll(token, "[redacted]");
+	if (/\b(?:authorization|bearer|api[_ -]?key|password|secret|access[_ -]?token|refresh[_ -]?token)\b/i.test(detail) || /\btoken\s*[=:]/i.test(detail) || /\b(?:gh[pousr]_|github_pat_|sk-|eyJ)[A-Za-z0-9_-]+/.test(detail) || /\b(?:request|payload|prompt|messages|input|headers|conversation)\b["']?\s*[:=]/i.test(detail)) return "";
+	return detail.trim().slice(0, 240);
+};
 const createClaudeWebSearchExecution = async (config, payload, requestedQuery, options = {}) => {
 	const backendModel = getWebSearchBackendModel(config);
 	const signal = createCopilotRequestSignal(options.signal, options.timeoutMs);
@@ -2878,8 +3010,8 @@ const createClaudeWebSearchExecution = async (config, payload, requestedQuery, o
 		timeoutMs: options.timeoutMs
 	});
 	if (!response.ok) {
-		const detail = await readCopilotText(response, signal, options.timeoutMs).catch(() => "");
-		return createFailedSearchExecution(payload, requestedQuery, backendModel, [`Copilot web search is not available for model ${backendModel}.`, detail ? `Upstream response: ${detail}` : ""].filter(Boolean).join("\n"));
+		const detail = getSearchFailureDetail(await readCopilotText(response, signal, options.timeoutMs).catch(() => ""), config.copilotToken);
+		return createFailedSearchExecution(payload, requestedQuery, backendModel, [`Copilot web search upstream ${getSearchFailureCategory(response.status)} (HTTP ${response.status}; model ${backendModel}).`, detail ? `Upstream response: ${detail}` : ""].filter(Boolean).join("\n"));
 	}
 	const upstream = await readCopilotJson(response, signal, options.timeoutMs);
 	const text = getResponseText(upstream);