diff --git a/apps/server/src/provider/Drivers/AntigravityDriver.test.ts b/apps/server/src/provider/Drivers/AntigravityDriver.test.ts index 922a27df5768..8e1b38688911 100644 --- a/apps/server/src/provider/Drivers/AntigravityDriver.test.ts +++ b/apps/server/src/provider/Drivers/AntigravityDriver.test.ts @@ -18,6 +18,7 @@ import * as Layer from "effect/Layer"; import * as Path from "effect/Path"; import * as Schema from "effect/Schema"; import * as TestClock from "effect/testing/TestClock"; +import { HttpClient } from "effect/unstable/http"; import * as ChildProcessSpawner from "effect/unstable/process/ChildProcessSpawner"; import * as BackgroundPolicy from "../../background/BackgroundPolicy.ts"; @@ -239,6 +240,12 @@ const testLayer = ServerConfig.layerTest(process.cwd(), { ), Layer.provideMerge(Layer.succeed(ProviderEventLoggers, NoOpProviderEventLoggers)), Layer.provideMerge(ModelManifest.layerTest), + Layer.provideMerge( + Layer.succeed( + HttpClient.HttpClient, + HttpClient.make(() => Effect.die("Disabled Antigravity must not make an HTTP request")), + ), + ), ); it.layer(testLayer)("AntigravityDriver", (it) => { diff --git a/apps/server/src/provider/Drivers/AntigravityDriver.ts b/apps/server/src/provider/Drivers/AntigravityDriver.ts index 65a8c97fe668..b0445799498d 100644 --- a/apps/server/src/provider/Drivers/AntigravityDriver.ts +++ b/apps/server/src/provider/Drivers/AntigravityDriver.ts @@ -1,6 +1,7 @@ import { AntigravitySettings, ProviderDriverKind, ProviderSetupError } from "@t3tools/contracts"; import { HostProcessPlatform } from "@t3tools/shared/hostProcess"; import * as Crypto from "effect/Crypto"; +import * as DateTime from "effect/DateTime"; import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; import * as Exit from "effect/Exit"; @@ -9,6 +10,7 @@ import * as Path from "effect/Path"; import * as Schema from "effect/Schema"; import * as Scope from "effect/Scope"; import * as Stream from "effect/Stream"; +import { HttpClient } from "effect/unstable/http"; import * as ChildProcessSpawner from "effect/unstable/process/ChildProcessSpawner"; import type { AcpError } from "effect-acp/errors"; @@ -41,6 +43,7 @@ import { removeAntigravitySessionFiles } from "../acp/AntigravitySessionFiles.ts import { ProviderDriverError } from "../Errors.ts"; import { makeAntigravityAdapter } from "../Layers/AntigravityAdapter.ts"; import { makeAntigravityProvider } from "../Layers/AntigravityProvider.ts"; +import { probeAntigravityUsageLimits } from "../Layers/antigravityUsageLimits.ts"; import { ProviderEventLoggers } from "../Layers/ProviderEventLoggers.ts"; import * as ModelManifest from "../ModelManifest.ts"; import { @@ -61,6 +64,7 @@ export type AntigravityDriverEnv = | ChildProcessSpawner.ChildProcessSpawner | Crypto.Crypto | FileSystem.FileSystem + | HttpClient.HttpClient | ModelManifest.ModelManifest | Path.Path | ProviderEventLoggers @@ -78,6 +82,7 @@ export const AntigravityDriver: ProviderDriver Effect.succeed(undefined)), + ); + const provider = yield* makeAntigravityProvider(settings, { stampIdentity: classifyModels, probe, @@ -285,6 +303,7 @@ export const AntigravityDriver: ProviderDriver false), ), + usageLimits, }).pipe( Effect.mapError( (cause) => diff --git a/apps/server/src/provider/Layers/AntigravityProvider.test.ts b/apps/server/src/provider/Layers/AntigravityProvider.test.ts index 34fd81a424ca..e373f10af9f0 100644 --- a/apps/server/src/provider/Layers/AntigravityProvider.test.ts +++ b/apps/server/src/provider/Layers/AntigravityProvider.test.ts @@ -5,6 +5,7 @@ import { ProviderDriverKind, ProviderInstanceId, ProviderSetupError, + type ServerProviderUsageLimits, } from "@t3tools/contracts"; import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; @@ -111,7 +112,11 @@ const testLayer = Layer.merge( type ProbeError = EffectAcpErrors.AcpError | ProviderSetupError; const makeHarness = Effect.fn("makeAntigravityProviderHarness")(function* ( - options: { readonly enabled?: boolean; readonly safe?: boolean } = {}, + options: { + readonly enabled?: boolean; + readonly safe?: boolean; + readonly usageLimits?: Effect.Effect; + } = {}, ) { const initialProbe = yield* Deferred.make(); const probeCalls = yield* Ref.make(0); @@ -132,6 +137,7 @@ const makeHarness = Effect.fn("makeAntigravityProviderHarness")(function* ( Effect.andThen(Ref.get(safety)), Effect.flatten, ), + ...(options.usageLimits ? { usageLimits: options.usageLimits } : {}), }, ); const initialUpdate = yield* Stream.toPull( @@ -641,4 +647,142 @@ it.layer(testLayer)("Antigravity provider snapshots", (it) => { }), ), ); + + it.effect("includes usageLimits in snapshot when available and clears on sign-out", () => + Effect.scoped( + Effect.gen(function* () { + const mockLimits: ServerProviderUsageLimits = { + checkedAt: "2026-09-09T00:00:00.000Z", + windows: [ + { + id: "plan_allowance", + kind: "session", + label: "Gemini Code Assist", + usedPercent: 0, + }, + ], + }; + const harness = yield* makeHarness({ + usageLimits: Effect.succeed(mockLimits), + }); + yield* harness.initialize; + const snapshot = yield* harness.provider.snapshot.getSnapshot; + expect(snapshot.usageLimits).toEqual(mockLimits); + + // onSessionStarted keeps/updates usage limits + yield* harness.provider.onSessionStarted(started); + const sessionSnapshot = yield* harness.provider.snapshot.getSnapshot; + expect(sessionSnapshot.usageLimits).toEqual(mockLimits); + + // onSignedOut clears usage limits + yield* harness.provider.onSignedOut; + const signedOutSnapshot = yield* harness.provider.snapshot.getSnapshot; + expect(signedOutSnapshot.usageLimits).toBeUndefined(); + }), + ), + ); + + it.effect("ignores stale onSessionStarted completion when newer session started", () => + Effect.scoped( + Effect.gen(function* () { + const session1Gate = yield* Deferred.make(); + const session1Continue = yield* Deferred.make(); + const session1Done = yield* Deferred.make(); + + let gateArmed = false; + let sessionCalls = 0; + const harness = yield* makeHarness({ + usageLimits: Effect.gen(function* () { + if (!gateArmed) { + return { checkedAt: "2026-09-09T00:00:00.000Z", windows: [] }; + } + sessionCalls++; + if (sessionCalls === 1) { + yield* Deferred.succeed(session1Gate, undefined); + yield* Deferred.await(session1Continue); + yield* Deferred.succeed(session1Done, undefined); + return { checkedAt: "2026-09-09T01:00:00.000Z", windows: [] }; + } + return { + checkedAt: "2026-09-09T02:00:00.000Z", + windows: [{ id: "new_session", kind: "session", label: "New", usedPercent: 10 }], + }; + }), + }); + yield* harness.initialize; + gateArmed = true; + + const session2LimitsApplied = yield* Stream.toPull( + harness.provider.snapshot.streamChanges.pipe( + Stream.filter((s) => s.usageLimits?.checkedAt === "2026-09-09T02:00:00.000Z"), + ), + ); + + const fiber1 = yield* harness.provider.onSessionStarted(started).pipe(Effect.forkChild); + yield* Deferred.await(session1Gate); + + const session2Started: AcpSessionRuntimeStartResult = { + ...started, + sessionSetupResult: { + sessionId: "session-2", + configOptions: [ + { ...modelConfig, currentValue: "gemini-pro-agent", options: [modelOptions[9]!] }, + ], + }, + }; + yield* harness.provider.onSessionStarted(session2Started); + yield* session2LimitsApplied; + + yield* Deferred.succeed(session1Continue, undefined); + yield* Deferred.await(session1Done); + yield* Fiber.join(fiber1); + + const finalSnapshot = yield* harness.provider.snapshot.getSnapshot; + expect(finalSnapshot.models.map((m) => m.slug)).toEqual(["gemini-pro-agent"]); + expect(finalSnapshot.usageLimits?.checkedAt).toBe("2026-09-09T02:00:00.000Z"); + }), + ), + ); + + it.effect("does not restore usageLimits if signed out before quota probe completes", () => + Effect.scoped( + Effect.gen(function* () { + const quotaGate = yield* Deferred.make(); + const quotaContinue = yield* Deferred.make(); + const quotaDone = yield* Deferred.make(); + + let gateArmed = false; + const harness = yield* makeHarness({ + usageLimits: Effect.gen(function* () { + if (!gateArmed) { + return undefined; + } + yield* Deferred.succeed(quotaGate, undefined); + yield* Deferred.await(quotaContinue); + yield* Deferred.succeed(quotaDone, undefined); + return { + checkedAt: "2026-09-09T03:00:00.000Z", + windows: [{ id: "stale_quota", kind: "session", label: "Stale", usedPercent: 50 }], + }; + }), + }); + yield* harness.initialize; + gateArmed = true; + + yield* harness.provider.onSessionStarted(started); + yield* Deferred.await(quotaGate); + + // Sign out while quota probe is still running in background + yield* harness.provider.onSignedOut; + + // Release the deferred quota effect + yield* Deferred.succeed(quotaContinue, undefined); + yield* Deferred.await(quotaDone); + + const finalSnapshot = yield* harness.provider.snapshot.getSnapshot; + expect(finalSnapshot.auth.status).toBe("unauthenticated"); + expect(finalSnapshot.usageLimits).toBeUndefined(); + }), + ), + ); }); diff --git a/apps/server/src/provider/Layers/AntigravityProvider.ts b/apps/server/src/provider/Layers/AntigravityProvider.ts index 956a5d81d3f4..997d753e7b66 100644 --- a/apps/server/src/provider/Layers/AntigravityProvider.ts +++ b/apps/server/src/provider/Layers/AntigravityProvider.ts @@ -6,6 +6,7 @@ import { type ServerProvider, type ServerProviderModel, type ServerProviderSlashCommand, + type ServerProviderUsageLimits, } from "@t3tools/contracts"; import { createModelCapabilities } from "@t3tools/shared/model"; import * as DateTime from "effect/DateTime"; @@ -126,6 +127,7 @@ interface AntigravityProviderOptions { readonly maintenanceCapabilities?: ProviderMaintenanceCapabilities; /** Auth type and label published once a session authenticates. */ readonly auth?: { readonly type: string; readonly label: string }; + readonly usageLimits?: Effect.Effect; } /** Health uses initialize only. Session callbacks supply account-specific metadata. */ @@ -133,6 +135,7 @@ export const makeAntigravityProvider = Effect.fn("makeAntigravityProvider")(func settings: AntigravitySettings, options: AntigravityProviderOptions, ) { + const scope = yield* Effect.scope; const checkedAt = DateTime.formatIso(yield* DateTime.now); const initialDraft = { ...buildServerProvider({ @@ -192,6 +195,10 @@ export const makeAntigravityProvider = Effect.fn("makeAntigravityProvider")(func : `Antigravity did not respond to its local health check within ${HEALTH_CHECK_TIMEOUT}.`; const supportsTextGeneration = initialized !== undefined ? yield* options.supportsTextGeneration : false; + const usageLimits = + initialized !== undefined && options.usageLimits + ? yield* options.usageLimits.pipe(Effect.catchCause(() => Effect.succeed(undefined))) + : undefined; const updatedAt = DateTime.formatIso(yield* DateTime.now); const next = yield* SubscriptionRef.updateAndGet(metadata, (state) => { if (state.authRevision !== before.authRevision) return state; @@ -212,6 +219,7 @@ export const makeAntigravityProvider = Effect.fn("makeAntigravityProvider")(func version: initialized?.agentInfo?.version || draft.version, status: errorMessage ? "error" : authenticated ? "ready" : "warning", checkedAt: updatedAt, + ...(usageLimits ? { usageLimits } : {}), ...(missingInstallation ? { models: [], @@ -262,47 +270,65 @@ export const makeAntigravityProvider = Effect.fn("makeAntigravityProvider")(func const before = yield* SubscriptionRef.get(metadata); const supportsTextGeneration = yield* options.supportsTextGeneration; const updatedAt = DateTime.formatIso(yield* DateTime.now); - yield* SubscriptionRef.update(metadata, (state) => { - if ( - state.authRevision !== before.authRevision && - state.draft.auth.status === "unauthenticated" - ) { - return state; + const updatedRevision = yield* SubscriptionRef.modify(metadata, (state) => { + if (state.authRevision !== before.authRevision) { + return [Option.none(), state] as const; } + const newRevision = state.authRevision + 1; const { message: _previousMessage, ...draft } = state.draft; const workspaces = draft.workspaceSnapshots ?? []; const workspace = cwd ? workspaces.find((entry) => entry.cwd === cwd) : undefined; - return { - authRevision: state.authRevision + 1, - draft: { - ...draft, - installed: true, - status: settings.enabled ? "ready" : "disabled", - version: started.initializeResult.agentInfo?.version || draft.version, - auth: { - status: "authenticated", - type: options.auth?.type ?? "oauth-personal", - label: options.auth?.label ?? "Google account", + return [ + Option.some(newRevision), + { + authRevision: newRevision, + draft: { + ...draft, + installed: true, + status: settings.enabled ? "ready" : "disabled", + version: started.initializeResult.agentInfo?.version || draft.version, + auth: { + status: "authenticated", + type: options.auth?.type ?? "oauth-personal", + label: options.auth?.label ?? "Google account", + }, + checkedAt: updatedAt, + models: buildAntigravityModelsFromSession(started.sessionSetupResult), + supportsTextGeneration, + ...(cwd + ? { + workspaceSnapshots: [ + ...workspaces.filter((entry) => entry.cwd !== cwd), + { + cwd, + checkedAt: updatedAt, + slashCommands: workspace?.slashCommands ?? draft.slashCommands, + skills: workspace?.skills ?? discoveredSkills.get(cwd) ?? [], + }, + ].slice(-MAX_WORKSPACE_SNAPSHOTS), + } + : {}), }, - checkedAt: updatedAt, - models: buildAntigravityModelsFromSession(started.sessionSetupResult), - supportsTextGeneration, - ...(cwd - ? { - workspaceSnapshots: [ - ...workspaces.filter((entry) => entry.cwd !== cwd), - { - cwd, - checkedAt: updatedAt, - slashCommands: workspace?.slashCommands ?? draft.slashCommands, - skills: workspace?.skills ?? discoveredSkills.get(cwd) ?? [], - }, - ].slice(-MAX_WORKSPACE_SNAPSHOTS), - } - : {}), - }, - } satisfies AntigravityProviderState; + } satisfies AntigravityProviderState, + ] as const; }); + + if (Option.isSome(updatedRevision) && options.usageLimits) { + const revision = updatedRevision.value; + yield* options.usageLimits.pipe( + Effect.catchCause(() => Effect.succeed(undefined)), + Effect.flatMap((usageLimits) => + usageLimits + ? SubscriptionRef.update(metadata, (state) => + state.authRevision === revision + ? { ...state, draft: { ...state.draft, usageLimits } } + : state, + ) + : Effect.void, + ), + Effect.forkIn(scope), + ); + } }); const onConfigOptionsUpdated = Effect.fn("AntigravityProvider.onConfigOptionsUpdated")(function* ( @@ -351,25 +377,24 @@ export const makeAntigravityProvider = Effect.fn("makeAntigravityProvider")(func const clearAccountMetadata = Effect.fn("AntigravityProvider.clearAccountMetadata")(function* () { const updatedAt = DateTime.formatIso(yield* DateTime.now); - yield* SubscriptionRef.update( - metadata, - (state) => - ({ - authRevision: state.authRevision + 1, - draft: { - ...state.draft, - auth: { status: "unauthenticated" }, - status: settings.enabled ? "warning" : "disabled", - message: SIGN_IN_MESSAGE, - checkedAt: updatedAt, - models: [], - slashCommands: [], - skills: [], - workspaceSnapshots: [], - supportsTextGeneration: false, - }, - }) satisfies AntigravityProviderState, - ); + yield* SubscriptionRef.update(metadata, (state) => { + const { usageLimits: _previousUsageLimits, ...draft } = state.draft; + return { + authRevision: state.authRevision + 1, + draft: { + ...draft, + auth: { status: "unauthenticated" }, + status: settings.enabled ? "warning" : "disabled", + message: SIGN_IN_MESSAGE, + checkedAt: updatedAt, + models: [], + slashCommands: [], + skills: [], + workspaceSnapshots: [], + supportsTextGeneration: false, + }, + } satisfies AntigravityProviderState; + }); discoveredSkills.clear(); }); diff --git a/apps/server/src/provider/Layers/antigravityUsageLimits.test.ts b/apps/server/src/provider/Layers/antigravityUsageLimits.test.ts new file mode 100644 index 000000000000..b3133b121238 --- /dev/null +++ b/apps/server/src/provider/Layers/antigravityUsageLimits.test.ts @@ -0,0 +1,580 @@ +import * as NodeServices from "@effect/platform-node/NodeServices"; +import * as Effect from "effect/Effect"; +import * as FileSystem from "effect/FileSystem"; +import * as Path from "effect/Path"; +import * as Ref from "effect/Ref"; +import { HttpClient, HttpClientResponse } from "effect/unstable/http"; +import { describe, expect, it } from "@effect/vitest"; + +import { + antigravityRateLimitsToLimits, + antigravityRateLimitsToUpdate, + antigravityRateLimitsToWindows, + probeAntigravityUsageLimits, + quotaSummaryToWindows, +} from "./antigravityUsageLimits.ts"; + +const checkedAt = "2026-09-09T10:00:00.000Z"; + +describe("antigravityRateLimitsToWindows", () => { + it("maps explicit quota windows properly", () => { + const windows = antigravityRateLimitsToWindows({ + windows: [ + { + id: "daily_quota", + kind: "session", + label: "Daily Quota", + usedPercent: 42, + windowDurationMins: 1440, + resetsAt: "2026-09-10T00:00:00.000Z", + }, + ], + }); + + expect(windows).toEqual([ + { + id: "daily_quota", + kind: "session", + label: "Daily Quota", + usedPercent: 42, + windowDurationMins: 1440, + resetsAt: "2026-09-10T00:00:00.000Z", + }, + ]); + }); + + it("calculates usage percentage from prompt credits", () => { + const windows = antigravityRateLimitsToWindows({ + monthlyPromptCredits: 1000, + availablePromptCredits: 600, + }); + + expect(windows).toEqual([ + { + id: "prompt_credits", + kind: "monthly", + label: "Prompt Credits", + usedPercent: 40, + windowDurationMins: 43200, + }, + ]); + }); + + it("calculates usage percentage from flow credits", () => { + const windows = antigravityRateLimitsToWindows({ + monthlyFlowCredits: 500, + availableFlowCredits: 100, + }); + + expect(windows).toEqual([ + { + id: "flow_credits", + kind: "monthly", + label: "Flow Credits", + usedPercent: 80, + windowDurationMins: 43200, + }, + ]); + }); + + it("creates a plan allowance window for recognized tiers", () => { + const windows = antigravityRateLimitsToWindows({ + allowedTiers: [ + { + id: "standard-tier", + name: "Gemini Code Assist", + description: "Unlimited coding assistant with the most powerful Gemini models", + }, + ], + }); + + expect(windows).toEqual([ + { + id: "plan_allowance", + kind: "session", + label: "Gemini Code Assist", + usedPercent: 0, + }, + ]); + }); +}); + +describe("antigravityRateLimitsToLimits", () => { + it("wraps windows into ServerProviderUsageLimits with checkedAt", () => { + const limits = antigravityRateLimitsToLimits({ + checkedAt, + snapshot: { + monthlyPromptCredits: 100, + availablePromptCredits: 75, + }, + }); + + expect(limits).toEqual({ + checkedAt, + windows: [ + { + id: "prompt_credits", + kind: "monthly", + label: "Prompt Credits", + usedPercent: 25, + windowDurationMins: 43200, + }, + ], + }); + }); +}); + +describe("antigravityRateLimitsToUpdate", () => { + it("produces a ProviderUsageLimitsUpdate when windows exist", () => { + const update = antigravityRateLimitsToUpdate({ + monthlyPromptCredits: 100, + availablePromptCredits: 50, + }); + + expect(update).toEqual({ + windows: [ + { + id: "prompt_credits", + kind: "monthly", + label: "Prompt Credits", + usedPercent: 50, + windowDurationMins: 43200, + }, + ], + }); + }); + + it("returns undefined when no windows exist", () => { + const update = antigravityRateLimitsToUpdate({}); + expect(update).toBeUndefined(); + }); +}); + +describe("quotaSummaryToWindows", () => { + it("maps Gemini and Claude/GPT model groups into 5-hour and weekly windows", () => { + const windows = quotaSummaryToWindows({ + groups: [ + { + displayName: "Gemini Models", + description: "Models within this group: Gemini Flash, Gemini Pro", + buckets: [ + { + bucketId: "gemini-weekly", + displayName: "Weekly Limit Remaining", + window: "weekly", + resetTime: "2026-09-11T06:06:12Z", + remainingFraction: 0.6348, + }, + { + bucketId: "gemini-5h", + displayName: "Five Hour Limit Remaining", + window: "5h", + resetTime: "2026-09-09T11:13:01Z", + remainingFraction: 0.5346, + }, + ], + }, + { + displayName: "Claude and GPT models", + description: "Models within this group: Claude Opus, Claude Sonnet, GPT-OSS", + buckets: [ + { + bucketId: "3p-weekly", + displayName: "Weekly Limit Remaining", + window: "weekly", + resetTime: "2026-09-16T06:14:10Z", + remainingFraction: 1, + }, + { + bucketId: "3p-5h", + displayName: "Five Hour Limit Remaining", + window: "5h", + resetTime: "2026-09-09T11:14:10Z", + remainingFraction: 1, + }, + ], + }, + ], + }); + + expect(windows).toEqual([ + { + id: "gemini-5h", + label: "Gemini (5-hour)", + kind: "session", + usedPercent: 47, + windowDurationMins: 300, + resetsAt: "2026-09-09T11:13:01.000Z", + }, + { + id: "gemini-weekly", + label: "Gemini (Weekly)", + kind: "weekly", + usedPercent: 37, + windowDurationMins: 10080, + resetsAt: "2026-09-11T06:06:12.000Z", + }, + { + id: "3p-5h", + label: "Claude & GPT (5-hour)", + kind: "session", + usedPercent: 0, + windowDurationMins: 300, + resetsAt: "2026-09-09T11:14:10.000Z", + }, + { + id: "3p-weekly", + label: "Claude & GPT (Weekly)", + kind: "weekly", + usedPercent: 0, + windowDurationMins: 10080, + resetsAt: "2026-09-16T06:14:10.000Z", + }, + ]); + }); + + it("ignores disabled buckets and handles malformed inputs safely", () => { + const windows = quotaSummaryToWindows({ + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { + bucketId: "disabled-bucket", + displayName: "Disabled Limit", + window: "5h", + disabled: true, + remainingFraction: 0.5, + }, + { + bucketId: "gemini-5h", + displayName: "Five Hour Limit Remaining", + window: "5h", + resetTime: "invalid-date", + remainingFraction: NaN, + }, + ], + }, + ], + }); + + expect(windows).toEqual([ + { + id: "gemini-5h", + label: "Gemini (5-hour)", + kind: "session", + usedPercent: 0, + windowDurationMins: 300, + }, + ]); + + expect(quotaSummaryToWindows(null)).toEqual([]); + expect(quotaSummaryToWindows(undefined)).toEqual([]); + }); +}); + +describe("probeAntigravityUsageLimits", () => { + it.effect("returns undefined when token file does not exist", () => + Effect.gen(function* () { + const mockHttp = HttpClient.make(() => Effect.die("HTTP should not be called")); + const result = yield* probeAntigravityUsageLimits({ + profileDirectory: "/nonexistent/profile", + checkedAt, + }).pipe( + Effect.provide(NodeServices.layer), + Effect.provideService(HttpClient.HttpClient, mockHttp), + ); + + expect(result).toBeUndefined(); + }), + ); + + it.effect("rejects untrusted token_uri to prevent SSRF", () => + Effect.gen(function* () { + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const tempDir = yield* fs.makeTempDirectoryScoped(); + const acpDir = path.join(tempDir, "antigravity-acp"); + yield* fs.makeDirectory(acpDir); + + yield* fs.writeFileString( + path.join(acpDir, "acp_token.json"), + // @effect-diagnostics-next-line preferSchemaOverJson:off + JSON.stringify({ + client_id: "id", + client_secret: "secret", + refresh_token: "refresh", + token_uri: "http://malicious-server.com/token", + }), + ); + + const calls = yield* Ref.make(0); + const mockHttp = HttpClient.make((request) => + Ref.update(calls, (count) => count + 1).pipe( + Effect.as(HttpClientResponse.fromWeb(request, Response.json({}))), + ), + ); + const result = yield* probeAntigravityUsageLimits({ + profileDirectory: tempDir, + checkedAt, + }).pipe(Effect.provideService(HttpClient.HttpClient, mockHttp)); + + expect(result).toBeUndefined(); + expect(yield* Ref.get(calls)).toBe(0); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), + ); + + it.effect("exchanges token and probes quota summary successfully", () => + Effect.gen(function* () { + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const tempDir = yield* fs.makeTempDirectoryScoped(); + const acpDir = path.join(tempDir, "antigravity-acp"); + yield* fs.makeDirectory(acpDir); + + yield* fs.writeFileString( + path.join(acpDir, "acp_token.json"), + // @effect-diagnostics-next-line preferSchemaOverJson:off + JSON.stringify({ + client_id: "id", + client_secret: "secret", + refresh_token: "refresh", + token_uri: "https://oauth2.googleapis.com/token", + }), + ); + + const mockHttp = HttpClient.make((request) => + Effect.sync(() => { + const url = request.url; + if (url === "https://oauth2.googleapis.com/token") { + return HttpClientResponse.fromWeb( + request, + Response.json({ access_token: "mock-access-token" }), + ); + } + if (url.includes("retrieveUserQuotaSummary")) { + return HttpClientResponse.fromWeb( + request, + Response.json({ + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { + bucketId: "gemini-5h", + displayName: "Five Hour Limit Remaining", + window: "5h", + resetTime: "2026-09-09T15:00:00.000Z", + remainingFraction: 0.75, + }, + ], + }, + ], + }), + ); + } + return HttpClientResponse.fromWeb(request, Response.json({}, { status: 404 })); + }), + ); + + const result = yield* probeAntigravityUsageLimits({ + profileDirectory: tempDir, + checkedAt, + }).pipe(Effect.provideService(HttpClient.HttpClient, mockHttp)); + + expect(result).toBeDefined(); + expect(result?.windows).toHaveLength(1); + expect(result?.windows[0]?.id).toBe("gemini-5h"); + expect(result?.windows[0]?.usedPercent).toBe(25); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), + ); + + it.effect("probes and orders multi-group quota summary buckets grouped by model family", () => + Effect.gen(function* () { + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const tempDir = yield* fs.makeTempDirectoryScoped(); + const acpDir = path.join(tempDir, "antigravity-acp"); + yield* fs.makeDirectory(acpDir); + + yield* fs.writeFileString( + path.join(acpDir, "acp_token.json"), + // @effect-diagnostics-next-line preferSchemaOverJson:off + JSON.stringify({ + client_id: "id", + client_secret: "secret", + refresh_token: "refresh", + token_uri: "https://oauth2.googleapis.com/token", + }), + ); + + const mockHttp = HttpClient.make((request) => + Effect.sync(() => { + const url = request.url; + if (url === "https://oauth2.googleapis.com/token") { + return HttpClientResponse.fromWeb( + request, + Response.json({ access_token: "mock-access-token" }), + ); + } + if (url.includes("retrieveUserQuotaSummary")) { + return HttpClientResponse.fromWeb( + request, + Response.json({ + groups: [ + { + displayName: "Gemini Models", + buckets: [ + { + bucketId: "gemini-weekly", + displayName: "Weekly Limit Remaining", + window: "weekly", + remainingFraction: 0.6, + }, + { + bucketId: "gemini-5h", + displayName: "Five Hour Limit Remaining", + window: "5h", + remainingFraction: 0.7, + }, + ], + }, + { + displayName: "Claude and GPT models", + buckets: [ + { + bucketId: "3p-weekly", + displayName: "Weekly Limit Remaining", + window: "weekly", + remainingFraction: 1, + }, + { + bucketId: "3p-5h", + displayName: "Five Hour Limit Remaining", + window: "5h", + remainingFraction: 1, + }, + ], + }, + ], + }), + ); + } + return HttpClientResponse.fromWeb(request, Response.json({}, { status: 404 })); + }), + ); + + const result = yield* probeAntigravityUsageLimits({ + profileDirectory: tempDir, + checkedAt, + }).pipe(Effect.provideService(HttpClient.HttpClient, mockHttp)); + + expect(result).toBeDefined(); + expect(result?.windows.map((w) => w.label)).toEqual([ + "Gemini (5-hour)", + "Gemini (Weekly)", + "Claude & GPT (5-hour)", + "Claude & GPT (Weekly)", + ]); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), + ); + + it.effect("falls back to loadCodeAssist when retrieveUserQuotaSummary fails", () => + Effect.gen(function* () { + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const tempDir = yield* fs.makeTempDirectoryScoped(); + const acpDir = path.join(tempDir, "antigravity-acp"); + yield* fs.makeDirectory(acpDir); + + yield* fs.writeFileString( + path.join(acpDir, "acp_token.json"), + // @effect-diagnostics-next-line preferSchemaOverJson:off + JSON.stringify({ + client_id: "id", + client_secret: "secret", + refresh_token: "refresh", + token_uri: "https://oauth2.googleapis.com/token", + }), + ); + + const mockHttp = HttpClient.make((request) => + Effect.sync(() => { + const url = request.url; + if (url === "https://oauth2.googleapis.com/token") { + return HttpClientResponse.fromWeb( + request, + Response.json({ access_token: "mock-access-token" }), + ); + } + if (url.includes("retrieveUserQuotaSummary")) { + return HttpClientResponse.fromWeb( + request, + Response.json({ error: "unavailable" }, { status: 500 }), + ); + } + if (url.includes("loadCodeAssist")) { + return HttpClientResponse.fromWeb( + request, + Response.json({ + allowedTiers: [{ id: "standard", name: "Gemini Code Assist" }], + quotaManagerState: { + monthlyPromptCredits: 100, + availablePromptCredits: 70, + }, + }), + ); + } + return HttpClientResponse.fromWeb(request, Response.json({}, { status: 404 })); + }), + ); + + const result = yield* probeAntigravityUsageLimits({ + profileDirectory: tempDir, + checkedAt, + }).pipe(Effect.provideService(HttpClient.HttpClient, mockHttp)); + + expect(result).toBeDefined(); + expect(result?.windows).toHaveLength(1); + expect(result?.windows[0]?.id).toBe("prompt_credits"); + expect(result?.windows[0]?.usedPercent).toBe(30); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), + ); + + it.effect("returns undefined when token exchange returns 401", () => + Effect.gen(function* () { + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const tempDir = yield* fs.makeTempDirectoryScoped(); + const acpDir = path.join(tempDir, "antigravity-acp"); + yield* fs.makeDirectory(acpDir); + + yield* fs.writeFileString( + path.join(acpDir, "acp_token.json"), + // @effect-diagnostics-next-line preferSchemaOverJson:off + JSON.stringify({ + client_id: "id", + client_secret: "secret", + refresh_token: "revoked-refresh-token", + token_uri: "https://oauth2.googleapis.com/token", + }), + ); + + const mockHttp = HttpClient.make((request) => + Effect.sync(() => { + return HttpClientResponse.fromWeb( + request, + Response.json({ error: "invalid_grant" }, { status: 401 }), + ); + }), + ); + + const result = yield* probeAntigravityUsageLimits({ + profileDirectory: tempDir, + checkedAt, + }).pipe(Effect.provideService(HttpClient.HttpClient, mockHttp)); + + expect(result).toBeUndefined(); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), + ); +}); diff --git a/apps/server/src/provider/Layers/antigravityUsageLimits.ts b/apps/server/src/provider/Layers/antigravityUsageLimits.ts new file mode 100644 index 000000000000..34e884ae69c6 --- /dev/null +++ b/apps/server/src/provider/Layers/antigravityUsageLimits.ts @@ -0,0 +1,412 @@ +/** + * Google Antigravity subscription usage and quota limits. + * + * Maps Google Code Assist / Antigravity tier responses and turn-driven quota + * updates into T3 Code's unified ServerProviderUsageLimits contract. + * + * @module provider/Layers/antigravityUsageLimits + */ +import type { + ProviderUsageLimitsUpdate, + ServerProviderUsageLimits, + ServerProviderUsageWindow, +} from "@t3tools/contracts"; +import { causeErrorTag } from "@t3tools/shared/observability"; +import * as Cause from "effect/Cause"; +import * as DateTime from "effect/DateTime"; +import * as Effect from "effect/Effect"; +import * as FileSystem from "effect/FileSystem"; +import * as Option from "effect/Option"; +import * as Path from "effect/Path"; +import * as Schema from "effect/Schema"; +import { + HttpClient, + HttpClientError, + HttpClientRequest, + HttpClientResponse, +} from "effect/unstable/http"; + +import { clampPercent, makeUsageLimits } from "../providerUsageLimits.ts"; + +export interface QuotaSummaryBucket { + readonly bucketId?: string | undefined; + readonly displayName?: string | undefined; + readonly description?: string | undefined; + readonly window?: string | undefined; + readonly remainingFraction?: number | undefined; + readonly remainingAmount?: number | undefined; + readonly disabled?: boolean | undefined; + readonly resetTime?: string | undefined; +} + +export interface QuotaSummaryGroup { + readonly displayName?: string | undefined; + readonly description?: string | undefined; + readonly buckets?: ReadonlyArray | undefined; +} + +export interface RetrieveUserQuotaSummaryResponse { + readonly groups?: ReadonlyArray | undefined; + readonly description?: string | undefined; +} + +export interface AntigravityQuotaWindow { + readonly id: string; + readonly label: string; + readonly kind: ServerProviderUsageWindow["kind"]; + readonly usedPercent: number; + readonly resetsAt?: string | null | undefined; + readonly windowDurationMins?: number | null | undefined; +} + +export interface AntigravityTierInfo { + readonly id?: string | null | undefined; + readonly name?: string | null | undefined; + readonly description?: string | null | undefined; +} + +export interface AntigravityQuotaSnapshot { + readonly quotaSummary?: RetrieveUserQuotaSummaryResponse | null | undefined; + readonly tier?: AntigravityTierInfo | null | undefined; + readonly allowedTiers?: ReadonlyArray | null | undefined; + readonly windows?: ReadonlyArray | null | undefined; + readonly monthlyPromptCredits?: number | null | undefined; + readonly availablePromptCredits?: number | null | undefined; + readonly monthlyFlowCredits?: number | null | undefined; + readonly availableFlowCredits?: number | null | undefined; +} + +const MONTH_MINS = 30 * 24 * 60; +const WEEK_MINS = 7 * 24 * 60; +const FIVE_HOUR_MINS = 5 * 60; + +const TokenFileSchema = Schema.Struct({ + client_id: Schema.String, + client_secret: Schema.String, + refresh_token: Schema.String, + token_uri: Schema.String, +}); + +const decodeTokenFile = Schema.decodeUnknownEffect(Schema.fromJsonString(TokenFileSchema)); + +function isSafeGoogleOAuthUri(uri: string): boolean { + try { + const parsed = new URL(uri); + return ( + parsed.protocol === "https:" && + (parsed.hostname === "oauth2.googleapis.com" || + parsed.hostname === "accounts.google.com" || + parsed.hostname.endsWith(".googleapis.com")) + ); + } catch { + return false; + } +} + +function isoFromString(value: string | null | undefined): string | undefined { + if (!value || typeof value !== "string") return undefined; + const dt = DateTime.make(value); + return Option.isSome(dt) && DateTime.toEpochMillis(dt.value) > 0 + ? DateTime.formatIso(dt.value) + : undefined; +} + +export function quotaSummaryToWindows( + response: RetrieveUserQuotaSummaryResponse | null | undefined, +): ReadonlyArray { + if (!response || typeof response !== "object") return []; + const windows: ServerProviderUsageWindow[] = []; + const rawGroups = Array.isArray(response.groups) ? response.groups : []; + const groups = [...rawGroups].sort((a, b) => { + const aGemini = /gemini/i.test(a?.displayName ?? "") ? 0 : 1; + const bGemini = /gemini/i.test(b?.displayName ?? "") ? 0 : 1; + return aGemini - bGemini; + }); + + for (const group of groups) { + if (!group || typeof group !== "object") continue; + const isGemini = /gemini/i.test(group.displayName ?? ""); + const is3P = /claude|gpt/i.test(group.displayName ?? ""); + const prefix = isGemini ? "Gemini" : is3P ? "Claude & GPT" : (group.displayName ?? "Model"); + + const rawBuckets = Array.isArray(group.buckets) ? group.buckets : []; + const buckets = [...rawBuckets].sort((a, b) => { + const aIs5h = a?.window === "5h" || /five hour|5h/i.test(a?.displayName ?? ""); + const bIs5h = b?.window === "5h" || /five hour|5h/i.test(b?.displayName ?? ""); + return (aIs5h ? 0 : 1) - (bIs5h ? 0 : 1); + }); + + for (const bucket of buckets) { + if (!bucket || typeof bucket !== "object" || bucket.disabled) continue; + const is5h = bucket.window === "5h" || /five hour|5h/i.test(bucket.displayName ?? ""); + const isWeekly = bucket.window === "weekly" || /weekly/i.test(bucket.displayName ?? ""); + const isMonthly = bucket.window === "monthly" || /monthly/i.test(bucket.displayName ?? ""); + + const windowSuffix = is5h + ? "5-hour" + : isWeekly + ? "Weekly" + : isMonthly + ? "Monthly" + : bucket.window || bucket.displayName || "Window"; + const label = `${prefix} (${windowSuffix})`; + const kind = is5h ? "session" : isWeekly ? "weekly" : isMonthly ? "monthly" : "session"; + const windowDurationMins = is5h + ? FIVE_HOUR_MINS + : isWeekly + ? WEEK_MINS + : isMonthly + ? MONTH_MINS + : undefined; + + const remainingFraction = + typeof bucket.remainingFraction === "number" && Number.isFinite(bucket.remainingFraction) + ? bucket.remainingFraction + : 1; + const usedPercent = clampPercent(Math.round((1 - remainingFraction) * 100)); + const resetsAt = isoFromString(bucket.resetTime); + + windows.push({ + id: + bucket.bucketId || + `${prefix.toLowerCase().replace(/[^a-z0-9]+/g, "_")}_${windowSuffix.toLowerCase()}`, + label, + kind, + usedPercent, + ...(windowDurationMins ? { windowDurationMins } : {}), + ...(resetsAt ? { resetsAt } : {}), + }); + } + } + return windows; +} + +export function antigravityRateLimitsToWindows( + snapshot: AntigravityQuotaSnapshot, +): ReadonlyArray { + if (snapshot.quotaSummary) { + const windows = quotaSummaryToWindows(snapshot.quotaSummary); + if (windows.length > 0) return windows; + } + + const windows: ServerProviderUsageWindow[] = []; + + if (snapshot.windows && snapshot.windows.length > 0) { + for (const w of snapshot.windows) { + windows.push({ + id: w.id, + kind: w.kind, + label: w.label, + usedPercent: clampPercent(w.usedPercent), + ...(typeof w.windowDurationMins === "number" + ? { windowDurationMins: w.windowDurationMins } + : {}), + ...(w.resetsAt ? { resetsAt: w.resetsAt } : {}), + }); + } + return windows; + } + + if ( + typeof snapshot.monthlyPromptCredits === "number" && + typeof snapshot.availablePromptCredits === "number" && + snapshot.monthlyPromptCredits > 0 + ) { + const used = snapshot.monthlyPromptCredits - snapshot.availablePromptCredits; + const usedPercent = (used / snapshot.monthlyPromptCredits) * 100; + windows.push({ + id: "prompt_credits", + kind: "monthly", + label: "Prompt Credits", + usedPercent: clampPercent(usedPercent), + windowDurationMins: MONTH_MINS, + }); + } + + if ( + typeof snapshot.monthlyFlowCredits === "number" && + typeof snapshot.availableFlowCredits === "number" && + snapshot.monthlyFlowCredits > 0 + ) { + const used = snapshot.monthlyFlowCredits - snapshot.availableFlowCredits; + const usedPercent = (used / snapshot.monthlyFlowCredits) * 100; + windows.push({ + id: "flow_credits", + kind: "monthly", + label: "Flow Credits", + usedPercent: clampPercent(usedPercent), + windowDurationMins: MONTH_MINS, + }); + } + + if ( + windows.length === 0 && + (snapshot.tier || (snapshot.allowedTiers && snapshot.allowedTiers.length > 0)) + ) { + const tier = snapshot.tier ?? snapshot.allowedTiers?.[0]; + const label = tier?.name ? `${tier.name}` : "Antigravity Plan"; + windows.push({ + id: "plan_allowance", + kind: "session", + label, + usedPercent: 0, + }); + } + + return windows; +} + +export function antigravityRateLimitsToLimits(input: { + readonly snapshot: AntigravityQuotaSnapshot; + readonly checkedAt: string; +}): ServerProviderUsageLimits { + const windows = antigravityRateLimitsToWindows(input.snapshot); + return makeUsageLimits({ + checkedAt: input.checkedAt, + windows, + }); +} + +export function antigravityRateLimitsToUpdate( + snapshot: AntigravityQuotaSnapshot, +): ProviderUsageLimitsUpdate | undefined { + const windows = antigravityRateLimitsToWindows(snapshot); + return windows.length > 0 ? { windows } : undefined; +} + +/** + * Probes Google Antigravity quota / tier info using the profile's stored OAuth credentials. + */ +export const probeAntigravityUsageLimits = (input: { + readonly profileDirectory: string; + readonly checkedAt: string; +}) => + Effect.gen(function* () { + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const client = yield* HttpClient.HttpClient; + + const tokenPath = path.join(input.profileDirectory, "antigravity-acp", "acp_token.json"); + if (!(yield* fs.exists(tokenPath))) { + return undefined; + } + + const token = yield* fs + .readFileString(tokenPath) + .pipe(Effect.flatMap(decodeTokenFile), Effect.option); + if (Option.isNone(token) || !isSafeGoogleOAuthUri(token.value.token_uri)) { + return undefined; + } + + // Exchange refresh token for access token + const tokenRequest = HttpClientRequest.post(token.value.token_uri).pipe( + HttpClientRequest.bodyUrlParams({ + client_id: token.value.client_id, + client_secret: token.value.client_secret, + refresh_token: token.value.refresh_token, + grant_type: "refresh_token", + }), + ); + + const tokenResponse = yield* client.execute(tokenRequest).pipe( + Effect.flatMap(HttpClientResponse.filterStatusOk), + Effect.flatMap((res) => res.json), + Effect.timeout("10 seconds"), + Effect.option, + ); + if (Option.isNone(tokenResponse)) return undefined; + + const accessToken = (tokenResponse.value as { access_token?: string })?.access_token; + if (!accessToken || typeof accessToken !== "string") return undefined; + + // 1. Probe granular multi-group quota buckets (Gemini & Claude/GPT models with 5h and weekly limits) + const quotaEndpoints = [ + "https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary", + "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary", + ]; + + for (const endpoint of quotaEndpoints) { + const quotaRequest = HttpClientRequest.post(endpoint).pipe( + HttpClientRequest.setHeader("Authorization", `Bearer ${accessToken}`), + HttpClientRequest.setHeader("Content-Type", "application/json"), + HttpClientRequest.setHeader("User-Agent", "antigravity/1.1.28"), + HttpClientRequest.bodyJsonUnsafe({ project: "default-cli-project" }), + ); + + const quotaResponse = yield* client.execute(quotaRequest).pipe( + Effect.flatMap(HttpClientResponse.filterStatusOk), + Effect.flatMap((res) => res.json), + Effect.timeout("5 seconds"), + Effect.option, + ); + + if (Option.isSome(quotaResponse)) { + const windows = quotaSummaryToWindows( + quotaResponse.value as RetrieveUserQuotaSummaryResponse, + ); + if (windows.length > 0) { + return makeUsageLimits({ + checkedAt: input.checkedAt, + windows, + }); + } + } + } + + // 2. Fall back to loadCodeAssist endpoint for tier and credit state + const apiRequest = HttpClientRequest.post( + "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist", + ).pipe( + HttpClientRequest.setHeader("Authorization", `Bearer ${accessToken}`), + HttpClientRequest.setHeader("Content-Type", "application/json"), + HttpClientRequest.setHeader("User-Agent", "antigravity/1.1.28"), + HttpClientRequest.bodyJsonUnsafe({ + metadata: { + ideType: "ANTIGRAVITY", + ideVersion: "1.0.0", + }, + }), + ); + + const apiResponse = yield* client.execute(apiRequest).pipe( + Effect.flatMap(HttpClientResponse.filterStatusOk), + Effect.flatMap((res) => res.json), + Effect.timeout("10 seconds"), + Effect.option, + ); + + if (Option.isNone(apiResponse)) return undefined; + + const response = apiResponse.value as { + allowedTiers?: ReadonlyArray; + quotaManagerState?: { + monthlyPromptCredits?: number; + availablePromptCredits?: number; + monthlyFlowCredits?: number; + availableFlowCredits?: number; + }; + }; + + const snapshot: AntigravityQuotaSnapshot = { + allowedTiers: response.allowedTiers, + monthlyPromptCredits: response.quotaManagerState?.monthlyPromptCredits, + availablePromptCredits: response.quotaManagerState?.availablePromptCredits, + monthlyFlowCredits: response.quotaManagerState?.monthlyFlowCredits, + availableFlowCredits: response.quotaManagerState?.availableFlowCredits, + }; + + return antigravityRateLimitsToLimits({ + snapshot, + checkedAt: input.checkedAt, + }); + }).pipe( + Effect.catchCause((cause) => { + const failure = Cause.findErrorOption(cause); + const err: unknown = Option.isSome(failure) ? failure.value : undefined; + const request = HttpClientError.isHttpClientError(err) ? err.request : undefined; + return Effect.logDebug("Antigravity usage-limit probe failed", { + errorTag: causeErrorTag(cause), + ...(request ? { method: request.method, url: request.url } : {}), + }).pipe(Effect.as(undefined)); + }), + ); diff --git a/apps/server/src/provider/makeManagedServerProvider.ts b/apps/server/src/provider/makeManagedServerProvider.ts index c180412ec42d..fee531fa8d2e 100644 --- a/apps/server/src/provider/makeManagedServerProvider.ts +++ b/apps/server/src/provider/makeManagedServerProvider.ts @@ -85,8 +85,17 @@ export const makeManagedServerProvider = Effect.fn("makeManagedServerProvider")( return [null, state] as const; } // Enrichment derives from the snapshot it was handed; a runtime usage - // update that landed since must not be reverted by it. - const merged = withUsageLimits(nextSnapshot, state.snapshot.usageLimits); + // update that landed since must not be reverted by it unless the enriched + // snapshot carries a fresher usageLimits timestamp or cleared it on sign-out. + const usageLimits = + nextSnapshot.usageLimits && + (!state.snapshot.usageLimits || + nextSnapshot.usageLimits.checkedAt >= state.snapshot.usageLimits.checkedAt) + ? nextSnapshot.usageLimits + : nextSnapshot.auth.status === "unauthenticated" + ? undefined + : state.snapshot.usageLimits; + const merged = withUsageLimits(nextSnapshot, usageLimits); if (Equal.equals(state.snapshot, merged)) { return [null, state] as const; } diff --git a/apps/server/src/provider/providerUsageLimits.test.ts b/apps/server/src/provider/providerUsageLimits.test.ts index 6e288ddd3a36..b554b8fdf752 100644 --- a/apps/server/src/provider/providerUsageLimits.test.ts +++ b/apps/server/src/provider/providerUsageLimits.test.ts @@ -1,6 +1,10 @@ import { describe, expect, it } from "vite-plus/test"; -import { applyUsageLimitsUpdate, resolveUsageLimitsAfterProbe } from "./providerUsageLimits.ts"; +import { + applyUsageLimitsUpdate, + makeUsageLimits, + resolveUsageLimitsAfterProbe, +} from "./providerUsageLimits.ts"; const checkedAt = "2026-09-03T12:00:00.000Z"; const session = { @@ -87,3 +91,24 @@ describe("resolveUsageLimitsAfterProbe", () => { expect(resolveUsageLimitsAfterProbe({ published: undefined, probed: failed })).toBe(failed); }); }); + +describe("makeUsageLimits", () => { + it("preserves group appearance order while sorting session before weekly within each group", () => { + const limits = makeUsageLimits({ + checkedAt, + windows: [ + { id: "gemini_weekly", kind: "weekly", label: "Gemini (Weekly)", usedPercent: 10 }, + { id: "gemini_5h", kind: "session", label: "Gemini (5-hour)", usedPercent: 20 }, + { id: "3p_weekly", kind: "weekly", label: "Claude & GPT (Weekly)", usedPercent: 30 }, + { id: "3p_5h", kind: "session", label: "Claude & GPT (5-hour)", usedPercent: 40 }, + ], + }); + + expect(limits.windows.map((w) => w.label)).toEqual([ + "Gemini (5-hour)", + "Gemini (Weekly)", + "Claude & GPT (5-hour)", + "Claude & GPT (Weekly)", + ]); + }); +}); diff --git a/apps/server/src/provider/providerUsageLimits.ts b/apps/server/src/provider/providerUsageLimits.ts index ea8d0d1d029f..e48ac69a453d 100644 --- a/apps/server/src/provider/providerUsageLimits.ts +++ b/apps/server/src/provider/providerUsageLimits.ts @@ -3,6 +3,7 @@ import type { ServerProviderUsageLimits, ServerProviderUsageWindow, } from "@t3tools/contracts"; +import { windowGroup } from "@t3tools/shared/usageLimits"; const WINDOW_KIND_ORDER: Record = { session: 0, @@ -18,11 +19,26 @@ export function clampPercent(value: number): number { function sortWindows( windows: Iterable, ): ReadonlyArray { - return [...windows].toSorted( - (left, right) => + const windowList = [...windows]; + const groupOrder = new Map(); + for (const window of windowList) { + const group = windowGroup(window); + if (!groupOrder.has(group)) { + groupOrder.set(group, groupOrder.size); + } + } + + return windowList.toSorted((left, right) => { + const leftGroupOrder = groupOrder.get(windowGroup(left)) ?? 0; + const rightGroupOrder = groupOrder.get(windowGroup(right)) ?? 0; + if (leftGroupOrder !== rightGroupOrder) { + return leftGroupOrder - rightGroupOrder; + } + return ( WINDOW_KIND_ORDER[left.kind] - WINDOW_KIND_ORDER[right.kind] || - left.id.localeCompare(right.id), - ); + left.id.localeCompare(right.id) + ); + }); } export function makeUsageLimits(input: { diff --git a/apps/web/src/components/usage/UsageLimits.tsx b/apps/web/src/components/usage/UsageLimits.tsx index e50547b66008..4fb8af6f908b 100644 --- a/apps/web/src/components/usage/UsageLimits.tsx +++ b/apps/web/src/components/usage/UsageLimits.tsx @@ -167,8 +167,8 @@ export function LimitWindows({
{windows.map((window) => { @@ -177,7 +177,9 @@ export function LimitWindows({ return ( - {window.label} + + {window.label} + {remainingPercent(window)}% left diff --git a/packages/shared/src/usageLimits.test.ts b/packages/shared/src/usageLimits.test.ts index b814e66da459..a4f38a78a530 100644 --- a/packages/shared/src/usageLimits.test.ts +++ b/packages/shared/src/usageLimits.test.ts @@ -782,6 +782,36 @@ describe("pools", () => { expect(session?.members.map((member) => member.account.key)).toEqual(["hub:a", "hub:b"]); expect(pools[0]?.accounts.map((account) => account.key)).toEqual(["hub:a", "hub:b"]); }); + + it("preserves group appearance order while sorting session before weekly within each group", () => { + const multiGroupAccount: LimitAccount = { + key: "native:antigravity", + driver: ProviderDriverKind.make("antigravity"), + displayName: "Antigravity", + email: undefined, + plan: undefined, + accentColor: undefined, + environments: [], + sourceLabel: null, + redeem: null, + limits: { + checkedAt, + windows: [ + { id: "gemini-5h", kind: "session", label: "Gemini (5-hour)", usedPercent: 40 }, + { id: "gemini-weekly", kind: "weekly", label: "Gemini (Weekly)", usedPercent: 20 }, + { id: "3p-5h", kind: "session", label: "Claude & GPT (5-hour)", usedPercent: 0 }, + { id: "3p-weekly", kind: "weekly", label: "Claude & GPT (Weekly)", usedPercent: 0 }, + ], + }, + }; + const pools = collectLimitPools([multiGroupAccount], now); + expect(pools[0]?.windows.map((w) => w.label)).toEqual([ + "Gemini (5-hour)", + "Gemini (Weekly)", + "Claude & GPT (5-hour)", + "Claude & GPT (Weekly)", + ]); + }); }); describe("pooled account columns", () => { diff --git a/packages/shared/src/usageLimits.ts b/packages/shared/src/usageLimits.ts index 5c32cc0343b7..09c4a7735093 100644 --- a/packages/shared/src/usageLimits.ts +++ b/packages/shared/src/usageLimits.ts @@ -433,6 +433,12 @@ function accountSortName(account: LimitAccount): string { return (account.displayName ?? account.email ?? account.key).toLowerCase(); } +export function windowGroup(window: { readonly label?: string | null }): string { + if (!window.label || typeof window.label !== "string") return ""; + const match = window.label.match(/^([^(]+?)\s*\(/); + return match ? match[1]!.trim().toLowerCase() : ""; +} + function poolWindows(accounts: readonly LimitAccount[], now: number): readonly LimitPoolWindow[] { const byKey = new Map(); for (const account of accounts) { @@ -485,7 +491,22 @@ function poolWindows(accounts: readonly LimitAccount[], now: number): readonly L resets, }; }); - return pools.sort((left, right) => WINDOW_KIND_ORDER[left.kind] - WINDOW_KIND_ORDER[right.kind]); + const groupOrder = new Map(); + for (const pool of pools) { + const group = windowGroup(pool); + if (!groupOrder.has(group)) { + groupOrder.set(group, groupOrder.size); + } + } + + return pools.sort((left, right) => { + const leftGroupOrder = groupOrder.get(windowGroup(left)) ?? 0; + const rightGroupOrder = groupOrder.get(windowGroup(right)) ?? 0; + if (leftGroupOrder !== rightGroupOrder) { + return leftGroupOrder - rightGroupOrder; + } + return WINDOW_KIND_ORDER[left.kind] - WINDOW_KIND_ORDER[right.kind]; + }); } /** The one-line status under a provider heading when there are no bars to draw. */