Separate API key sign-in into its own screen

Sources/CodexReviewHost/AccountRuntimeTransitionCoordinator.swift

@@ -141,7 +141,7 @@ final class AccountRuntimeTransitionCoordinator {
             generation: UInt64,
             phase: RuntimeAuthPhase
         )
-        case loginAdmission(id: UUID)
+        case loginAdmission(id: UUID, cancellationRequested: Bool)
         case primaryLogin(id: UUID)
 
         var id: UUID {
@@ -150,7 +150,7 @@ final class AccountRuntimeTransitionCoordinator {
                  .primaryReconciliation(let id, _),
                  .explicitRuntimeStart(let id, _, _),
                  .runtimeAuthReconciliation(let id, _, _),
-                 .loginAdmission(let id),
+                 .loginAdmission(let id, _),
                  .primaryLogin(let id):
                 return id
             }
@@ -210,28 +210,52 @@ final class AccountRuntimeTransitionCoordinator {
             throw CodexReviewAuthenticationFailure.accountMutationBlockedByAuthentication
         }
         let id = UUID()
-        activeTransition = .loginAdmission(id: id)
+        activeTransition = .loginAdmission(id: id, cancellationRequested: false)
         return .init(id: id)
     }
 
     func canCommitLoginAdmission(_ admission: LoginAdmission) -> Bool {
-        guard case .loginAdmission(let id) = activeTransition,
-              id == admission.id else {
+        guard case .loginAdmission(let id, let cancellationRequested) = activeTransition,
+              id == admission.id,
+              cancellationRequested == false else {
             return false
         }
         return canPublish
     }
 
+    func requestLoginAdmissionCancellation() -> LoginAdmission? {
+        guard case .loginAdmission(let id, cancellationRequested: false) = activeTransition else {
+            return nil
+        }
+        activeTransition = .loginAdmission(id: id, cancellationRequested: true)
+        return .init(id: id)
+    }
+
+    func isLoginAdmissionCancellationRequested(_ admission: LoginAdmission) -> Bool {
+        guard case .loginAdmission(let id, let cancellationRequested) = activeTransition,
+              id == admission.id else {
+            return false
+        }
+        return cancellationRequested
+    }
+
+    func waitForLoginAdmissionCompletion(_ admission: LoginAdmission) async {
+        guard activeTransition?.id == admission.id else {
+            return
+        }
+        await waitForTransitionCompletion()
+    }
+
     func finishLoginAdmission(_ admission: LoginAdmission) {
-        guard case .loginAdmission(let id) = activeTransition,
+        guard case .loginAdmission(let id, _) = activeTransition,
               id == admission.id else {
             preconditionFailure("Only the active login admission can finish.")
         }
         finishTransition(id: id)
     }
 
     func retainPrimaryLoginAdmission(_ admission: LoginAdmission) {
-        guard case .loginAdmission(let id) = activeTransition,
+        guard case .loginAdmission(let id, cancellationRequested: false) = activeTransition,
               id == admission.id else {
             preconditionFailure("Only the active login admission can retain primary runtime ownership.")
         }

Sources/CodexReviewHost/LiveCodexReviewStoreBackend.swift

@@ -1447,6 +1447,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend {
 
     func cancelAuthentication(auth _: CodexReviewAuthModel) async {
         guard let session = loginSession else {
+            if let admission = accountRuntimeTransitionCoordinator
+                .requestLoginAdmissionCancellation() {
+                await accountRuntimeTransitionCoordinator.waitForLoginAdmissionCompletion(admission)
+                return
+            }
             if let activePrimaryAuthenticationReconciliation {
                 _ = await activePrimaryAuthenticationReconciliation.finalResult.wait()
             }
@@ -2002,6 +2007,11 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend {
         do {
             try await attachedStore?.requireReviewThreadRetentionAcceptance()
             guard accountRuntimeTransitionCoordinator.canCommitLoginAdmission(loginAdmission) else {
+                if accountRuntimeTransitionCoordinator
+                    .isLoginAdmissionCancellationRequested(loginAdmission) {
+                    accountRuntimeTransitionCoordinator.finishLoginAdmission(loginAdmission)
+                    return
+                }
                 throw CodexReviewAuthenticationFailure.accountMutationBlockedByAuthentication
             }
             authenticationMutation = try await accountRegistry.beginAuthenticationMutation(
@@ -2012,10 +2022,18 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend {
             throw error
         }
         let mutationLease = authenticationMutation.lease
+        let cancellationRequested = accountRuntimeTransitionCoordinator
+            .isLoginAdmissionCancellationRequested(loginAdmission)
         guard accountRuntimeTransitionCoordinator.canCommitLoginAdmission(loginAdmission),
               loginSession == nil else {
+            if cancellationRequested {
+                await accountRegistry.requestAuthenticationCancellation(mutationLease)
+            }
             await accountRegistry.finishMutation(mutationLease)
             accountRuntimeTransitionCoordinator.finishLoginAdmission(loginAdmission)
+            if cancellationRequested {
+                return
+            }
             throw CodexReviewAuthenticationFailure.accountMutationBlockedByAuthentication
         }
         let purpose = authenticationMutation.purpose
@@ -2162,6 +2180,10 @@ private final class LiveCodexReviewStoreBackend: CodexReviewStoreBackend {
                         return finish(.cancelled)
                     }
                     await self?.authenticationOperationDidBind?()
+                    guard case .proceed = await operationState.claimAPIKeyRequest() else {
+                        await startCompletion.resolve(.success(()))
+                        return finish(.cancelled)
+                    }
                     auth?.updatePhase(.signingIn(.init(
                         title: "Sign in to Codex",
                         detail: "Authenticating with API key."

Sources/CodexReviewHost/LoginSession.swift

@@ -191,6 +191,7 @@ actor LoginOperationState {
         case runtimeBound(LoginRuntime)
         case loginPending(LoginRuntime, CodexLoginHandle)
         case apiKeyPending(LoginRuntime)
+        case apiKeyRequestClaimed(LoginRuntime)
         case resourcesTaken
     }
 
@@ -214,7 +215,7 @@ actor LoginOperationState {
         switch phase {
         case .loginPending(_, let handle):
             action = .chatGPT(handle)
-        case .apiKeyPending:
+        case .apiKeyPending, .apiKeyRequestClaimed:
             action = .apiKeyRootTask
         case .acquiringRuntime, .runtimeBound, .resourcesTaken:
             return nil
@@ -282,6 +283,18 @@ actor LoginOperationState {
         return .cancel
     }
 
+    func claimAPIKeyRequest() -> BindDisposition {
+        guard case .apiKeyPending(let runtime) = phase else {
+            preconditionFailure("An API-key request can be claimed only once after binding its runtime.")
+        }
+        guard cancellationRequested == false else {
+            cancellationClaimed = true
+            return .cancel
+        }
+        phase = .apiKeyRequestClaimed(runtime)
+        return .proceed
+    }
+
     func claimURLPresentation(handle: CodexLoginHandle) -> BindDisposition {
         guard case .loginPending(_, let boundHandle) = phase,
               boundHandle == handle else {
@@ -299,7 +312,8 @@ actor LoginOperationState {
         switch phase {
         case .acquiringRuntime, .resourcesTaken:
             return nil
-        case .runtimeBound(let runtime), .loginPending(let runtime, _), .apiKeyPending(let runtime):
+        case .runtimeBound(let runtime), .loginPending(let runtime, _),
+             .apiKeyPending(let runtime), .apiKeyRequestClaimed(let runtime):
             return runtime
         }
     }
@@ -313,7 +327,8 @@ actor LoginOperationState {
 
     func takeOwnedRuntime() -> LoginRuntime? {
         switch phase {
-        case .runtimeBound(let runtime), .loginPending(let runtime, _), .apiKeyPending(let runtime):
+        case .runtimeBound(let runtime), .loginPending(let runtime, _),
+             .apiKeyPending(let runtime), .apiKeyRequestClaimed(let runtime):
             guard runtime.usesPrimaryRuntime == false else {
                 return nil
             }

Sources/CodexReviewKit/Store/CodexReviewStore.swift

@@ -240,6 +240,7 @@ public final class CodexReviewStore {
         else {
             return
         }
+        try Task.checkCancellation()
         try await signIn(using: method)
     }