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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -148,6 +148,13 @@ type OpenAICompletionsCompletionTokenDetails = {
148
148
  reasoning_tokens?: unknown;
149
149
  };
150
150
 
151
+ function firstPositiveNumber(...values: unknown[]): number {
152
+ for (const value of values) {
153
+ if (typeof value === "number" && value > 0) return value;
154
+ }
155
+ return 0;
156
+ }
157
+
151
158
  /**
152
159
  * Normalize tool call ID for Mistral.
153
160
  * Mistral requires tool IDs to be exactly 9 alphanumeric characters (a-z, A-Z, 0-9).
@@ -1558,11 +1565,7 @@ export function parseChunkUsage(
1558
1565
  const accounting = calculateOpenAIUsageAccounting({
1559
1566
  promptTokens: typeof promptTokens === "number" ? promptTokens : 0,
1560
1567
  outputTokens,
1561
- cachedTokens:
1562
- (typeof cachedTokens === "number" ? cachedTokens : undefined) ??
1563
- (typeof promptCacheHitTokens === "number" ? promptCacheHitTokens : undefined) ??
1564
- (typeof promptTokenCachedTokens === "number" ? promptTokenCachedTokens : undefined) ??
1565
- 0,
1568
+ cachedTokens: firstPositiveNumber(cachedTokens, promptCacheHitTokens, promptTokenCachedTokens),
1566
1569
  reasoningTokens: typeof completionReasoningTokens === "number" ? completionReasoningTokens : 0,
1567
1570
  cacheWriteOpenRouter: typeof cacheWriteTokens === "number" ? cacheWriteTokens : undefined,
1568
1571
  cacheWriteDeepSeek: typeof promptCacheMissTokens === "number" ? promptCacheMissTokens : undefined,
@@ -1786,12 +1789,12 @@ export function convertMessages(
1786
1789
  if (compat.requiresThinkingAsText) {
1787
1790
  const thinkingText = nonEmptyThinkingBlocks
1788
1791
  .map(b => renderDemotedThinking(model.id, b.thinking))
1789
- .join("");
1792
+ .join(" ");
1790
1793
  // `content` is a plain string at this point (set above) or null —
1791
1794
  // never an array. Prepend the demoted thinking to the string form.
1792
1795
  assistantMsg.content =
1793
1796
  typeof assistantMsg.content === "string" && assistantMsg.content.length > 0
1794
- ? `${thinkingText}${assistantMsg.content}`
1797
+ ? `${thinkingText} ${assistantMsg.content}`
1795
1798
  : thinkingText;
1796
1799
  } else if (compat.requiresReasoningContentForToolCalls) {
1797
1800
  // Use the streamed signature when the backend accepts whichever
@@ -341,13 +341,18 @@ export function transformMessages<TApi extends Api>(
341
341
  const isLatestSurvivingAssistant = index === latestSurvivingAssistantIndex;
342
342
  // Signature policy is a second axis. Anthropic cryptographically
343
343
  // binds reasoning signatures to its key+session+model, so cross-model
344
- // signatures must be stripped whenever official Anthropic is on
345
- // either end of the replay:
346
- // * official 3p: the 3p target can't reverify the signature;
347
- // keeping it leaks private continuation metadata for no benefit.
348
- // * 3p → official: official rejects a foreign signature outright.
349
- // * official official cross-model: the new model rejects the
350
- // previous model's signature.
344
+ // signatures must be stripped whenever a signing Anthropic endpoint
345
+ // is on either end of the replay:
346
+ // * official Anthropic (source): the 3p target can't reverify a
347
+ // foreign signature and keeping it leaks continuation metadata
348
+ // for no benefit.
349
+ // * signing Anthropic (target): official Anthropic, GitHub Copilot,
350
+ // ZenMux, Cloudflare AI Gateway `/anthropic`, and Google Vertex
351
+ // `publishers/anthropic/…` all forward to signature-enforcing
352
+ // Anthropic. Any stale/cross-model signature on the wire triggers
353
+ // `400 Invalid signature in thinking block` — same failure class
354
+ // whether `officialEndpoint` is true or the endpoint is one of
355
+ // the known signing proxies (#4297).
351
356
  // 3p ↔ 3p replays preserve signatures because compatible providers
352
357
  // (Z.AI, DeepSeek, custom `models.yaml` providers) treat them as
353
358
  // opaque continuation hints rather than verified material; stripping
@@ -358,8 +363,8 @@ export function transformMessages<TApi extends Api>(
358
363
  // a custom proxy via `models.yaml` will see signatures stripped, the
359
364
  // conservative direction (degraded reasoning, not broken requests).
360
365
  const isOfficialAnthropicSource = isAnthropicReplay && assistantMsg.provider === "anthropic";
361
- const isOfficialAnthropicTarget = isAnthropicTarget && model.compat.officialEndpoint;
362
- const officialAnthropicInvolved = isOfficialAnthropicSource || isOfficialAnthropicTarget;
366
+ const isSigningAnthropicTarget = isAnthropicTarget && model.compat.signingEndpoint;
367
+ const signingAnthropicInvolved = isOfficialAnthropicSource || isSigningAnthropicTarget;
363
368
  // Compatible Anthropic-messages reasoning targets that accept
364
369
  // unsigned thinking natively (Z.AI, DeepSeek, the generic
365
370
  // `reasoning && !official` case in the compat builder). Used to keep
@@ -421,7 +426,7 @@ export function transformMessages<TApi extends Api>(
421
426
  if (
422
427
  !isLatestSurvivingAssistant &&
423
428
  !isSameModel &&
424
- officialAnthropicInvolved &&
429
+ signingAnthropicInvolved &&
425
430
  sanitized.thinkingSignature
426
431
  ) {
427
432
  sanitized = { ...sanitized, thinkingSignature: undefined };
@@ -438,7 +443,7 @@ export function transformMessages<TApi extends Api>(
438
443
  // textual thinking dialect; keep demotion for signatures stripped
439
444
  // by the untrustworthy-turn recovery above and for literal thinking
440
445
  // envelopes that never carried a signature field.
441
- if (isSameModel && isOfficialAnthropicTarget && sanitized.thinkingSignature?.trim() === "") {
446
+ if (isSameModel && isSigningAnthropicTarget && sanitized.thinkingSignature?.trim() === "") {
442
447
  return [];
443
448
  }
444
449
  return sanitized;
@@ -1,17 +1,20 @@
1
1
  import { scheduler } from "node:timers/promises";
2
+ import { bareModelId, parseAnthropicModel } from "@oh-my-pi/pi-catalog/identity";
2
3
  import { toNumber } from "@oh-my-pi/pi-catalog/utils";
3
4
  import * as AIError from "../error";
4
5
  import { claudeCodeVersion } from "../providers/anthropic";
5
- import type {
6
- CredentialRankingStrategy,
7
- UsageAmount,
8
- UsageFetchContext,
9
- UsageFetchParams,
10
- UsageLimit,
11
- UsageProvider,
12
- UsageReport,
13
- UsageStatus,
14
- UsageWindow,
6
+ import {
7
+ type CredentialRankingContext,
8
+ type CredentialRankingStrategy,
9
+ resolveUsedFraction,
10
+ type UsageAmount,
11
+ type UsageFetchContext,
12
+ type UsageFetchParams,
13
+ type UsageLimit,
14
+ type UsageProvider,
15
+ type UsageReport,
16
+ type UsageStatus,
17
+ type UsageWindow,
15
18
  } from "../usage";
16
19
  import { isRecord } from "../utils";
17
20
 
@@ -60,13 +63,37 @@ interface ParsedUsageBucket {
60
63
  utilization?: number;
61
64
  resetsAt?: number;
62
65
  }
63
- type ClaudeUnifiedWindow = "5h" | "7d";
66
+ type ClaudeUnifiedWindow = "5h" | "7d" | "7d_oi";
67
+ type ClaudeModelKind = "opus" | "sonnet" | "fable" | "mythos";
64
68
 
65
69
  interface ClaudeUsageResponse {
66
70
  five_hour?: ClaudeUsageBucket | null;
67
71
  seven_day?: ClaudeUsageBucket | null;
68
72
  seven_day_opus?: ClaudeUsageBucket | null;
69
73
  seven_day_sonnet?: ClaudeUsageBucket | null;
74
+ limits?: unknown;
75
+ }
76
+
77
+ interface ClaudeApiLimitModelScope {
78
+ display_name?: string | null;
79
+ }
80
+
81
+ interface ClaudeApiLimitScope {
82
+ model?: ClaudeApiLimitModelScope | null;
83
+ }
84
+
85
+ interface ClaudeApiLimitEntry {
86
+ kind?: string;
87
+ percent?: unknown;
88
+ resets_at?: string | null;
89
+ scope?: ClaudeApiLimitScope | null;
90
+ is_active?: boolean;
91
+ }
92
+
93
+ interface ParsedApiLimitEntry {
94
+ kind: string;
95
+ bucket: ParsedUsageBucket;
96
+ displayName?: string;
70
97
  }
71
98
 
72
99
  type ClaudeUsagePayload = {
@@ -89,6 +116,43 @@ function parseBucket(bucket: unknown): ParsedUsageBucket | undefined {
89
116
  }
90
117
  return { utilization, resetsAt };
91
118
  }
119
+
120
+ function getApiLimitDisplayName(scope: unknown): string | undefined {
121
+ if (!isRecord(scope)) return undefined;
122
+ const model = scope.model;
123
+ if (!isRecord(model)) return undefined;
124
+ const displayName = model.display_name;
125
+ return typeof displayName === "string" && displayName.trim() ? displayName.trim() : undefined;
126
+ }
127
+
128
+ /**
129
+ * Anthropic kept the legacy account-wide buckets populated, but as of
130
+ * 2026-07-02 the legacy per-model weekly buckets (`seven_day_opus` /
131
+ * `seven_day_sonnet`) are permanently null. Model-scoped weekly caps now arrive
132
+ * only through generic `limits[]` entries (`kind: "weekly_scoped"`) with the
133
+ * model family named by `scope.model.display_name`.
134
+ */
135
+ function parseApiLimitEntries(raw: unknown): ParsedApiLimitEntry[] {
136
+ if (!Array.isArray(raw)) return [];
137
+ const entries: ParsedApiLimitEntry[] = [];
138
+ for (const rawEntry of raw) {
139
+ if (!isRecord(rawEntry)) continue;
140
+ const entry = rawEntry as ClaudeApiLimitEntry;
141
+ if (typeof entry.kind !== "string") continue;
142
+ if (entry.is_active === false) continue;
143
+ const utilization = toNumber(entry.percent);
144
+ const resetsAt = parseIsoTime(typeof entry.resets_at === "string" ? entry.resets_at : undefined);
145
+ if (utilization === undefined && resetsAt === undefined) continue;
146
+ const displayName = getApiLimitDisplayName(entry.scope);
147
+ entries.push({
148
+ kind: entry.kind,
149
+ bucket: { utilization, resetsAt },
150
+ ...(displayName ? { displayName } : {}),
151
+ });
152
+ }
153
+ return entries;
154
+ }
155
+
92
156
  function parseUnifiedWindow(
93
157
  headers: Record<string, string>,
94
158
  window: ClaudeUnifiedWindow,
@@ -144,12 +208,17 @@ function hasUsageData(payload: ClaudeUsageResponse): boolean {
144
208
  parseBucket(payload.five_hour)?.utilization !== undefined ||
145
209
  parseBucket(payload.seven_day)?.utilization !== undefined ||
146
210
  parseBucket(payload.seven_day_opus)?.utilization !== undefined ||
147
- parseBucket(payload.seven_day_sonnet)?.utilization !== undefined
211
+ parseBucket(payload.seven_day_sonnet)?.utilization !== undefined ||
212
+ parseApiLimitEntries(payload.limits).some(entry => entry.bucket.utilization !== undefined)
148
213
  );
149
214
  }
