diff --git a/apps/server/src/provider/CodexChatGptAuth.test.ts b/apps/server/src/provider/CodexChatGptAuth.test.ts index 3e521103d56cf..1ceb3330e1abb 100644 --- a/apps/server/src/provider/CodexChatGptAuth.test.ts +++ b/apps/server/src/provider/CodexChatGptAuth.test.ts @@ -1522,6 +1522,25 @@ it.effect("primary completes OAuth and destination imports and owns the refresh ), ); +it.effect("profile import stops credential users before taking the auth access lock", () => + provision( + Effect.gen(function* () { + const h = yield* makeHarness; + yield* h.signIn; + yield* h.phase("succeeded"); + const profile = yield* h.auth.exportProfile; + const destination = yield* h.destination; + yield* destination.auth.controller.importProfile!(profile, Effect.void); + + const imported = yield* destination.auth.controller.importProfile!( + profile, + destination.auth.access.pipe(Effect.asVoid), + ); + assert.strictEqual(imported.phase, "succeeded"); + }), + ), +); + it.effect("transferred profiles reject mismatched identities without overwriting credentials", () => provision( Effect.gen(function* () { @@ -1539,11 +1558,15 @@ it.effect("transferred profiles reject mismatched identities without overwriting { ...profile.credentials, scopes: ["openid"] }, { ...profile.credentials, expiresAt: 0 }, ]) { + let stopped = false; const result = yield* destination.auth.controller.importProfile!( { ...profile, credentials }, - Effect.void, + Effect.sync(() => { + stopped = true; + }), ).pipe(Effect.result); assert.strictEqual(result._tag, "Failure"); + assert.isFalse(stopped); assert.deepEqual(Option.getOrThrow(yield* destination.auth.read), original); } }), diff --git a/apps/server/src/provider/CodexChatGptAuth.ts b/apps/server/src/provider/CodexChatGptAuth.ts index 1051f72814390..e0f2df6593ba5 100644 --- a/apps/server/src/provider/CodexChatGptAuth.ts +++ b/apps/server/src/provider/CodexChatGptAuth.ts @@ -992,13 +992,15 @@ export const makeCodexChatGptAuth = Effect.fn("makeCodexChatGptAuth")(function* yield* saveRegistration(registration); yield* save(profile.credentials, true); }); - return lock.withPermit( - track( - "transfer", - { flow: "primary_handoff" }, - validate.pipe( - Effect.andThen(controller.adoptCredentials!(withSessionLock(saveImported), stopSessions)), - ), + // Claim credential adoption before validation, but release the auth lock before + // session teardown so a closing managed Codex session can drain token renewal. + return track( + "transfer", + { flow: "primary_handoff" }, + controller.adoptCredentials!( + lock.withPermit(withSessionLock(saveImported)), + stopSessions, + lock.withPermit(validate), ), ); }; diff --git a/packages/provider-core/src/server/auth.ts b/packages/provider-core/src/server/auth.ts index 25f30b8fd0c43..df39ebeeeca52 100644 --- a/packages/provider-core/src/server/auth.ts +++ b/packages/provider-core/src/server/auth.ts @@ -21,6 +21,7 @@ export interface ProviderAuthController { readonly adoptCredentials?: ( update: Effect.Effect, stopSessions: Effect.Effect, + prepare?: Effect.Effect, ) => Effect.Effect; readonly isChangingCredentials?: Effect.Effect; readonly invalidate?: Effect.Effect; diff --git a/packages/provider-core/src/server/providerAuthFlow.test.ts b/packages/provider-core/src/server/providerAuthFlow.test.ts index 968eb04f78b4d..4f9ae2ae7b37f 100644 --- a/packages/provider-core/src/server/providerAuthFlow.test.ts +++ b/packages/provider-core/src/server/providerAuthFlow.test.ts @@ -516,3 +516,82 @@ it.effect("profile import stops owned sessions and gates access until credential ); }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), ); + +it.effect("profile import claims credential changes before pre-stop validation", () => + Effect.gen(function* () { + const validationStarted = yield* Deferred.make(); + const releaseValidation = yield* Deferred.make(); + let stopped = false; + let loggedOut = false; + const controller = yield* ProviderAuthFlow.make({ + instanceId, + credentialBinding: { owner: "t3", key: "handoff-validation" }, + methods: Effect.succeed([method]), + authenticate: () => Effect.void, + logout: Effect.sync(() => { + loggedOut = true; + }), + }); + const imported = yield* controller.adoptCredentials!( + Effect.void, + Effect.sync(() => { + stopped = true; + }), + Deferred.succeed(validationStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseValidation)), + Effect.andThen( + Effect.fail( + new ProviderSetupError({ + instanceId, + operation: "import", + detail: "invalid transferred credentials", + }), + ), + ), + ), + ).pipe(Effect.result, Effect.forkChild); + + yield* Deferred.await(validationStarted); + assert.isTrue(yield* controller.isChangingCredentials!); + const concurrentLogout = yield* controller.logout(Effect.void).pipe(Effect.result); + assert.strictEqual(concurrentLogout._tag, "Failure"); + assert.isFalse(loggedOut); + assert.isFalse(stopped); + + yield* Deferred.succeed(releaseValidation, undefined); + assert.strictEqual((yield* Fiber.join(imported))._tag, "Failure"); + assert.isFalse(stopped); + + yield* controller.logout(Effect.void); + assert.isTrue(loggedOut); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), +); + +it.effect("cancelling pre-stop validation releases the credential-change claim", () => + Effect.gen(function* () { + const validationStarted = yield* Deferred.make(); + let stopped = false; + const controller = yield* ProviderAuthFlow.make({ + instanceId, + credentialBinding: { owner: "t3", key: "handoff-validation-cancel" }, + methods: Effect.succeed([method]), + authenticate: () => Effect.void, + logout: Effect.void, + }); + const imported = yield* controller.adoptCredentials!( + Effect.void, + Effect.sync(() => { + stopped = true; + }), + Deferred.succeed(validationStarted, undefined).pipe(Effect.andThen(Effect.never)), + ).pipe(Effect.forkChild); + + yield* Deferred.await(validationStarted); + assert.isTrue(yield* controller.isChangingCredentials!); + yield* Fiber.interrupt(imported); + + assert.isFalse(stopped); + assert.isFalse(yield* controller.isChangingCredentials!); + assert.strictEqual(yield* controller.withAccess!(Effect.succeed("available")), "available"); + }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)), +); diff --git a/packages/provider-core/src/server/providerAuthFlow.ts b/packages/provider-core/src/server/providerAuthFlow.ts index 6dbd06d976f03..69c3bbb9f837f 100644 --- a/packages/provider-core/src/server/providerAuthFlow.ts +++ b/packages/provider-core/src/server/providerAuthFlow.ts @@ -143,39 +143,47 @@ export const make = Effect.fn("ProviderAuthFlow.make")(function* (options: { adoptCredentials: Effect.fn("providerAuthFlow.adoptCredentials")(function* ( update, stopSessions, + prepare = Effect.void, ) { - yield* lock.withPermit( - Effect.gen(function* () { - if (operation !== "idle") - return yield* new ProviderSetupError({ - instanceId: options.instanceId, - operation: "import", - detail: "Finish or cancel the existing sign-in first.", - }); - operation = "stopping"; - }), - ); - const result = yield* Effect.gen(function* () { - yield* stopOwnedSessions; - yield* stopSessions; - yield* update; - yield* refreshMethods; - }).pipe(Effect.exit); - const state: ProviderAuthState = { - ...empty, - methods: snapshot.value.state.methods ?? [], - phase: Exit.isSuccess(result) ? "succeeded" : "failed", - message: Exit.isSuccess(result) ? null : failureMessage(result.cause), - }; - yield* lock.withPermit( + return yield* Effect.uninterruptibleMask((restore) => Effect.gen(function* () { - yield* SubscriptionRef.set(snapshot, { owner: null, state }); - operation = "idle"; + yield* lock.withPermit( + Effect.gen(function* () { + if (operation !== "idle") + return yield* new ProviderSetupError({ + instanceId: options.instanceId, + operation: "import", + detail: "Finish or cancel the existing sign-in first.", + }); + operation = "stopping"; + }), + ); + const result = yield* restore( + Effect.gen(function* () { + yield* prepare; + yield* stopOwnedSessions; + yield* stopSessions; + yield* update; + yield* refreshMethods; + }), + ).pipe(Effect.exit); + const state: ProviderAuthState = { + ...empty, + methods: snapshot.value.state.methods ?? [], + phase: Exit.isSuccess(result) ? "succeeded" : "failed", + message: Exit.isSuccess(result) ? null : failureMessage(result.cause), + }; + yield* lock.withPermit( + Effect.gen(function* () { + yield* SubscriptionRef.set(snapshot, { owner: null, state }); + operation = "idle"; + }), + ); + if (Exit.isFailure(result)) return yield* Effect.failCause(result.cause); + return state; }), ); - if (Exit.isFailure(result)) return yield* Effect.failCause(result.cause); - return state; - }, Effect.uninterruptible), + }), refreshMethods, invalidate: lock.withPermit( Effect.gen(function* () {