Просмотр исходного кода

fix(account): coalesce concurrent console token refreshes (#20503)

Kit Langton 2 недель назад
Родитель
Сommit
c619caefdd

+ 26 - 3
packages/opencode/src/account/index.ts

@@ -1,4 +1,4 @@
-import { Clock, Duration, Effect, Layer, Option, Schema, SchemaGetter, ServiceMap } from "effect"
+import { Cache, Clock, Duration, Effect, Layer, Option, Schema, SchemaGetter, ServiceMap } from "effect"
 import { FetchHttpClient, HttpClient, HttpClientRequest, HttpClientResponse } from "effect/unstable/http"
 
 import { makeRuntime } from "@/effect/run-service"
@@ -175,9 +175,8 @@ export namespace Account {
           mapAccountServiceError("HTTP request failed"),
         )
 
-      const resolveToken = Effect.fnUntraced(function* (row: AccountRow) {
+      const refreshToken = Effect.fnUntraced(function* (row: AccountRow) {
         const now = yield* Clock.currentTimeMillis
-        if (row.token_expiry && row.token_expiry > now) return row.access_token
 
         const response = yield* executeEffectOk(
           HttpClientRequest.post(`${row.url}/auth/device/token`).pipe(
@@ -208,6 +207,30 @@ export namespace Account {
         return parsed.access_token
       })
 
+      const refreshTokenCache = yield* Cache.make<AccountID, AccessToken, AccountError>({
+        capacity: Number.POSITIVE_INFINITY,
+        timeToLive: Duration.zero,
+        lookup: Effect.fnUntraced(function* (accountID) {
+          const maybeAccount = yield* repo.getRow(accountID)
+          if (Option.isNone(maybeAccount)) {
+            return yield* Effect.fail(new AccountServiceError({ message: "Account not found during token refresh" }))
+          }
+
+          const account = maybeAccount.value
+          const now = yield* Clock.currentTimeMillis
+          if (account.token_expiry && account.token_expiry > now) return account.access_token
+
+          return yield* refreshToken(account)
+        }),
+      })
+
+      const resolveToken = Effect.fnUntraced(function* (row: AccountRow) {
+        const now = yield* Clock.currentTimeMillis
+        if (row.token_expiry && row.token_expiry > now) return row.access_token
+
+        return yield* Cache.get(refreshTokenCache, row.id)
+      })
+
       const resolveAccess = Effect.fnUntraced(function* (accountID: AccountID) {
         const maybeAccount = yield* repo.getRow(accountID)
         if (Option.isNone(maybeAccount)) return Option.none()

+ 64 - 0
packages/opencode/test/account/service.test.ts

@@ -148,6 +148,70 @@ it.live("token refresh persists the new token", () =>
   }),
 )
 
+it.live("concurrent config and token requests coalesce token refresh", () =>
+  Effect.gen(function* () {
+    const id = AccountID.make("user-1")
+
+    yield* AccountRepo.use((r) =>
+      r.persistAccount({
+        id,
+        email: "[email protected]",
+        url: "https://one.example.com",
+        accessToken: AccessToken.make("at_old"),
+        refreshToken: RefreshToken.make("rt_old"),
+        expiry: Date.now() - 1_000,
+        orgID: Option.some(OrgID.make("org-9")),
+      }),
+    )
+
+    let refreshCalls = 0
+    const client = HttpClient.make((req) =>
+      Effect.promise(async () => {
+        if (req.url === "https://one.example.com/auth/device/token") {
+          refreshCalls += 1
+
+          if (refreshCalls === 1) {
+            await new Promise((resolve) => setTimeout(resolve, 25))
+            return json(req, {
+              access_token: "at_new",
+              refresh_token: "rt_new",
+              expires_in: 60,
+            })
+          }
+
+          return json(
+            req,
+            {
+              error: "invalid_grant",
+              error_description: "refresh token already used",
+            },
+            400,
+          )
+        }
+
+        if (req.url === "https://one.example.com/api/config") {
+          return json(req, { config: { theme: "light", seats: 5 } })
+        }
+
+        return json(req, {}, 404)
+      }),
+    )
+
+    const [cfg, token] = yield* Account.Service.use((s) =>
+      Effect.all([s.config(id, OrgID.make("org-9")), s.token(id)], { concurrency: 2 }),
+    ).pipe(Effect.provide(live(client)))
+
+    expect(Option.getOrThrow(cfg)).toEqual({ theme: "light", seats: 5 })
+    expect(String(Option.getOrThrow(token))).toBe("at_new")
+    expect(refreshCalls).toBe(1)
+
+    const row = yield* AccountRepo.use((r) => r.getRow(id))
+    const value = Option.getOrThrow(row)
+    expect(value.access_token).toBe(AccessToken.make("at_new"))
+    expect(value.refresh_token).toBe(RefreshToken.make("rt_new"))
+  }),
+)
+
 it.live("config sends the selected org header", () =>
   Effect.gen(function* () {
     const id = AccountID.make("user-1")