@oh-my-pi/pi-ai 16.3.2 → 16.3.4

package/src/providers/openai-completions.ts

@@ -148,6 +148,13 @@ type OpenAICompletionsCompletionTokenDetails = {
 	reasoning_tokens?: unknown;
 };
 
+function firstPositiveNumber(...values: unknown[]): number {
+	for (const value of values) {
+		if (typeof value === "number" && value > 0) return value;
+	}
+	return 0;
+}
+
 /**
  * Normalize tool call ID for Mistral.
  * Mistral requires tool IDs to be exactly 9 alphanumeric characters (a-z, A-Z, 0-9).
@@ -1558,11 +1565,7 @@ export function parseChunkUsage(
 	const accounting = calculateOpenAIUsageAccounting({
 		promptTokens: typeof promptTokens === "number" ? promptTokens : 0,
 		outputTokens,
-		cachedTokens:
-			(typeof cachedTokens === "number" ? cachedTokens : undefined) ??
-			(typeof promptCacheHitTokens === "number" ? promptCacheHitTokens : undefined) ??
-			(typeof promptTokenCachedTokens === "number" ? promptTokenCachedTokens : undefined) ??
-			0,
+		cachedTokens: firstPositiveNumber(cachedTokens, promptCacheHitTokens, promptTokenCachedTokens),
 		reasoningTokens: typeof completionReasoningTokens === "number" ? completionReasoningTokens : 0,
 		cacheWriteOpenRouter: typeof cacheWriteTokens === "number" ? cacheWriteTokens : undefined,
 		cacheWriteDeepSeek: typeof promptCacheMissTokens === "number" ? promptCacheMissTokens : undefined,
@@ -1786,12 +1789,12 @@ export function convertMessages(
 				if (compat.requiresThinkingAsText) {
 					const thinkingText = nonEmptyThinkingBlocks
 						.map(b => renderDemotedThinking(model.id, b.thinking))
-						.join("");
+						.join(" ");
 					// `content` is a plain string at this point (set above) or null —
 					// never an array. Prepend the demoted thinking to the string form.
 					assistantMsg.content =
 						typeof assistantMsg.content === "string" && assistantMsg.content.length > 0
-							? `${thinkingText}${assistantMsg.content}`
+							? `${thinkingText} ${assistantMsg.content}`
 							: thinkingText;
 				} else if (compat.requiresReasoningContentForToolCalls) {
 					// Use the streamed signature when the backend accepts whichever

package/src/providers/transform-messages.ts

@@ -341,13 +341,18 @@ export function transformMessages<TApi extends Api>(
 			const isLatestSurvivingAssistant = index === latestSurvivingAssistantIndex;
 			// Signature policy is a second axis. Anthropic cryptographically
 			// binds reasoning signatures to its key+session+model, so cross-model
-			// signatures must be stripped whenever official Anthropic is on
-			// either end of the replay:
-			//   * official → 3p: the 3p target can't reverify the signature;
-			//     keeping it leaks private continuation metadata for no benefit.
-			//   * 3p → official: official rejects a foreign signature outright.
-			//   * official → official cross-model: the new model rejects the
-			//     previous model's signature.
+			// signatures must be stripped whenever a signing Anthropic endpoint
+			// is on either end of the replay:
+			//   * official Anthropic (source): the 3p target can't reverify a
+			//     foreign signature and keeping it leaks continuation metadata
+			//     for no benefit.
+			//   * signing Anthropic (target): official Anthropic, GitHub Copilot,
+			//     ZenMux, Cloudflare AI Gateway `/anthropic`, and Google Vertex
+			//     `publishers/anthropic/…` all forward to signature-enforcing
+			//     Anthropic. Any stale/cross-model signature on the wire triggers
+			//     `400 Invalid signature in thinking block` — same failure class
+			//     whether `officialEndpoint` is true or the endpoint is one of
+			//     the known signing proxies (#4297).
 			// 3p ↔ 3p replays preserve signatures because compatible providers
 			// (Z.AI, DeepSeek, custom `models.yaml` providers) treat them as
 			// opaque continuation hints rather than verified material; stripping
@@ -358,8 +363,8 @@ export function transformMessages<TApi extends Api>(
 			// a custom proxy via `models.yaml` will see signatures stripped, the
 			// conservative direction (degraded reasoning, not broken requests).
 			const isOfficialAnthropicSource = isAnthropicReplay && assistantMsg.provider === "anthropic";
-			const isOfficialAnthropicTarget = isAnthropicTarget && model.compat.officialEndpoint;
-			const officialAnthropicInvolved = isOfficialAnthropicSource || isOfficialAnthropicTarget;
+			const isSigningAnthropicTarget = isAnthropicTarget && model.compat.signingEndpoint;
+			const signingAnthropicInvolved = isOfficialAnthropicSource || isSigningAnthropicTarget;
 			// Compatible Anthropic-messages reasoning targets that accept
 			// unsigned thinking natively (Z.AI, DeepSeek, the generic
 			// `reasoning && !official` case in the compat builder). Used to keep
@@ -421,7 +426,7 @@ export function transformMessages<TApi extends Api>(
 						if (
 							!isLatestSurvivingAssistant &&
 							!isSameModel &&
-							officialAnthropicInvolved &&
+							signingAnthropicInvolved &&
 							sanitized.thinkingSignature
 						) {
 							sanitized = { ...sanitized, thinkingSignature: undefined };
@@ -438,7 +443,7 @@ export function transformMessages<TApi extends Api>(
 						// textual thinking dialect; keep demotion for signatures stripped
 						// by the untrustworthy-turn recovery above and for literal thinking
 						// envelopes that never carried a signature field.
-						if (isSameModel && isOfficialAnthropicTarget && sanitized.thinkingSignature?.trim() === "") {
+						if (isSameModel && isSigningAnthropicTarget && sanitized.thinkingSignature?.trim() === "") {
 							return [];
 						}
 						return sanitized;

package/src/registry/baseten.ts

@@ -0,0 +1,22 @@
+import { createApiKeyLogin } from "./api-key-login";
+import type { OAuthLoginCallbacks } from "./oauth/types";
+import type { ProviderDefinition } from "./types";
+
+export const loginBaseten = createApiKeyLogin({
+	providerLabel: "Baseten",
+	authUrl: "https://app.baseten.co/settings/api_keys",
+	instructions: "Copy your API key from the Baseten dashboard",
+	promptMessage: "Paste your Baseten API key",
+	placeholder: "bt_...",
+	validation: {
+		kind: "models-endpoint",
+		provider: "Baseten",
+		modelsUrl: "https://inference.baseten.co/v1/models",
+	},
+});
+
+export const basetenProvider = {
+	id: "baseten",
+	name: "Baseten",
+	login: (cb: OAuthLoginCallbacks) => loginBaseten(cb),
+} as const satisfies ProviderDefinition;

package/src/registry/registry.ts

@@ -4,6 +4,7 @@ import { alibabaCodingPlanProvider } from "./alibaba-coding-plan";
 import { amazonBedrockProvider } from "./amazon-bedrock";
 import { anthropicProvider } from "./anthropic";
 import { azureProvider } from "./azure";
+import { basetenProvider } from "./baseten";
 import { cerebrasProvider } from "./cerebras";
 import { cloudflareAiGatewayProvider } from "./cloudflare-ai-gateway";
 import { coreWeaveProvider } from "./coreweave";
@@ -105,6 +106,7 @@ const ALL = [
 	deepseekProvider,
 	moonshotProvider,
 	cerebrasProvider,
+	basetenProvider,
 	fireworksProvider,
 	togetherProvider,
 	nvidiaProvider,

package/src/usage/claude.ts

@@ -1,17 +1,20 @@
 import { scheduler } from "node:timers/promises";
+import { bareModelId, parseAnthropicModel } from "@oh-my-pi/pi-catalog/identity";
 import { toNumber } from "@oh-my-pi/pi-catalog/utils";
 import * as AIError from "../error";
 import { claudeCodeVersion } from "../providers/anthropic";
-import type {
-	CredentialRankingStrategy,
-	UsageAmount,
-	UsageFetchContext,
-	UsageFetchParams,
-	UsageLimit,
-	UsageProvider,
-	UsageReport,
-	UsageStatus,
-	UsageWindow,
+import {
+	type CredentialRankingContext,
+	type CredentialRankingStrategy,
+	resolveUsedFraction,
+	type UsageAmount,
+	type UsageFetchContext,
+	type UsageFetchParams,
+	type UsageLimit,
+	type UsageProvider,
+	type UsageReport,
+	type UsageStatus,
+	type UsageWindow,
 } from "../usage";
 import { isRecord } from "../utils";
 
@@ -60,13 +63,37 @@ interface ParsedUsageBucket {
 	utilization?: number;
 	resetsAt?: number;
 }
-type ClaudeUnifiedWindow = "5h" | "7d";
+type ClaudeUnifiedWindow = "5h" | "7d" | "7d_oi";
+type ClaudeModelKind = "opus" | "sonnet" | "fable" | "mythos";
 
 interface ClaudeUsageResponse {
 	five_hour?: ClaudeUsageBucket | null;
 	seven_day?: ClaudeUsageBucket | null;
 	seven_day_opus?: ClaudeUsageBucket | null;
 	seven_day_sonnet?: ClaudeUsageBucket | null;
+	limits?: unknown;
+}
+
+interface ClaudeApiLimitModelScope {
+	display_name?: string | null;
+}
+
+interface ClaudeApiLimitScope {
+	model?: ClaudeApiLimitModelScope | null;
+}
+
+interface ClaudeApiLimitEntry {
+	kind?: string;
+	percent?: unknown;
+	resets_at?: string | null;
+	scope?: ClaudeApiLimitScope | null;
+	is_active?: boolean;
+}
+
+interface ParsedApiLimitEntry {
+	kind: string;
+	bucket: ParsedUsageBucket;
+	displayName?: string;
 }
 
 type ClaudeUsagePayload = {
@@ -89,6 +116,43 @@ function parseBucket(bucket: unknown): ParsedUsageBucket | undefined {
 	}
 	return { utilization, resetsAt };
 }
+
+function getApiLimitDisplayName(scope: unknown): string | undefined {
+	if (!isRecord(scope)) return undefined;
+	const model = scope.model;
+	if (!isRecord(model)) return undefined;
+	const displayName = model.display_name;
+	return typeof displayName === "string" && displayName.trim() ? displayName.trim() : undefined;
+}
+
+/**
+ * Anthropic kept the legacy account-wide buckets populated, but as of
+ * 2026-07-02 the legacy per-model weekly buckets (`seven_day_opus` /
+ * `seven_day_sonnet`) are permanently null. Model-scoped weekly caps now arrive
+ * only through generic `limits[]` entries (`kind: "weekly_scoped"`) with the
+ * model family named by `scope.model.display_name`.
+ */
+function parseApiLimitEntries(raw: unknown): ParsedApiLimitEntry[] {
+	if (!Array.isArray(raw)) return [];
+	const entries: ParsedApiLimitEntry[] = [];
+	for (const rawEntry of raw) {
+		if (!isRecord(rawEntry)) continue;
+		const entry = rawEntry as ClaudeApiLimitEntry;
+		if (typeof entry.kind !== "string") continue;
+		if (entry.is_active === false) continue;
+		const utilization = toNumber(entry.percent);
+		const resetsAt = parseIsoTime(typeof entry.resets_at === "string" ? entry.resets_at : undefined);
+		if (utilization === undefined && resetsAt === undefined) continue;
+		const displayName = getApiLimitDisplayName(entry.scope);
+		entries.push({
+			kind: entry.kind,
+			bucket: { utilization, resetsAt },
+			...(displayName ? { displayName } : {}),
+		});
+	}
+	return entries;
+}
+
 function parseUnifiedWindow(
 	headers: Record<string, string>,
 	window: ClaudeUnifiedWindow,
@@ -144,12 +208,17 @@ function hasUsageData(payload: ClaudeUsageResponse): boolean {
 		parseBucket(payload.five_hour)?.utilization !== undefined ||
 		parseBucket(payload.seven_day)?.utilization !== undefined ||
 		parseBucket(payload.seven_day_opus)?.utilization !== undefined ||
-		parseBucket(payload.seven_day_sonnet)?.utilization !== undefined
+		parseBucket(payload.seven_day_sonnet)?.utilization !== undefined ||
+		parseApiLimitEntries(payload.limits).some(entry => entry.bucket.utilization !== undefined)
 	);
 }
 
 function isRetryableStatus(status: number): boolean {
-	return AIError.isTransientStatus(status);
+	// Exclude 429: the usage endpoint is informational and rate-limited per
+	// source IP, so retrying a rate_limit_error inside a single fetch can't
+	// succeed and only deepens the throttle (3 attempts per poll). Fall through
+	// to the caller's failure cool-down and retry on the next poll instead.
+	return AIError.isTransientStatus(status) && status !== 429;
 }
 
 function isAbortError(error: unknown, signal?: AbortSignal): boolean {
@@ -334,8 +403,8 @@ function buildUsageLimit(args: {
 		scope: {
 			provider: args.provider,
 			windowId: args.windowId,
-			tier: args.tier,
-			shared: args.shared,
+			...(args.tier !== undefined ? { tier: args.tier } : {}),
+			...(args.shared !== undefined ? { shared: args.shared } : {}),
 		},
 		window,
 		amount,
@@ -343,9 +412,47 @@ function buildUsageLimit(args: {
 	};
 }
 
+function slugifyClaudeLimitDisplayName(displayName: string): string {
+	return displayName
+		.trim()
+		.toLowerCase()
+		.replace(/[^a-z0-9]+/g, "-")
+		.replace(/^-+|-+$/g, "");
+}
+
+/**
+ * Scoped weekly rows are per-model-family counters, not account-wide windows.
+ * They deliberately leave `scope.shared` unset so credential-wide exhaustion
+ * gating only considers the shared umbrella windows; an exhausted Fable weekly
+ * cap must not block Opus or Sonnet requests on the same credential.
+ */
+function buildScopedWeeklyUsageLimits(entries: readonly ParsedApiLimitEntry[]): UsageLimit[] {
+	const seenSlugs = new Set<string>();
+	const limits: UsageLimit[] = [];
+	for (const entry of entries) {
+		if (entry.kind !== "weekly_scoped" || !entry.displayName) continue;
+		const slug = slugifyClaudeLimitDisplayName(entry.displayName);
+		if (!slug || seenSlugs.has(slug)) continue;
+		seenSlugs.add(slug);
+		const limit = buildUsageLimit({
+			id: `anthropic:7d:${slug}`,
+			label: `Claude 7 Day (${entry.displayName})`,
+			windowId: "7d",
+			windowLabel: "7 Day",
+			durationMs: SEVEN_DAYS_MS,
+			bucket: entry.bucket,
+			provider: "anthropic",
+			tier: slug,
+		});
+		if (limit) limits.push(limit);
+	}
+	return limits;
+}
+
 export function parseClaudeRateLimitHeaders(headers: Record<string, string>, now = Date.now()): UsageReport | null {
 	const fiveHour = parseUnifiedWindow(headers, "5h");
 	const sevenDay = parseUnifiedWindow(headers, "7d");
+	const modelScopedSevenDay = parseUnifiedWindow(headers, "7d_oi");
 	const limits = [
 		buildUsageLimit({
 			id: "anthropic:5h",
@@ -367,6 +474,16 @@ export function parseClaudeRateLimitHeaders(headers: Record<string, string>, now
 			provider: "anthropic",
 			shared: true,
 		}),
+		buildUsageLimit({
+			id: "anthropic:7d:fable",
+			label: "Claude 7 Day (Fable)",
+			windowId: "7d",
+			windowLabel: "7 Day",
+			durationMs: SEVEN_DAYS_MS,
+			bucket: modelScopedSevenDay,
+			provider: "anthropic",
+			tier: "fable",
+		}),
 	].filter((limit): limit is UsageLimit => limit !== null);
 
 	if (limits.length === 0) return null;
@@ -394,8 +511,10 @@ async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext
 	if (!payloadResult || !isRecord(payloadResult.payload)) return null;
 	const { payload, orgId } = payloadResult;
 
-	const fiveHour = parseBucket(payload.five_hour);
-	const sevenDay = parseBucket(payload.seven_day);
+	const apiLimitEntries = parseApiLimitEntries(payload.limits);
+	const fiveHour = parseBucket(payload.five_hour) ?? apiLimitEntries.find(entry => entry.kind === "session")?.bucket;
+	const sevenDay =
+		parseBucket(payload.seven_day) ?? apiLimitEntries.find(entry => entry.kind === "weekly_all")?.bucket;
 	const sevenDayOpus = parseBucket(payload.seven_day_opus);
 	const sevenDaySonnet = parseBucket(payload.seven_day_sonnet);
 
@@ -440,6 +559,7 @@ async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext
 			provider: "anthropic",
 			tier: "sonnet",
 		}),
+		...buildScopedWeeklyUsageLimits(apiLimitEntries),
 	].filter((limit): limit is UsageLimit => limit !== null);
 
 	if (limits.length === 0) return null;
@@ -475,11 +595,101 @@ export const claudeUsageProvider: UsageProvider = {
 	supports: params => params.provider === "anthropic" && params.credential.type === "oauth",
 };
 
+function getClaudeModelKind(context: CredentialRankingContext | undefined): ClaudeModelKind | undefined {
+	const modelId = context?.modelId;
+	if (!modelId) return undefined;
+	return parseAnthropicModel(bareModelId(modelId))?.kind;
+}
+
+/**
+ * Claude model-scoped rows are only relevant to the matching model family.
+ * Credential-wide exhaustion checks stay on shared umbrella windows unless the
+ * request model parses to a concrete Anthropic kind, preventing a Fable cap from
+ * suppressing unrelated Opus/Sonnet traffic.
+ */
+function scopeClaudeLimitsForModel(report: UsageReport, context: CredentialRankingContext | undefined): UsageLimit[] {
+	const kind = getClaudeModelKind(context);
+	return report.limits.filter(
+		limit => limit.scope.shared === true || (kind !== undefined && limit.scope.tier === kind),
+	);
+}
+
+/**
+ * Exclude Fable and Mythos tier weekly caps from proactive hard-blocking
+ * (gating) because they are notoriously unreliable (they report 100% exhausted
+ * while the account can still serve requests). They stay available for ranking
+ * pressure in findWindowLimits via scopeClaudeLimitsForModel.
+ */
+function scopeClaudeLimitsForModelHardBlock(
+	report: UsageReport,
+	context: CredentialRankingContext | undefined,
+): UsageLimit[] {
+	const kind = getClaudeModelKind(context);
+	const excludeHardBlock = kind === "fable" || kind === "mythos";
+	return report.limits.filter(
+		limit => limit.scope.shared === true || (kind !== undefined && limit.scope.tier === kind && !excludeHardBlock),
+	);
+}
+
+function rankingUsedFraction(limit: UsageLimit): number {
+	const fraction = resolveUsedFraction(limit);
+	if (typeof fraction !== "number" || !Number.isFinite(fraction)) return 0.5;
+	return Math.min(Math.max(fraction, 0), 1);
+}
+
+function rankingDrainRate(limit: UsageLimit, nowMs: number): number {
+	const usedFraction = rankingUsedFraction(limit);
+	const durationMs = limit.window?.durationMs ?? SEVEN_DAYS_MS;
+	if (!Number.isFinite(durationMs) || durationMs <= 0) return usedFraction;
+	const resetAt = limit.window?.resetsAt;
+	if (typeof resetAt !== "number" || !Number.isFinite(resetAt)) return usedFraction;
+	const remainingWindowMs = resetAt - nowMs;
+	const clampedRemainingWindowMs = Math.min(Math.max(remainingWindowMs, 0), durationMs);
+	const elapsedMs = durationMs - clampedRemainingWindowMs;
+	if (elapsedMs <= 0) return usedFraction;
+	const elapsedHours = elapsedMs / (60 * 60 * 1000);
+	if (!Number.isFinite(elapsedHours) || elapsedHours <= 0) return usedFraction;
+	return usedFraction / elapsedHours;
+}
+
+function morePressuredLimit(
+	left: UsageLimit | undefined,
+	right: UsageLimit | undefined,
+	nowMs: number,
+): UsageLimit | undefined {
+	if (!left) return right;
+	if (!right) return left;
+	const leftDrainRate = rankingDrainRate(left, nowMs);
+	const rightDrainRate = rankingDrainRate(right, nowMs);
+	if (rightDrainRate !== leftDrainRate) return rightDrainRate > leftDrainRate ? right : left;
+	return rankingUsedFraction(right) > rankingUsedFraction(left) ? right : left;
+}
+
+function findClaudeSecondaryLimit(
+	report: UsageReport,
+	context: CredentialRankingContext | undefined,
+): UsageLimit | undefined {
+	const nowMs = Date.now();
+	return scopeClaudeLimitsForModel(report, context)
+		.filter(limit => limit.scope.windowId === "7d" || limit.window?.id === "7d")
+		.reduce<UsageLimit | undefined>((selected, limit) => morePressuredLimit(selected, limit, nowMs), undefined);
+}
+
 export const claudeRankingStrategy: CredentialRankingStrategy = {
-	findWindowLimits(report) {
-		const primary = report.limits.find(l => l.id === "anthropic:5h");
-		const secondary = report.limits.find(l => l.id === "anthropic:7d");
+	findWindowLimits(report, context) {
+		const primary = report.limits.find(limit => limit.id === "anthropic:5h");
+		const secondary = findClaudeSecondaryLimit(report, context);
 		return { primary, secondary };
 	},
+	scopeLimits: scopeClaudeLimitsForModelHardBlock,
+	/**
+	 * Fable/Mythos usage-limit errors map to tier-local weekly counters. Scope
+	 * reactive backoff blocks for those tiers, mirroring the per-counter
+	 * precedent in packages/ai/src/usage/google-antigravity.ts:466-497.
+	 */
+	blockScope(context) {
+		const kind = getClaudeModelKind(context);
+		return kind === "fable" || kind === "mythos" ? `tier:${kind}` : undefined;
+	},
 	windowDefaults: { primaryMs: 5 * 60 * 60 * 1000, secondaryMs: 7 * 24 * 60 * 60 * 1000 },
 };

package/src/utils/block-symbols.ts

@@ -30,3 +30,19 @@ export const kStreamingArgumentsDone = Symbol("provider.block.argumentsDone");
 
 /** Classifies Cursor's in-flight tool-call kind without leaking provider-private state. */
 export const kStreamingBlockKind = Symbol("provider.block.kind");
+
+/**
+ * Marks a `toolCall` content block that Cursor's exec channel already
+ * executed server-side (via the coding-agent bridge) and whose result is
+ * buffered separately for emission via the assistant-loop stream.
+ *
+ * `agent-loop.ts` MUST skip execution of blocks carrying this marker —
+ * treating them as a fresh runnable tool call would run the same
+ * side-effecting tool (bash, write, delete, …) a second time. Symbol-keyed
+ * so it never persists across the JSONL round-trip, where rebuild instead
+ * pairs the block with its already-persisted `toolResult` message by id.
+ */
+export const kCursorExecResolved = Symbol("provider.block.cursorExecResolved");
+
+/** Carries the resolved marker without exposing a string-keyed property. */
+export type CursorExecResolvedCarrier = object & { [kCursorExecResolved]?: true };

package/src/utils/proxy.ts

@@ -162,11 +162,26 @@ export function wrapFetchForProxy(fetchImpl: FetchImpl, provider: string): Fetch
 	return wrapped;
 }
 
+export interface ConnectProxiedSocketOptions {
+	/** Caller cancellation for the proxy TCP/TLS handshake and CONNECT tunnel. */
+	signal?: AbortSignal;
+	/** Maximum wall-clock time to establish the final TLS tunnel. Disabled when absent or non-positive. */
+	timeoutMs?: number;
+}
+
 /**
  * Tunnel a socket connection through an HTTP CONNECT proxy.
  * This is used specifically to wrap Node's `http2.connect(baseUrl, { createConnection })` for Cursor.
  */
-export async function connectProxiedSocket(proxyUrlStr: string, targetUrlStr: string): Promise<tls.TLSSocket> {
+export async function connectProxiedSocket(
+	proxyUrlStr: string,
+	targetUrlStr: string,
+	options?: ConnectProxiedSocketOptions,
+): Promise<tls.TLSSocket> {
+	if (options?.signal?.aborted) {
+		throw new AIError.AbortError("Proxy tunnel aborted");
+	}
+
 	const proxyUrl = new URL(proxyUrlStr);
 	const targetUrl = new URL(targetUrlStr);
 
@@ -179,23 +194,73 @@ export async function connectProxiedSocket(proxyUrlStr: string, targetUrlStr: st
 
 	const { promise, resolve, reject } = Promise.withResolvers<tls.TLSSocket>();
 
-	let rawSocket: net.Socket;
-	if (useProxySsl) {
-		rawSocket = tls.connect({
-			host: proxyHost,
-			port: proxyPort,
-		});
-	} else {
-		rawSocket = net.connect({
-			host: proxyHost,
-			port: proxyPort,
-		});
-	}
-
-	rawSocket.once("error", reject);
-
 	const readyEvent = useProxySsl ? "secureConnect" : "connect";
-	rawSocket.once(readyEvent, () => {
+	let rawSocket: net.Socket | undefined;
+	let tunnelSocket: tls.TLSSocket | undefined;
+	let timeout: NodeJS.Timeout | undefined;
+	let responseData = "";
+	let settled = false;
+
+	const cleanup = (): void => {
+		if (timeout) {
+			clearTimeout(timeout);
+			timeout = undefined;
+		}
+		options?.signal?.removeEventListener("abort", onAbort);
+		rawSocket?.off("error", onRawError);
+		rawSocket?.off(readyEvent, onProxyReady);
+		rawSocket?.off("data", onProxyData);
+		tunnelSocket?.off("secureConnect", onTunnelReady);
+		tunnelSocket?.off("error", onTunnelError);
+	};
+	const destroyInProgress = (): void => {
+		tunnelSocket?.destroy();
+		rawSocket?.destroy();
+	};
+	const rejectOnce = (error: Error): void => {
+		if (settled) return;
+		settled = true;
+		cleanup();
+		destroyInProgress();
+		reject(error);
+	};
+	const resolveOnce = (socket: tls.TLSSocket): void => {
+		if (settled) return;
+		settled = true;
+		cleanup();
+		resolve(socket);
+	};
+	const onAbort = (): void => rejectOnce(new AIError.AbortError("Proxy tunnel aborted"));
+	const onRawError = (error: Error): void => rejectOnce(error);
+	const onTunnelError = (error: Error): void => rejectOnce(error);
+	const onTunnelReady = (): void => {
+		if (!tunnelSocket) return;
+		resolveOnce(tunnelSocket);
+	};
+	const onProxyData = (chunk: Buffer): void => {
+		if (!rawSocket) return;
+		responseData += chunk.toString("binary");
+		if (!responseData.includes("\r\n\r\n")) return;
+
+		rawSocket.off("data", onProxyData);
+		rawSocket.off("error", onRawError);
+
+		const firstLine = responseData.split("\r\n")[0];
+		if (!firstLine.includes(" 200 ")) {
+			rejectOnce(new AIError.ValidationError(`Proxy tunnel failed: ${firstLine}`));
+			return;
+		}
+
+		tunnelSocket = tls.connect({
+			socket: rawSocket,
+			servername: targetHost,
+			ALPNProtocols: ["h2"],
+		});
+		tunnelSocket.once("secureConnect", onTunnelReady);
+		tunnelSocket.once("error", onTunnelError);
+	};
+	const onProxyReady = (): void => {
+		if (!rawSocket) return;
 		let connectReq = `CONNECT ${targetHost}:${targetPort} HTTP/1.1\r\n` + `Host: ${targetHost}:${targetPort}\r\n`;
 
 		if (proxyUrl.username || proxyUrl.password) {
@@ -207,34 +272,29 @@ export async function connectProxiedSocket(proxyUrlStr: string, targetUrlStr: st
 		connectReq += "\r\n";
 
 		rawSocket.write(connectReq);
+		rawSocket.on("data", onProxyData);
+	};
 
-		let responseData = "";
-		const onData = (chunk: Buffer) => {
-			responseData += chunk.toString("binary");
-			if (responseData.includes("\r\n\r\n")) {
-				rawSocket.off("data", onData);
-				rawSocket.off("error", reject);
-
-				const firstLine = responseData.split("\r\n")[0];
-				if (firstLine.includes(" 200 ")) {
-					const tlsSocket = tls.connect({
-						socket: rawSocket,
-						servername: targetHost,
-						ALPNProtocols: ["h2"],
-					});
-
-					tlsSocket.once("secureConnect", () => {
-						resolve(tlsSocket);
-					});
-					tlsSocket.once("error", reject);
-				} else {
-					rawSocket.destroy();
-					reject(new AIError.ValidationError(`Proxy tunnel failed: ${firstLine}`));
-				}
-			}
-		};
-		rawSocket.on("data", onData);
-	});
+	options?.signal?.addEventListener("abort", onAbort, { once: true });
+	if (options?.timeoutMs !== undefined && Number.isFinite(options.timeoutMs) && options.timeoutMs > 0) {
+		const timeoutMs = Math.trunc(options.timeoutMs);
+		timeout = setTimeout(() => {
+			rejectOnce(new AIError.StreamTimeoutError(`Proxy tunnel timed out after ${timeoutMs}ms`));
+		}, timeoutMs);
+		timeout.unref?.();
+	}
+
+	rawSocket = useProxySsl
+		? tls.connect({
+				host: proxyHost,
+				port: proxyPort,
+			})
+		: net.connect({
+				host: proxyHost,
+				port: proxyPort,
+			});
+	rawSocket.once("error", onRawError);
+	rawSocket.once(readyEvent, onProxyReady);
 
 	return promise;
 }