packages/storage/src/usage-stores.ts

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

import type {
  ModelCallUsageBuckets,
  ModelCallUsageLogs,
  ModelCallUsageSummary,
} from '@maka/core/model-call-usage-projection';
import type {
  PricingConfig,
  UsageBucket,
  UsageGroupBy,
  UsageLogRow,
  UsageQuery,
  UsageSummaryV2,
} from '@maka/core/usage-stats/types';
import type { UsageScreenRequest, UsageScreenResult } from '@maka/core/settings';
import { createUsageScreenReader } from './usage-screen.js';
import { throwDeduplicatedFailures } from './failure-utils.js';
import {
  createSqliteModelCallLedger,
  type CatchUpModelCallProjectionInput,
  type CatchUpModelCallProjectionResult,
  ModelCallLedgerClosedError,
  ModelCallLedgerPublicationError,
  type ModelCallLedger,
  type ModelCallLedgerReader,
} from './model-call-ledger.js';
import {
  PricingCommitUnknownError,
  PricingRevisionConflictError,
  PricingStoreClosedError,
  PricingStoreNotLoadedError,
  PricingStorePublicationError,
  PricingValidationError,
  type PricingMutationResult,
  type PricingSnapshot,
  type PricingStore,
} from './pricing-store.js';
import {
  runWithStorageRootLease,
  StorageRootAuthorityError,
  type StorageRootAuthorityErrorCode,
  type StorageRootLease,
} from './root-authority.js';
import {
  TelemetryQueryValidationError,
  TelemetryRepoClosedError,
  TelemetryRepoNotLoadedError,
  TelemetryRepoPublicationError,
  type PersistedLlmCallRecord,
  type PersistedToolInvocationRecord,
  type TelemetryRepo,
  type ToolUsageQuery,
} from './telemetry-repo.js';
import { createSqlitePricingStore, createSqliteTelemetryRepo } from './sqlite-usage-store.js';

const readerBrand: unique symbol = Symbol('InteractiveUsageStoresReader');
const writerBrand: unique symbol = Symbol('InteractiveUsageStoresWriter');
const readers = new WeakSet<object>();
const writers = new WeakSet<object>();
const writerByLease = new WeakMap<object, InteractiveUsageStoresWriter>();
const writerOpeningByLease = new WeakMap<object, Promise<InteractiveUsageStoresWriter>>();

export interface TelemetryIndexReader {
  summary(query: UsageQuery): Promise<UsageSummaryV2>;
  toolSummary(query: UsageQuery): Promise<{ requests: number; durationMs: number }>;
  buckets(query: UsageQuery, groupBy: UsageGroupBy): Promise<UsageBucket[]>;
  logs(
    query: UsageQuery,
    offset?: number,
    limit?: number,
  ): Promise<{ rows: UsageLogRow[]; total: number }>;
  toolLogs(
    query: ToolUsageQuery,
    offset?: number,
    limit?: number,
  ): Promise<{ rows: PersistedToolInvocationRecord[]; total: number }>;
  latestLlmRuntimeProbe(connectionSlug: string, modelId?: string): Promise<UsageLogRow | undefined>;
}

export interface TelemetryIndexWriter extends TelemetryIndexReader {
  recordLlmCall(record: PersistedLlmCallRecord): Promise<void>;
  recordToolInvocation(record: PersistedToolInvocationRecord): Promise<void>;
}

/**
 * Read side of the canonical model-call ledger (#1679). Async here — unlike the
 * synchronous store beneath it — because every authority read goes through the
 * storage-root lease.
 */
/** One Usage answer from the canonical ledger, with the rows it could not read. */
export interface ModelCallLedgerResult<T> {
  readonly projection: T;
  readonly unreadableRecords: number;
}

export interface ModelCallIndexReader {
  modelCallSummary(
    query: UsageQuery,
    now: number,
  ): Promise<ModelCallLedgerResult<ModelCallUsageSummary>>;
  modelCallBuckets(
    query: UsageQuery,
    groupBy: UsageGroupBy,
    now: number,
  ): Promise<ModelCallLedgerResult<ModelCallUsageBuckets>>;
  modelCallLogs(
    query: UsageQuery,
    now: number,
    offset: number,
    limit: number,
  ): Promise<ModelCallLedgerResult<ModelCallUsageLogs>>;
}

