makelore/electron/services/works-square-runtime.ts

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<unknown>;
  };
  imageWorkspace?: {
    closeEventSessions?(options?: {
      accessToken?: string;
      tolerateRemoteFailure?: boolean;
    }): Promise<unknown>;
  };
};

let cleanupFlight: Promise<void> | null = null;
let cleanupRequired = false;

async function runManagedWorksSquareRuntimeCleanup(
  ctx: WorksSquareRuntimeContext,
  accessToken?: string,
  tolerateRemoteFailure = false,
): Promise<void> {
  // 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<void> {
  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<void> {
  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<void> {
  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;
}