diff --git a/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts b/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts index 520fea7c0575b..eb66918c913bf 100644 --- a/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts +++ b/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts @@ -23,7 +23,10 @@ import { RunAttemptId, RunId, RuntimeRequestId, + type ServerProvider, TurnItemId, + UsageLimitSourceId, + type UsageLimitSourceSnapshot, } from "@t3tools/contracts"; import * as Cause from "effect/Cause"; import * as DateTime from "effect/DateTime"; @@ -48,6 +51,7 @@ import * as ProjectionStore from "./ProjectionStore.ts"; import * as ProviderEventIngestor from "./ProviderEventIngestor.ts"; import * as ProviderUsageAnalytics from "../telemetry/ProviderUsageAnalytics.ts"; import * as ThreadCommandExecutor from "./ThreadCommandExecutor.ts"; +import { UsageLimitSources } from "../usage/UsageLimitSources.ts"; import * as ProviderInstanceRegistry from "../provider/ProviderInstanceRegistry.ts"; import * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; import type { ProviderInstance } from "@t3tools/provider-core/server/driver"; @@ -326,6 +330,137 @@ it.effect("refreshes the failing instance snapshot after a committed usage_limit ), ); +it.effect( + "takes a proxied usage limit's missing reset from the configured hub's exhausted pool", + () => + Effect.scoped( + Effect.gen(function* () { + const instanceId = ProviderInstanceId.make("claudeAgent"); + const driverKind = ProviderDriverKind.make("claudeAgent"); + const adapter = yield* ProviderAdapter.ProviderAdapterV2.pipe( + Effect.provide( + Layer.mock(ProviderAdapter.ProviderAdapterV2)({ instanceId, driver: driverKind }), + ), + ); + const instance: ProviderInstance = { + instanceId, + driverKind, + continuationIdentity: { driverKind, continuationKey: "proxy" }, + enabled: true, + displayName: undefined, + orchestrationAdapter: adapter, + textGeneration: { + generateBranchName: () => Effect.die("unused"), + generateCommitMessage: () => Effect.die("unused"), + generatePrContent: () => Effect.die("unused"), + generateThreadTitle: () => Effect.die("unused"), + generateReviewGuide: () => Effect.die("unused"), + generateTranscriptionPostProcessing: () => Effect.die("unused"), + }, + snapshot: { + refreshUsage: Effect.void, + // A Claude instance pointed at a hub through ANTHROPIC_BASE_URL. + getSnapshot: Effect.succeed({ + auth: { + status: "authenticated", + type: "customEndpoint", + label: "Custom API endpoint", + }, + } as unknown as ServerProvider), + refresh: Effect.die("unused"), + resolveMaintenance: () => Effect.die("unused"), + applyUsageLimits: () => Effect.die("unused"), + streamChanges: Stream.empty, + }, + }; + const instances = Layer.mock(ProviderInstanceRegistry.ProviderInstanceRegistry)({ + getInstance: (id) => Effect.succeed(id === instanceId ? instance : undefined), + }); + const firstRefill = "2033-05-19T07:20:00.000Z"; + const exhausted = (id: string, resetsAt: string) => ({ + id, + driver: driverKind, + usageLimits: { + checkedAt: "2033-05-19T05:00:00.000Z", + windows: [ + { id: "five_hour", kind: "session", label: "Session", usedPercent: 100, resetsAt }, + ], + }, + }); + const hub: UsageLimitSourceSnapshot = { + id: UsageLimitSourceId.make("hub"), + kind: "cliproxy", + label: "Hub", + checkedAt: "2033-05-19T05:00:00.000Z", + accounts: [ + exhausted("a.json", "2033-05-19T09:00:00.000Z"), + exhausted("b.json", firstRefill), + ], + } as UsageLimitSourceSnapshot; + // The background snapshot predates the request that exhausted the pool. + let refreshed = false; + const sources = Layer.mock(UsageLimitSources)({ + current: Effect.sync(() => (refreshed ? [hub] : [])), + refresh: Effect.sync(() => { + refreshed = true; + }), + streamChanges: Stream.empty, + }); + yield* Effect.gen(function* () { + const now = yield* DateTime.now; + const eventSink = yield* EventSink.EventSinkV2; + const ingestor = yield* ProviderEventIngestor.ProviderEventIngestorV2; + const idAllocator = yield* IdAllocator.IdAllocatorV2; + const threadEvent = yield* threadCreatedEvent(now); + yield* eventSink.write({ events: [threadEvent] }); + const stored = yield* ingestor.ingestNormalized({ + providerSessionId: yield* idAllocator.allocate.providerSession({ + providerInstanceId: instanceId, + threadId: threadEvent.threadId, + }), + providerInstanceId: instanceId, + threadId: threadEvent.threadId, + event: { + type: "turn.terminal", + driver: driverKind, + providerThreadId: idAllocator.derive.providerThread({ + driver: driverKind, + nativeThreadId: "limited", + }), + providerTurnId: idAllocator.derive.providerTurn({ + driver: driverKind, + nativeTurnId: "pool-cooldown", + }), + runOrdinal: 1, + failureItemOrdinal: 1, + status: "failed", + // regression: #26574 the CLI drops the hub's Retry-After, so the native failure has no reset. + failure: makeProviderFailure({ + class: "usage_limit", + message: "Claude API rate limit reached. Try again later.", + code: "api_error_429", + retryable: true, + resetAt: null, + }), + threadDisposition: "reusable", + }, + }); + const item = stored[0]?.event; + assert.equal( + item?.type === "turn-item.updated" && item.payload.type === "error" + ? item.payload.failure.resetAt + : undefined, + firstRefill, + ); + }).pipe( + Effect.provide( + makeLayerTest(layerTestEventSink, instances).pipe(Layer.provideMerge(sources)), + ), + ); + }), + ), +); + it.effect("records accepted billed turn usage once without billing the context window", () => { const recorded: Array = []; const layerAnalytics = Layer.succeed(ProviderUsageAnalytics.ProviderTurnAnalytics, { diff --git a/apps/server/src/orchestration-v2/ProviderEventIngestor.ts b/apps/server/src/orchestration-v2/ProviderEventIngestor.ts index 623fddb491108..fe619196d5e7a 100644 --- a/apps/server/src/orchestration-v2/ProviderEventIngestor.ts +++ b/apps/server/src/orchestration-v2/ProviderEventIngestor.ts @@ -2,6 +2,7 @@ import { NodeId, CommandId, OrchestrationV2DomainEvent, + type OrchestrationV2ProviderFailure, OrchestrationV2StoredEvent, type OrchestrationV2PlanArtifact, type OrchestrationV2Run, @@ -21,6 +22,7 @@ import * as Context from "effect/Context"; import * as DateTime from "effect/DateTime"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; +import * as Option from "effect/Option"; import * as Schema from "effect/Schema"; import * as SqlError from "effect/sql/SqlError"; @@ -32,6 +34,7 @@ import { type ProviderTurnAnalyticsContext, } from "../telemetry/ProviderUsageAnalytics.ts"; import * as GitVcsDriver from "../vcs/GitVcsDriver.ts"; +import { exhaustedPoolRefillAt, UsageLimitSources } from "../usage/UsageLimitSources.ts"; import * as ProviderInstanceRegistry from "../provider/ProviderInstanceRegistry.ts"; import * as EventSink from "./EventSink.ts"; import * as ProjectionStore from "./ProjectionStore.ts"; @@ -255,9 +258,38 @@ export const layer: Layer.Layer< const analytics = yield* ProviderTurnAnalytics; const agentScope = yield* AgentScope; const instances = yield* ProviderInstanceRegistry.ProviderInstanceRegistry; + const usageLimitSources = yield* Effect.serviceOption(UsageLimitSources); const scope = yield* Effect.scope; const completedTurnAnalytics = new Set(); + // Proxied CLIs drop the hub's Retry-After, so take the reset from the hub's own snapshot. + const withProxiedPoolReset = Effect.fnUntraced(function* ( + input: ProviderEventIngestInput, + failure: OrchestrationV2ProviderFailure, + ) { + if ( + failure.class !== "usage_limit" || + failure.resetAt != null || + Option.isNone(usageLimitSources) + ) + return failure; + const instance = yield* instances.getInstance(input.providerInstanceId); + if (!instance) return failure; + const provider = yield* instance.snapshot.getSnapshot; + const proxied = + provider.auth.type === "customEndpoint" || + provider.usageLimits?.unavailable?.reason === "unsupported"; + if (!proxied) return failure; + // The cached snapshot usually predates the request that exhausted the pool. + yield* usageLimitSources.value.refresh.pipe(Effect.timeout("5 seconds"), Effect.ignore); + const resetAt = exhaustedPoolRefillAt( + yield* usageLimitSources.value.current, + instance.driverKind, + DateTime.toEpochMillis(yield* DateTime.now), + ); + return resetAt === null ? failure : { ...failure, resetAt }; + }); + const makeDomainEvent = /* This helper runs for individual orchestration events; the enclosing operation owns the span. */ Effect.fnUntraced(function* ( @@ -659,7 +691,7 @@ export const layer: Layer.Layer< const occurredAt = yield* DateTime.now; const failure = (yield* agentScope.oomKilled(input.threadId)) ? { ...input.event.failure, message: "Killed: out of memory.", code: "oom_kill" } - : input.event.failure; + : yield* withProxiedPoolReset(input, input.event.failure); return [ ...dismissed, yield* makeDomainEvent(input, { diff --git a/apps/server/src/usage/UsageLimitSources.ts b/apps/server/src/usage/UsageLimitSources.ts index 7d40b0fb9ecea..8209d7aa671f3 100644 --- a/apps/server/src/usage/UsageLimitSources.ts +++ b/apps/server/src/usage/UsageLimitSources.ts @@ -3,6 +3,7 @@ import { UsageLimitSourceError, type UsageLimitSourceConsumeResetCreditInput, type ProviderConsumeResetCreditResult, + type ProviderDriverKind, type ServerSettings, UsageLimitSourceConfig, settingsSecretLiteral, @@ -40,6 +41,27 @@ export class UsageLimitSources extends Context.Service< } >()("t3/usage/UsageLimitSources") {} +/** When a hub has exhausted every pooled account for `driver`, the earliest moment one refills. */ +export function exhaustedPoolRefillAt( + sources: ReadonlyArray, + driver: ProviderDriverKind, + now: number, +): string | null { + let earliest: number | null = null; + for (const account of sources.flatMap((source) => source.accounts)) { + if (account.driver !== driver) continue; + const refills = account.usageLimits.windows + .filter((window) => window.usedPercent >= 100 && window.resetsAt !== undefined) + .map((window) => Date.parse(window.resetsAt!)) + .filter((at) => at > now); + // An account with no exhausted window may still serve, so the pool has no reset. + if (refills.length === 0) return null; + const refill = Math.max(...refills); + if (earliest === null || refill < earliest) earliest = refill; + } + return earliest === null ? null : DateTime.formatIso(DateTime.makeUnsafe(earliest)); +} + function sourceLabel(id: string, config: UsageLimitSourceConfig): string { if (config.label) return config.label; const fallback = id.trim() || "Usage limit source";