export interface ModelCallIndexWriter extends ModelCallIndexReader {
  catchUpModelCallProjection(
    input?: CatchUpModelCallProjectionInput,
  ): Promise<CatchUpModelCallProjectionResult>;
}

export interface PricingAuthorityReader {
  snapshot(): Promise<PricingSnapshot>;
}

export interface PricingAuthorityWriter extends PricingAuthorityReader {
  upsert(expectedRevision: number, pricing: PricingConfig): Promise<PricingMutationResult>;
  delete(expectedRevision: number, modelKey: string): Promise<PricingMutationResult>;
}

export interface InteractiveUsageStoresReader {
  readUsageScreen(input: UsageScreenRequest): Promise<UsageScreenResult>;
  readonly kind: 'interactive';
  readonly access: 'read';
  readonly [readerBrand]: true;
  readonly telemetry: Readonly<TelemetryIndexReader>;
  readonly modelCalls: Readonly<ModelCallIndexReader>;
  readonly pricing: Readonly<PricingAuthorityReader>;
  close(): Promise<void>;
}

export interface InteractiveUsageStoresWriter {
  readUsageScreen(input: UsageScreenRequest): Promise<UsageScreenResult>;
  readonly kind: 'interactive';
  readonly access: 'write';
  readonly [writerBrand]: true;
  readonly telemetry: Readonly<TelemetryIndexWriter>;
  readonly modelCalls: Readonly<ModelCallIndexWriter>;
  readonly pricing: Readonly<PricingAuthorityWriter>;
  subscribeSessionUsageChanges(listener: (sessionId: string) => void): () => void;
  beginDrain(): Promise<void>;
  flush(): Promise<void>;
  close(): Promise<void>;
}

export class InteractiveUsageStoresClosedError extends Error {
  constructor() {
    super('Interactive usage stores are draining or closed');
    this.name = 'InteractiveUsageStoresClosedError';
  }
}

export type InteractiveUsageStoresFailureClassification =
  | { readonly kind: 'lifecycle' }
  | {
      readonly kind: 'revision_conflict';
      readonly expectedRevision: number;
      readonly actualRevision: number;
    }
  | { readonly kind: 'invalid_request' }
  | { readonly kind: 'commit_outcome_unknown'; readonly needsDrain: true }
  | { readonly kind: 'persistence_failed'; readonly needsDrain: boolean }
  | { readonly kind: 'unknown'; readonly error: unknown };

export function classifyInteractiveUsageStoresFailure(
  error: unknown,
): InteractiveUsageStoresFailureClassification {
  if (error instanceof PricingRevisionConflictError) {
    return {
      kind: 'revision_conflict',
      expectedRevision: error.expectedRevision,
      actualRevision: error.actualRevision,
    };
  }
  if (error instanceof PricingValidationError || error instanceof TelemetryQueryValidationError) {
    return { kind: 'invalid_request' };
  }
  if (
    error instanceof InteractiveUsageStoresClosedError ||
    error instanceof PricingStoreClosedError ||
    error instanceof TelemetryRepoClosedError ||
    error instanceof ModelCallLedgerClosedError ||
    (error instanceof StorageRootAuthorityError &&
      (error.code === 'invalid_lease' || error.code === 'invalid_owner'))
  ) {
    return { kind: 'lifecycle' };
  }
  if (
    error instanceof PricingCommitUnknownError ||
    (error instanceof TelemetryRepoPublicationError && error.commitUnknown) ||
    (error instanceof ModelCallLedgerPublicationError && error.commitUnknown)
  ) {
    return { kind: 'commit_outcome_unknown', needsDrain: true };
  }
  if (
    error instanceof PricingStorePublicationError ||
    error instanceof TelemetryRepoPublicationError ||
    error instanceof ModelCallLedgerPublicationError
  ) {
    return { kind: 'persistence_failed', needsDrain: true };
  }
  if (error instanceof PricingStoreNotLoadedError || error instanceof TelemetryRepoNotLoadedError) {
    return { kind: 'persistence_failed', needsDrain: false };
  }
  if (error instanceof StorageRootAuthorityError) {
    return {
      kind: 'persistence_failed',
      needsDrain: rootAuthorityFailureNeedsDrain(error.code),
    };
  }
  return { kind: 'unknown', error };
}

