@meistrari/agent-core 0.2.0-experimental.0 → 0.2.0-experimental.2

package/src/agents/codex/codex-run-auth.ts

@@ -0,0 +1,113 @@
+import type { CodexAuthProvider, CodexAuthRefreshInput, CodexAuthTokenRequest, CodexAuthTokens } from './codex-auth'
+import { AgentProviderProtocolError } from '../agent-runtime-error'
+import { SubscriptionTokenRefresh } from '../subscription-token-refresh'
+
+export class CodexRunAuthClosedError extends AgentProviderProtocolError {
+    constructor() {
+        super({ message: 'Codex run authentication is closed.' })
+    }
+}
+
+/** One initial account selection; later replacements are owned by the central refresh callback. */
+export class CodexRunAuth {
+    private readonly abort = new AbortController()
+    private readonly ready = Promise.withResolvers<void>()
+    private initialAcquisition: Promise<CodexAuthTokens> | undefined
+    private current: CodexAuthTokens | undefined
+    private readonly refresher = new SubscriptionTokenRefresh<CodexAuthTokens>({
+        installed: async () => {
+            const installed = await this.getCurrentTokens({})
+            this.abort.signal.throwIfAborted()
+            return installed
+        },
+        current: () => {
+            if (!this.current)
+                throw authSourceClosed()
+            return this.current
+        },
+        install: (tokens) => {
+            this.current = tokens
+        },
+        refresh: async (rejected) => {
+            const tokens = await waitForAuth(this.provider.refreshAuthTokens({ reason: 'unauthorized', ...rejected }), this.abort.signal)
+            this.abort.signal.throwIfAborted()
+            return tokens
+        },
+        accountId: tokens => tokens.chatgptAccountId,
+        conflictingIdentity: () => new AgentProviderProtocolError({ message: 'Codex authentication generation has conflicting token identities.' }),
+    })
+
+    constructor(private readonly provider: CodexAuthProvider) {
+        void this.ready.promise.catch(() => undefined)
+    }
+
+    async acquireInitialTokens(input: CodexAuthTokenRequest): Promise<CodexAuthTokens> {
+        this.abort.signal.throwIfAborted()
+        this.initialAcquisition ??= this.acquire(input)
+        return await this.initialAcquisition
+    }
+
+    async getCurrentTokens(input: CodexAuthTokenRequest): Promise<CodexAuthTokens> {
+        const signal = this.operationSignal(input.signal)
+        await waitForAuth(this.ready.promise, signal)
+        signal.throwIfAborted()
+        if (!this.current)
+            throw authSourceClosed()
+        return this.current
+    }
+
+    async refreshAuthTokens(input: CodexAuthRefreshInput): Promise<CodexAuthTokens> {
+        this.abort.signal.throwIfAborted()
+        return await this.refresher.refresh(input)
+    }
+
+    close(): void {
+        this.current = undefined
+        this.initialAcquisition = undefined
+        this.refresher.clear()
+        this.abort.abort(authSourceClosed())
+        this.ready.reject(authSourceClosed())
+    }
+
+    private async acquire(input: CodexAuthTokenRequest): Promise<CodexAuthTokens> {
+        const signal = this.operationSignal(input.signal)
+        try {
+            signal.throwIfAborted()
+            const tokens = await waitForAuth(this.provider.getAuthTokens({ signal }), signal)
+            signal.throwIfAborted()
+            this.current = tokens
+            this.ready.resolve()
+            return tokens
+        }
+        catch (error) {
+            this.ready.reject(error)
+            throw error
+        }
+    }
+
+    private operationSignal(signal: AbortSignal | undefined): AbortSignal {
+        return signal ? AbortSignal.any([signal, this.abort.signal]) : this.abort.signal
+    }
+}
+
+async function waitForAuth<T>(operation: Promise<T>, signal: AbortSignal): Promise<T> {
+    signal.throwIfAborted()
+    let abort: (() => void) | undefined
+    try {
+        return await Promise.race([
+            operation,
+            new Promise<never>((_, reject) => {
+                abort = () => reject(signal.reason)
+                signal.addEventListener('abort', abort, { once: true })
+            }),
+        ])
+    }
+    finally {
+        if (abort)
+            signal.removeEventListener('abort', abort)
+    }
+}
+
+function authSourceClosed(): CodexRunAuthClosedError {
+    return new CodexRunAuthClosedError()
+}