diff --git a/apps/server/src/provider/antigravityUsageLimits.test.ts b/apps/server/src/provider/antigravityUsageLimits.test.ts index 77d6be133143c27e8df53da79317ae46ce96aef5..00c2dd21f05721743923e79ccd6bab921f8e81e7 100644 --- a/apps/server/src/provider/antigravityUsageLimits.test.ts +++ b/apps/server/src/provider/antigravityUsageLimits.test.ts @@ -53,6 +53,9 @@ const fixture = Effect.fn("AntigravityUsageTest.fixture")(function* ( models?: unknown; summaryStatus?: number; tokenStatus?: number; + refreshToken?: string; + /** Google account id from userinfo; omit to make userinfo fail. */ + userId?: string; missingToken?: boolean; authMethod?: "oauth-personal" | "gemini-api-key"; } = {}, @@ -72,7 +75,7 @@ const fixture = Effect.fn("AntigravityUsageTest.fixture")(function* ( project_id: "project-id", }), ); - if (!options.missingToken) yield* writeToken("refresh-a"); + if (!options.missingToken) yield* writeToken(options.refreshToken ?? "refresh-a"); const requests: Array<{ url: string; body: string }> = []; const http = HttpClient.make((request, _url, _signal, fiber) => Effect.sync(() => { @@ -92,6 +95,14 @@ const fixture = Effect.fn("AntigravityUsageTest.fixture")(function* ( ); } expect(request.headers.authorization).toBe("Bearer access-token"); + if (request.url.endsWith("/userinfo")) { + return HttpClientResponse.fromWeb( + request, + options.userId + ? Response.json({ id: options.userId, email: "someone@example.com" }) + : Response.json({}, { status: 403 }), + ); + } expect(JSON.parse(body)).toEqual({ project: "project-id" }); return HttpClientResponse.fromWeb( request, @@ -142,6 +153,7 @@ describe("Antigravity subscription limits", () => { const first = yield* probe; yield* writeToken("refresh-b"); const second = yield* probe; + // Without a readable Google account, the credential itself names the account. expect(second.credentialFingerprint).not.toBe(first.credentialFingerprint); expect(requests.filter((request) => request.url.endsWith("/token"))).toHaveLength(2); yield* (yield* FileSystem.FileSystem).remove(tokenPath); @@ -155,6 +167,20 @@ describe("Antigravity subscription limits", () => { }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), ); + it.effect("names one Google account the same way from every sign-in", () => + Effect.gen(function* () { + const read = (refreshToken: string, userId: string) => + fixture({ refreshToken, userId }).pipe(Effect.flatMap((usage) => usage.probe)); + const machineA = yield* read("refresh-a", "google-user-1"); + const machineB = yield* read("refresh-b", "google-user-1"); + const otherAccount = yield* read("refresh-c", "google-user-2"); + expect(machineA.credentialFingerprint).toMatch(/^[0-9a-f]{64}$/); + expect(machineB.credentialFingerprint).toBe(machineA.credentialFingerprint); + expect(otherAccount.credentialFingerprint).not.toBe(machineA.credentialFingerprint); + expect(JSON.stringify(machineA)).not.toMatch(/google-user-1|someone@example.com/); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), + ); + it.effect("falls back to model quotas without inventing a window duration", () => Effect.gen(function* () { const { probe } = yield* fixture({ diff --git a/apps/server/src/provider/antigravityUsageLimits.ts b/apps/server/src/provider/antigravityUsageLimits.ts index f0d7f8ead52fa65eef21801d11f4c2d2dda03557..185a0c44c1447612f7e232ba984de7dfe3904339 100644 --- a/apps/server/src/provider/antigravityUsageLimits.ts +++ b/apps/server/src/provider/antigravityUsageLimits.ts @@ -30,6 +30,7 @@ const TokenResponse = Schema.Struct({ access_token: Schema.NonEmptyString, expires_in: Schema.Number.check(Schema.isGreaterThan(0)), }); +const UserInfo = Schema.Struct({ id: Schema.NonEmptyString }); const Quota = Schema.Struct({ remainingFraction: Schema.optional(Schema.Number), resetTime: Schema.optional(Schema.String), @@ -65,6 +66,7 @@ const Models = Schema.Struct({ const decodeTokenFile = Schema.decodeEffect(Schema.fromJsonString(TokenFile)); const decodeTokenResponse = Schema.decodeUnknownEffect(TokenResponse); +const decodeUserInfo = Schema.decodeUnknownEffect(UserInfo); const decodeSummary = Schema.decodeUnknownEffect(Summary); const decodeModels = Schema.decodeUnknownEffect(Models); @@ -168,9 +170,20 @@ export const makeAntigravityUsageProbe = Effect.fn("makeAntigravityUsageProbe")( const platform = yield* HostProcess.Platform; const architecture = yield* HostProcess.Architecture; const tokenPath = path.join(input.profileDirectory, "antigravity-acp", "acp_token.json"); - let cached: { fingerprint: string; accessToken: string; expiresAt: number } | undefined; + let cached: + | { fingerprint: string; account: string; accessToken: string; expiresAt: number } + | undefined; let generation = 0; + const sha256 = (text: string) => + crypto + .digest("SHA-256", new TextEncoder().encode(text)) + .pipe( + Effect.map((digest) => + Array.from(digest, (byte) => byte.toString(16).padStart(2, "0")).join(""), + ), + ); + const request = (request: HttpClientRequest.HttpClientRequest) => client.execute(request).pipe( Effect.provideService(FetchHttpClient.RequestInit, MANUAL_REDIRECT), @@ -193,17 +206,12 @@ export const makeAntigravityUsageProbe = Effect.fn("makeAntigravityUsageProbe")( }); } const token = yield* fs.readFileString(tokenPath).pipe(Effect.flatMap(decodeTokenFile)); - const fingerprint = Array.from( - yield* crypto.digest( - "SHA-256", - new TextEncoder().encode( - `${token.client_id}\0${token.refresh_token}\0${token.project_id ?? ""}`, - ), - ), - (byte) => byte.toString(16).padStart(2, "0"), - ).join(""); + const fingerprint = yield* sha256( + `${token.client_id}\0${token.refresh_token}\0${token.project_id ?? ""}`, + ); const now = yield* Clock.currentTimeMillis; let accessToken: string; + let account: string; if (cached?.fingerprint !== fingerprint || cached.expiresAt <= now + 60_000) { cached = undefined; const refreshed = yield* request( @@ -217,11 +225,28 @@ export const makeAntigravityUsageProbe = Effect.fn("makeAntigravityUsageProbe")( ), ).pipe(Effect.flatMap(decodeTokenResponse)); accessToken = refreshed.access_token; + // Each machine's sign-in has its own refresh token, so the Google + // account id is what lets clients merge one account across environments. + account = yield* request( + HttpClientRequest.get("https://www.googleapis.com/oauth2/v2/userinfo").pipe( + HttpClientRequest.setHeaders({ Authorization: `Bearer ${accessToken}` }), + ), + ).pipe( + Effect.flatMap(decodeUserInfo), + Effect.flatMap((info) => sha256(`google-account\0${info.id}`)), + Effect.orElseSucceed(() => fingerprint), + ); if (generation === startedGeneration) { - cached = { fingerprint, accessToken, expiresAt: now + refreshed.expires_in * 1000 }; + cached = { + fingerprint, + account, + accessToken, + expiresAt: now + refreshed.expires_in * 1000, + }; } } else { accessToken = cached.accessToken; + account = cached.account; } const fetchWindows = Effect.fn("AntigravityUsageProbe.fetchWindows")(function* ( operation: string, @@ -259,7 +284,7 @@ export const makeAntigravityUsageProbe = Effect.fn("makeAntigravityUsageProbe")( message: "Antigravity did not return usable quota information.", }); } - return { ...makeUsageLimits({ checkedAt, windows }), credentialFingerprint: fingerprint }; + return { ...makeUsageLimits({ checkedAt, windows }), credentialFingerprint: account }; }); const probe = read.pipe( Effect.timeout("30 seconds"),