function rootAuthorityFailureNeedsDrain(code: StorageRootAuthorityErrorCode): boolean {
  switch (code) {
    case 'root_unmarked':
    case 'invalid_marker':
    case 'root_identity_collision':
    case 'root_identity_changed':
      return true;
    case 'invalid_root':
    case 'invalid_root_kind':
    case 'root_not_found':
    case 'invalid_repair':
    case 'invalid_capability':
    case 'invalid_lease':
    case 'invalid_owner':
    case 'invalid_lock_artifact':
    case 'insecure_control_directory':
    case 'root_io_failed':
    case 'control_io_failed':
    case 'lock_failed':
      return false;
  }
}

export function authenticateInteractiveUsageStoresReader(
  stores: InteractiveUsageStoresReader,
): InteractiveUsageStoresReader {
  if (!readers.has(stores)) throw new TypeError('Expected an authentic interactive usage reader');
  return stores;
}

export function authenticateInteractiveUsageStoresWriter(
  stores: InteractiveUsageStoresWriter,
): InteractiveUsageStoresWriter {
  if (!writers.has(stores)) throw new TypeError('Expected an authentic interactive usage writer');
  return stores;
}

export async function openInteractiveUsageStoresForRead(
  lease: StorageRootLease<'interactive', 'read'>,
): Promise<InteractiveUsageStoresReader> {
  const repos = await runWithStorageRootLease(lease, 'interactive', 'read', (root) =>
    openRepos(root, false),
  );
  let closed = false;
  let closePromise: Promise<void> | undefined;
  const run = <T>(operation: () => T | Promise<T>): Promise<T> => {
    if (closed) return Promise.reject(new InteractiveUsageStoresClosedError());
    return runWithStorageRootLease(lease, 'interactive', 'read', async () => operation());
  };
  const stores: InteractiveUsageStoresReader = {
    readUsageScreen: (input) => run(() => repos.screen.read(input)),
    kind: 'interactive',
    access: 'read',
    [readerBrand]: true,
    telemetry: telemetryReader(repos.telemetry, run),
    modelCalls: modelCallReader(repos.modelCalls, run),
    pricing: pricingReader(repos.pricing, run),
    close: () => {
      if (closePromise) return closePromise;
      closed = true;
      repos.screen.close();
      closePromise = closeRepos(repos.telemetry, repos.modelCalls, repos.pricing);
      return closePromise;
    },
  };
  freezeFacade(stores);
  readers.add(stores);
  return stores;
}

export async function openInteractiveUsageStoresForWrite(
  lease: StorageRootLease<'interactive', 'write'>,
): Promise<InteractiveUsageStoresWriter> {
  const existing = writerByLease.get(lease);
  if (existing) return existing;
  const opening = writerOpeningByLease.get(lease);
  if (opening) return opening;
  const pending = runWithStorageRootLease(lease, 'interactive', 'write', async (root) => {
    const repos = await openRepos(root, true);
    const stores = createWriterFacade(
      lease,
      repos.telemetry,
      repos.modelCalls,
      repos.pricing,
      repos.screen,
    );
    writers.add(stores);
    writerByLease.set(lease, stores);
    return stores;
  });
  writerOpeningByLease.set(lease, pending);
  try {
    return await pending;
  } finally {
    if (writerOpeningByLease.get(lease) === pending) writerOpeningByLease.delete(lease);
  }
}

async function openRepos(
  root: string,
  createIfMissing: boolean,
): Promise<{
  telemetry: TelemetryRepo;
  modelCalls: ModelCallLedger;
  pricing: PricingStore;
  screen: ReturnType<typeof createUsageScreenReader>;
}> {
  const telemetry = createSqliteTelemetryRepo(root, { createIfMissing, managePricing: false });
  await telemetry.load();
  const modelCalls = createSqliteModelCallLedger(root);
  const pricing = createSqlitePricingStore(root, { createIfMissing });
  try {
    await pricing.load();
    return { telemetry, modelCalls, pricing, screen: createUsageScreenReader(root) };
  } catch (error) {
    const closed = await Promise.allSettled([
      telemetry.close(),
      modelCalls.close(),
      pricing.close(),
    ]);
    const failures = [error, ...rejectedReasons(closed)];
    throwDeduplicatedFailures('Unable to open interactive usage stores', failures);
    throw error;
  }
}

