import { NIANCODE_USER_MODEL_ACCOUNT_ID } from '../../shared/user-model-config'; import { logger } from '../utils/logger'; import { clearWorksSquareAIGatewayCredential } from './works-square-ai-gateway'; import { getProviderService } from './providers/provider-service'; export type WorksSquareRuntimeContext = { opencodeManager: { stop(): Promise; }; imageWorkspace?: { closeEventSessions?(options?: { accessToken?: string; tolerateRemoteFailure?: boolean; }): Promise; }; }; let cleanupFlight: Promise | null = null; let cleanupRequired = false; async function runManagedWorksSquareRuntimeCleanup( ctx: WorksSquareRuntimeContext, accessToken?: string, tolerateRemoteFailure = false, ): Promise { // Revoke the in-memory derived credential synchronously; slower resource cleanup follows. clearWorksSquareAIGatewayCredential(); const results = await Promise.allSettled([ (async () => await ctx.imageWorkspace?.closeEventSessions?.({ accessToken, tolerateRemoteFailure, }))(), (async () => await ctx.opencodeManager.stop())(), (async () => { const deleted = await getProviderService().deleteAccountApiKey( NIANCODE_USER_MODEL_ACCOUNT_ID, ); if (!deleted) throw new Error('provider API key storage rejected deletion'); })(), ]); const errors = results .filter((result): result is PromiseRejectedResult => result.status === 'rejected') .map((result) => ( result.reason instanceof Error ? result.reason.message : String(result.reason) )); if (errors.length > 0) { throw new Error(`Failed to clear managed Works Square runtime state: ${errors.join('; ')}`); } } export function clearManagedWorksSquareRuntime( ctx: WorksSquareRuntimeContext, accessToken?: string, tolerateRemoteFailure = false, ): Promise { if (cleanupFlight) return cleanupFlight; cleanupRequired = true; const flight = runManagedWorksSquareRuntimeCleanup( ctx, accessToken, tolerateRemoteFailure, ) .then(() => { cleanupRequired = false; }) .finally(() => { if (cleanupFlight === flight) cleanupFlight = null; }); cleanupFlight = flight; return flight; } export async function ensureManagedWorksSquareRuntimeClean( ctx: WorksSquareRuntimeContext, ): Promise { if (cleanupFlight) { try { await cleanupFlight; } catch { // Retry below with the current runtime context. } } if (cleanupRequired) await clearManagedWorksSquareRuntime(ctx); } export async function clearManagedWorksSquareRuntimeBestEffort( ctx: WorksSquareRuntimeContext, reason: string, accessToken?: string, ): Promise { try { await clearManagedWorksSquareRuntime(ctx, accessToken, true); } catch (error) { logger.error(`[auth] Failed to clear managed runtime after ${reason}`, error); } } export function resetManagedWorksSquareRuntimeForTests(): void { cleanupFlight = null; cleanupRequired = false; }