150
215
 
151
216
  function isRetryableStatus(status: number): boolean {
152
- return AIError.isTransientStatus(status);
217
+ // Exclude 429: the usage endpoint is informational and rate-limited per
218
+ // source IP, so retrying a rate_limit_error inside a single fetch can't
219
+ // succeed and only deepens the throttle (3 attempts per poll). Fall through
220
+ // to the caller's failure cool-down and retry on the next poll instead.
221
+ return AIError.isTransientStatus(status) && status !== 429;
153
222
  }
154
223
 
155
224
  function isAbortError(error: unknown, signal?: AbortSignal): boolean {
@@ -334,8 +403,8 @@ function buildUsageLimit(args: {
334
403
  scope: {
335
404
  provider: args.provider,
336
405
  windowId: args.windowId,
337
- tier: args.tier,
338
- shared: args.shared,
406
+ ...(args.tier !== undefined ? { tier: args.tier } : {}),
407
+ ...(args.shared !== undefined ? { shared: args.shared } : {}),
339
408
  },
340
409
  window,
341
410
  amount,
@@ -343,9 +412,47 @@ function buildUsageLimit(args: {
343
412
  };
344
413
  }
345
414
 
415
+ function slugifyClaudeLimitDisplayName(displayName: string): string {
416
+ return displayName
417
+ .trim()
418
+ .toLowerCase()
419
+ .replace(/[^a-z0-9]+/g, "-")
420
+ .replace(/^-+|-+$/g, "");
421
+ }
422
+
423
+ /**
424
+ * Scoped weekly rows are per-model-family counters, not account-wide windows.
425
+ * They deliberately leave `scope.shared` unset so credential-wide exhaustion
426
+ * gating only considers the shared umbrella windows; an exhausted Fable weekly
427
+ * cap must not block Opus or Sonnet requests on the same credential.
428
+ */
429
+ function buildScopedWeeklyUsageLimits(entries: readonly ParsedApiLimitEntry[]): UsageLimit[] {
430
+ const seenSlugs = new Set<string>();
431
+ const limits: UsageLimit[] = [];
432
+ for (const entry of entries) {
433
+ if (entry.kind !== "weekly_scoped" || !entry.displayName) continue;
434
+ const slug = slugifyClaudeLimitDisplayName(entry.displayName);
435
+ if (!slug || seenSlugs.has(slug)) continue;
436
+ seenSlugs.add(slug);
437
+ const limit = buildUsageLimit({
438
+ id: `anthropic:7d:${slug}`,
439
+ label: `Claude 7 Day (${entry.displayName})`,
440
+ windowId: "7d",
441
+ windowLabel: "7 Day",
442
+ durationMs: SEVEN_DAYS_MS,
443
+ bucket: entry.bucket,
444
+ provider: "anthropic",
445
+ tier: slug,
446
+ });
447
+ if (limit) limits.push(limit);
448
+ }
449
+ return limits;
450
+ }
451
+
346
452
  export function parseClaudeRateLimitHeaders(headers: Record<string, string>, now = Date.now()): UsageReport | null {
347
453
  const fiveHour = parseUnifiedWindow(headers, "5h");
348
454
  const sevenDay = parseUnifiedWindow(headers, "7d");
455
+ const modelScopedSevenDay = parseUnifiedWindow(headers, "7d_oi");
349
456
  const limits = [
350
457
  buildUsageLimit({
351
458
  id: "anthropic:5h",
@@ -367,6 +474,16 @@ export function parseClaudeRateLimitHeaders(headers: Record<string, string>, now
367
474
  provider: "anthropic",
368
475
  shared: true,
369
476
  }),
477
+ buildUsageLimit({
478
+ id: "anthropic:7d:fable",
479
+ label: "Claude 7 Day (Fable)",
480
+ windowId: "7d",
481
+ windowLabel: "7 Day",
482
+ durationMs: SEVEN_DAYS_MS,
483
+ bucket: modelScopedSevenDay,
484
+ provider: "anthropic",
485
+ tier: "fable",
486
+ }),
370
487
  ].filter((limit): limit is UsageLimit => limit !== null);
371
488
 
372
489
  if (limits.length === 0) return null;
@@ -394,8 +511,10 @@ async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext
394
511
  if (!payloadResult || !isRecord(payloadResult.payload)) return null;
395
512
  const { payload, orgId } = payloadResult;
396
513
 
397
- const fiveHour = parseBucket(payload.five_hour);
398
- const sevenDay = parseBucket(payload.seven_day);
514
+ const apiLimitEntries = parseApiLimitEntries(payload.limits);
515
+ const fiveHour = parseBucket(payload.five_hour) ?? apiLimitEntries.find(entry => entry.kind === "session")?.bucket;
516
+ const sevenDay =
517
+ parseBucket(payload.seven_day) ?? apiLimitEntries.find(entry => entry.kind === "weekly_all")?.bucket;
399
518
  const sevenDayOpus = parseBucket(payload.seven_day_opus);
400
519
  const sevenDaySonnet = parseBucket(payload.seven_day_sonnet);
401
520
 
@@ -440,6 +559,7 @@ async function fetchClaudeUsage(params: UsageFetchParams, ctx: UsageFetchContext
440
559
  provider: "anthropic",
441
560
  tier: "sonnet",
442
561
  }),
562
+ ...buildScopedWeeklyUsageLimits(apiLimitEntries),
443
563
  ].filter((limit): limit is UsageLimit => limit !== null);
444
564
 
445
565
  if (limits.length === 0) return null;
@@ -475,11 +595,84 @@ export const claudeUsageProvider: UsageProvider = {
475
595
  supports: params => params.provider === "anthropic" && params.credential.type === "oauth",
476
596
  };
477
597
 
598
+ function getClaudeModelKind(context: CredentialRankingContext | undefined): ClaudeModelKind | undefined {
599
+ const modelId = context?.modelId;
600
+ if (!modelId) return undefined;
601
+ return parseAnthropicModel(bareModelId(modelId))?.kind;
602
+ }
603
+
604
+ /**
605
+ * Claude model-scoped rows are only relevant to the matching model family.
606
+ * Credential-wide exhaustion checks stay on shared umbrella windows unless the
607
+ * request model parses to a concrete Anthropic kind, preventing a Fable cap from
608
+ * suppressing unrelated Opus/Sonnet traffic.
609
+ */
610
+ function scopeClaudeLimitsForModel(report: UsageReport, context: CredentialRankingContext | undefined): UsageLimit[] {
611
+ const kind = getClaudeModelKind(context);
612
+ return report.limits.filter(
613
+ limit => limit.scope.shared === true || (kind !== undefined && limit.scope.tier === kind),
614
+ );
615
+ }
616
+
617
+ function rankingUsedFraction(limit: UsageLimit): number {
618
+ const fraction = resolveUsedFraction(limit);
619
+ if (typeof fraction !== "number" || !Number.isFinite(fraction)) return 0.5;
620
+ return Math.min(Math.max(fraction, 0), 1);
621
+ }
622
+
623
+ function rankingDrainRate(limit: UsageLimit, nowMs: number): number {
624
+ const usedFraction = rankingUsedFraction(limit);
625
+ const durationMs = limit.window?.durationMs ?? SEVEN_DAYS_MS;
626
+ if (!Number.isFinite(durationMs) || durationMs <= 0) return usedFraction;
627
+ const resetAt = limit.window?.resetsAt;
628
+ if (typeof resetAt !== "number" || !Number.isFinite(resetAt)) return usedFraction;
629
+ const remainingWindowMs = resetAt - nowMs;
630
+ const clampedRemainingWindowMs = Math.min(Math.max(remainingWindowMs, 0), durationMs);
631
+ const elapsedMs = durationMs - clampedRemainingWindowMs;
632
+ if (elapsedMs <= 0) return usedFraction;
633
+ const elapsedHours = elapsedMs / (60 * 60 * 1000);
634
+ if (!Number.isFinite(elapsedHours) || elapsedHours <= 0) return usedFraction;
635
+ return usedFraction / elapsedHours;
636
+ }
637
+
638
+ function morePressuredLimit(
639
+ left: UsageLimit | undefined,
640
+ right: UsageLimit | undefined,
641
+ nowMs: number,
642
+ ): UsageLimit | undefined {
643
+ if (!left) return right;
644
+ if (!right) return left;
645
+ const leftDrainRate = rankingDrainRate(left, nowMs);
646
+ const rightDrainRate = rankingDrainRate(right, nowMs);
647
+ if (rightDrainRate !== leftDrainRate) return rightDrainRate > leftDrainRate ? right : left;
648
+ return rankingUsedFraction(right) > rankingUsedFraction(left) ? right : left;
649
+ }
650
+
651
+ function findClaudeSecondaryLimit(
652
+ report: UsageReport,
653
+ context: CredentialRankingContext | undefined,
654
+ ): UsageLimit | undefined {
655
+ const nowMs = Date.now();
656
+ return scopeClaudeLimitsForModel(report, context)
657
+ .filter(limit => limit.scope.windowId === "7d" || limit.window?.id === "7d")
658
+ .reduce<UsageLimit | undefined>((selected, limit) => morePressuredLimit(selected, limit, nowMs), undefined);
659
+ }
660
+
478
661
  export const claudeRankingStrategy: CredentialRankingStrategy = {
479
- findWindowLimits(report) {
480
- const primary = report.limits.find(l => l.id === "anthropic:5h");
481
- const secondary = report.limits.find(l => l.id === "anthropic:7d");
662
+ findWindowLimits(report, context) {
663
+ const primary = report.limits.find(limit => limit.id === "anthropic:5h");
664
+ const secondary = findClaudeSecondaryLimit(report, context);
482
665
  return { primary, secondary };
483
666
  },
667
+ scopeLimits: scopeClaudeLimitsForModel,
668
+ /**
669
+ * Fable/Mythos usage-limit errors map to tier-local weekly counters. Scope
670
+ * reactive backoff blocks for those tiers, mirroring the per-counter
671
+ * precedent in packages/ai/src/usage/google-antigravity.ts:466-497.
672
+ */
673
+ blockScope(context) {
674
+ const kind = getClaudeModelKind(context);
675
+ return kind === "fable" || kind === "mythos" ? `tier:${kind}` : undefined;
676
+ },
484
677
  windowDefaults: { primaryMs: 5 * 60 * 60 * 1000, secondaryMs: 7 * 24 * 60 * 60 * 1000 },
485
678
  };
@@ -30,3 +30,19 @@ export const kStreamingArgumentsDone = Symbol("provider.block.argumentsDone");
30
30
 
31
31
  /** Classifies Cursor's in-flight tool-call kind without leaking provider-private state. */
32
32
  export const kStreamingBlockKind = Symbol("provider.block.kind");
33
+
34
+ /**
35
+ * Marks a `toolCall` content block that Cursor's exec channel already
36
+ * executed server-side (via the coding-agent bridge) and whose result is
37
+ * buffered separately for emission via the assistant-loop stream.
38
+ *
39
+ * `agent-loop.ts` MUST skip execution of blocks carrying this marker —
40
+ * treating them as a fresh runnable tool call would run the same
41
+ * side-effecting tool (bash, write, delete, …) a second time. Symbol-keyed
42
+ * so it never persists across the JSONL round-trip, where rebuild instead
43
+ * pairs the block with its already-persisted `toolResult` message by id.
44
+ */
45
+ export const kCursorExecResolved = Symbol("provider.block.cursorExecResolved");
46
+
47
+ /** Carries the resolved marker without exposing a string-keyed property. */
48
+ export type CursorExecResolvedCarrier = object & { [kCursorExecResolved]?: true };
@@ -162,11 +162,26 @@ export function wrapFetchForProxy(fetchImpl: FetchImpl, provider: string): Fetch
162
162
  return wrapped;
163
163
  }
164
164
 
165
+ export interface ConnectProxiedSocketOptions {
166
+ /** Caller cancellation for the proxy TCP/TLS handshake and CONNECT tunnel. */
167
+ signal?: AbortSignal;
168
+ /** Maximum wall-clock time to establish the final TLS tunnel. Disabled when absent or non-positive. */
169
+ timeoutMs?: number;
170
+ }
171
+
165
172
  /**
166
173
  * Tunnel a socket connection through an HTTP CONNECT proxy.
167
174
  * This is used specifically to wrap Node's `http2.connect(baseUrl, { createConnection })` for Cursor.
168
175
  */
169
- export async function connectProxiedSocket(proxyUrlStr: string, targetUrlStr: string): Promise<tls.TLSSocket> {
176
+ export async function connectProxiedSocket(
177
+ proxyUrlStr: string,
178
+ targetUrlStr: string,
179
+ options?: ConnectProxiedSocketOptions,
180
+ ): Promise<tls.TLSSocket> {
181
+ if (options?.signal?.aborted) {
182
+ throw new AIError.AbortError("Proxy tunnel aborted");
183
+ }
184
+
170
185
  const proxyUrl = new URL(proxyUrlStr);
171
186
  const targetUrl = new URL(targetUrlStr);
172
187
 
@@ -179,23 +194,73 @@ export async function connectProxiedSocket(proxyUrlStr: string, targetUrlStr: st
179
194
 
180
195
  const { promise, resolve, reject } = Promise.withResolvers<tls.TLSSocket>();
181
196
 
182
- let rawSocket: net.Socket;
183
- if (useProxySsl) {
184
- rawSocket = tls.connect({
185
- host: proxyHost,
186
- port: proxyPort,
187
- });
188
- } else {
189
- rawSocket = net.connect({
190
- host: proxyHost,
191
- port: proxyPort,
192
- });
193
- }
194
-
195
- rawSocket.once("error", reject);
196
-
197
197
  const readyEvent = useProxySsl ? "secureConnect" : "connect";
198
- rawSocket.once(readyEvent, () => {
198
+ let rawSocket: net.Socket | undefined;
199
+ let tunnelSocket: tls.TLSSocket | undefined;
200
+ let timeout: NodeJS.Timeout | undefined;
201
+ let responseData = "";
202
+ let settled = false;
203
+
204
+ const cleanup = (): void => {
205
+ if (timeout) {
206
+ clearTimeout(timeout);
207
+ timeout = undefined;
208
+ }
209
+ options?.signal?.removeEventListener("abort", onAbort);
210
+ rawSocket?.off("error", onRawError);
211
+ rawSocket?.off(readyEvent, onProxyReady);
212
+ rawSocket?.off("data", onProxyData);
213
+ tunnelSocket?.off("secureConnect", onTunnelReady);
214
+ tunnelSocket?.off("error", onTunnelError);
215
+ };
216
+ const destroyInProgress = (): void => {
217
+ tunnelSocket?.destroy();
218
+ rawSocket?.destroy();
219
+ };
220
+ const rejectOnce = (error: Error): void => {
221
+ if (settled) return;
222
+ settled = true;
223
+ cleanup();
224
+ destroyInProgress();
225
+ reject(error);
226
+ };
227
+ const resolveOnce = (socket: tls.TLSSocket): void => {
228
+ if (settled) return;
229
+ settled = true;
230
+ cleanup();
231
+ resolve(socket);
232
+ };
233
+ const onAbort = (): void => rejectOnce(new AIError.AbortError("Proxy tunnel aborted"));
234
+ const onRawError = (error: Error): void => rejectOnce(error);
235
+ const onTunnelError = (error: Error): void => rejectOnce(error);
236
+ const onTunnelReady = (): void => {
237
+ if (!tunnelSocket) return;
238
+ resolveOnce(tunnelSocket);
239
+ };
240
+ const onProxyData = (chunk: Buffer): void => {
241
+ if (!rawSocket) return;
242
+ responseData += chunk.toString("binary");
243
+ if (!responseData.includes("\r\n\r\n")) return;
244
+
245
+ rawSocket.off("data", onProxyData);
246
+ rawSocket.off("error", onRawError);
247
+
248
+ const firstLine = responseData.split("\r\n")[0];
249
+ if (!firstLine.includes(" 200 ")) {
250
+ rejectOnce(new AIError.ValidationError(`Proxy tunnel failed: ${firstLine}`));
251
+ return;
252
+ }
253
+
254
+ tunnelSocket = tls.connect({
255
+ socket: rawSocket,
256
+ servername: targetHost,
257
+ ALPNProtocols: ["h2"],
258
+ });
259
+ tunnelSocket.once("secureConnect", onTunnelReady);
260
+ tunnelSocket.once("error", onTunnelError);
261
+ };
262
+ const onProxyReady = (): void => {
263
+ if (!rawSocket) return;
199
264
  let connectReq = `CONNECT ${targetHost}:${targetPort} HTTP/1.1\r\n` + `Host: ${targetHost}:${targetPort}\r\n`;
200
265
 
201
266
  if (proxyUrl.username || proxyUrl.password) {
@@ -207,34 +272,29 @@ export async function connectProxiedSocket(proxyUrlStr: string, targetUrlStr: st
207
272
  connectReq += "\r\n";
208
273
 
209
274
  rawSocket.write(connectReq);
275
+ rawSocket.on("data", onProxyData);
276
+ };
210
277
 
211
- let responseData = "";
212
- const onData = (chunk: Buffer) => {
213
- responseData += chunk.toString("binary");
214
- if (responseData.includes("\r\n\r\n")) {
215
- rawSocket.off("data", onData);
216
- rawSocket.off("error", reject);
217
-
218
- const firstLine = responseData.split("\r\n")[0];
219
- if (firstLine.includes(" 200 ")) {
220
- const tlsSocket = tls.connect({
221
- socket: rawSocket,
222
- servername: targetHost,
223
- ALPNProtocols: ["h2"],
224
- });
225
-
226
- tlsSocket.once("secureConnect", () => {
227
- resolve(tlsSocket);
228
- });
229
- tlsSocket.once("error", reject);
230
- } else {
231
- rawSocket.destroy();
232
- reject(new AIError.ValidationError(`Proxy tunnel failed: ${firstLine}`));
233
- }
234
- }
235
- };
236
- rawSocket.on("data", onData);
237
- });
278
+ options?.signal?.addEventListener("abort", onAbort, { once: true });
279
+ if (options?.timeoutMs !== undefined && Number.isFinite(options.timeoutMs) && options.timeoutMs > 0) {
280
+ const timeoutMs = Math.trunc(options.timeoutMs);
281
+ timeout = setTimeout(() => {
282
+ rejectOnce(new AIError.StreamTimeoutError(`Proxy tunnel timed out after ${timeoutMs}ms`));
283
+ }, timeoutMs);
284
+ timeout.unref?.();
285
+ }
286
+
287
+ rawSocket = useProxySsl
288
+ ? tls.connect({
289
+ host: proxyHost,
290
+ port: proxyPort,
291
+ })
292
+ : net.connect({
293
+ host: proxyHost,
294
+ port: proxyPort,
295
+ });
296
+ rawSocket.once("error", onRawError);
297
+ rawSocket.once(readyEvent, onProxyReady);
238
298
 
239
299
  return promise;
240
300
  }