function createWriterFacade(
  lease: StorageRootLease<'interactive', 'write'>,
  telemetry: TelemetryRepo,
  modelCalls: ModelCallLedger,
  pricing: PricingStore,
  screen: ReturnType<typeof createUsageScreenReader>,
): InteractiveUsageStoresWriter {
  const run = <T>(operation: () => T | Promise<T>): Promise<T> =>
    runWithStorageRootLease(lease, 'interactive', 'write', async () => operation());
  let state: 'open' | 'draining' | 'closed' = 'open';
  let barrier: Promise<void> = Promise.resolve();
  const failures: unknown[] = [];
  let drainPromise: Promise<void> | undefined;
  let closePromise: Promise<void> | undefined;
  const sessionUsageChangeListeners = new Set<(sessionId: string) => void>();

  const publishSessionUsageChange = (sessionId: string): void => {
    for (const listener of sessionUsageChangeListeners) {
      try {
        listener(sessionId);
      } catch {
        /* observers cannot perturb the Usage authority */
      }
    }
  };

  const assertOpen = () => {
    if (state !== 'open') throw new InteractiveUsageStoresClosedError();
  };
  const admit = <T>(
    operation: () => Promise<T>,
    expectedFailure: (error: unknown) => boolean = () => false,
  ): Promise<T> => {
    assertOpen();
    const admitted = Promise.resolve().then(operation);
    const observed = admitted.then(
      () => undefined,
      (error: unknown) => {
        if (!expectedFailure(error)) failures.push(error);
      },
    );
    barrier = Promise.all([barrier, observed]).then(() => undefined);
    return admitted;
  };
  const admitSessionUsageMutation = <T>(
    sessionId: string | undefined,
    operation: () => T | Promise<T>,
  ): Promise<T> =>
    admit(async () => {
      const result = await run(operation);
      if (sessionId) publishSessionUsageChange(sessionId);
      return result;
    });
  const admitSessionUsageChange = (
    sessionId: string,
    operation: () => Promise<boolean>,
  ): Promise<void> =>
    admit(async () => {
      if (await run(operation)) publishSessionUsageChange(sessionId);
    });
  const admitModelCallProjectionCatchUp = (
    input?: CatchUpModelCallProjectionInput,
  ): Promise<CatchUpModelCallProjectionResult> =>
    admit(async () => {
      const result = await run(() => modelCalls.catchUpProjection(input));
      for (const sessionId of result.changedSessionIds) publishSessionUsageChange(sessionId);
      return result;
    });
  const read = <T>(operation: () => T): Promise<T> => {
    assertOpen();
    return run(operation);
  };

  const beginDrain = (): Promise<void> => {
    if (drainPromise) return drainPromise;
    state = 'draining';
    const accepted = barrier;
    drainPromise = accepted.then(() =>
      throwDeduplicatedFailures('Interactive usage store drain failed', failures),
    );
    return drainPromise;
  };

  const flush = async (): Promise<void> => {
    const accepted = barrier;
    await accepted;
    const flushed = await Promise.allSettled([
      run(() => telemetry.flush()),
      run(() => modelCalls.flush()),
      run(() => pricing.flush()),
    ]);
    throwDeduplicatedFailures('Interactive usage store flush failed', [
      ...failures,
      ...rejectedReasons(flushed),
    ]);
  };

  const close = (): Promise<void> => {
    if (closePromise) return closePromise;
    state = 'draining';
    if (writerByLease.get(lease) === stores) writerByLease.delete(lease);
    writers.delete(stores);
    const accepted = barrier;
    closePromise = accepted
      .then(async () => {
        const closed = await Promise.allSettled([
          telemetry.close(),
          modelCalls.close(),
          pricing.close(),
        ]);
        throwDeduplicatedFailures('Interactive usage stores close failed', [
          ...failures,
          ...rejectedReasons(closed),
        ]);
      })
      .finally(() => {
        state = 'closed';
        sessionUsageChangeListeners.clear();
        screen.close();
      });
    return closePromise;
  };

  const stores: InteractiveUsageStoresWriter = {
    readUsageScreen: (input) => read(() => screen.read(input)),
    kind: 'interactive',
    access: 'write',
    [writerBrand]: true,
    telemetry: {
      summary: (query) => read(() => telemetry.summary(query)),
      toolSummary: (query) => read(() => telemetry.toolSummary(query)),
      buckets: (query, groupBy) => read(() => telemetry.buckets(query, groupBy)),
      logs: (query, offset, limit) => read(() => telemetry.logs(query, offset, limit)),
      toolLogs: (query, offset, limit) => read(() => telemetry.toolLogs(query, offset, limit)),
      latestLlmRuntimeProbe: (connectionSlug, modelId) =>
        read(() => telemetry.latestLlmRuntimeProbe(connectionSlug, modelId)),
      recordLlmCall: (record) =>
        admitSessionUsageMutation(record.sessionId, () => telemetry.insertLlmCall(record)),
      recordToolInvocation: (record) =>
        admitSessionUsageMutation(record.sessionId, () => telemetry.insertToolInvocation(record)),
    },
    modelCalls: {
      modelCallSummary: (query, now) => read(() => modelCalls.summary(query, now)),
      modelCallBuckets: (query, groupBy, now) =>
        read(() => modelCalls.buckets(query, groupBy, now)),
      modelCallLogs: (query, now, offset, limit) =>
        read(() => modelCalls.logs(query, now, offset, limit)),
      catchUpModelCallProjection: admitModelCallProjectionCatchUp,
    },
    pricing: {
      snapshot: () => read(() => pricing.snapshot()),
      upsert: (expectedRevision, value) =>
        admit(() => run(() => pricing.upsert(expectedRevision, value)), isExpectedPricingFailure),
      delete: (expectedRevision, modelKey) =>
        admit(
          () => run(() => pricing.delete(expectedRevision, modelKey)),
          isExpectedPricingFailure,
        ),
    },
    subscribeSessionUsageChanges(listener) {
      assertOpen();
      sessionUsageChangeListeners.add(listener);
      return () => sessionUsageChangeListeners.delete(listener);
    },
    beginDrain,
    flush,
    close,
  };
  freezeFacade(stores);
  return stores;
}

function telemetryReader(
  repo: TelemetryRepo,
  run: <T>(operation: () => T | Promise<T>) => Promise<T>,
): Readonly<TelemetryIndexReader> {
  return Object.freeze({
    summary: (query: UsageQuery) => run(() => repo.summary(query)),
    toolSummary: (query: UsageQuery) => run(() => repo.toolSummary(query)),
    buckets: (query: UsageQuery, groupBy: UsageGroupBy) => run(() => repo.buckets(query, groupBy)),
    logs: (query: UsageQuery, offset?: number, limit?: number) =>
      run(() => repo.logs(query, offset, limit)),
    toolLogs: (query: ToolUsageQuery, offset?: number, limit?: number) =>
      run(() => repo.toolLogs(query, offset, limit)),
    latestLlmRuntimeProbe: (connectionSlug: string, modelId?: string) =>
      run(() => repo.latestLlmRuntimeProbe(connectionSlug, modelId)),
  });
}

function modelCallReader(
  ledger: ModelCallLedgerReader,
  run: <T>(operation: () => T | Promise<T>) => Promise<T>,
): Readonly<ModelCallIndexReader> {
  return Object.freeze({
    modelCallSummary: (query: UsageQuery, now: number) => run(() => ledger.summary(query, now)),
    modelCallBuckets: (query: UsageQuery, groupBy: UsageGroupBy, now: number) =>
      run(() => ledger.buckets(query, groupBy, now)),
    modelCallLogs: (query: UsageQuery, now: number, offset: number, limit: number) =>
      run(() => ledger.logs(query, now, offset, limit)),
  });
}

function pricingReader(
  store: PricingStore,
  run: <T>(operation: () => T | Promise<T>) => Promise<T>,
): Readonly<PricingAuthorityReader> {
  return Object.freeze({ snapshot: () => run(() => store.snapshot()) });
}

function isExpectedPricingFailure(error: unknown): boolean {
  return error instanceof PricingRevisionConflictError || error instanceof PricingValidationError;
}

async function closeRepos(
  telemetry: TelemetryRepo,
  modelCalls: ModelCallLedger,
  pricing: PricingStore,
): Promise<void> {
  const closed = await Promise.allSettled([telemetry.close(), modelCalls.close(), pricing.close()]);
  throwDeduplicatedFailures('Unable to close interactive usage stores', rejectedReasons(closed));
}

function rejectedReasons(results: readonly PromiseSettledResult<unknown>[]): unknown[] {
  return results.flatMap((result) => (result.status === 'rejected' ? [result.reason] : []));
}

function freezeFacade(stores: InteractiveUsageStoresReader | InteractiveUsageStoresWriter): void {
  Object.freeze(stores.telemetry);
  Object.freeze(stores.modelCalls);
  Object.freeze(stores.pricing);
  Object.freeze(stores);
}