paseo/plugins/muse-provider/server/usage.ts

import { homedir } from "node:os";
import path from "node:path";
import { z } from "zod";
import type { UsageSourceRegistration } from "@getpaseo/plugin/server";
import type { ProviderLaunch } from "@getpaseo/plugin/server/provider";
import {
  hashAccountKey,
  unavailable,
  windowFromUsedPct,
  windowFromReportedDuration,
  toneFromUsedPct,
} from "@getpaseo/plugin/server/usage";

import { MspConnection } from "./connection.js";
import { usageSchema } from "./wire.js";

const inputSchema = z.object({ account: z.string() }).strict();
type Observation = z.infer<typeof usageSchema>;

function accountKey(env: ProviderLaunch["env"]): string {
  // MSP 1.4.1 has no stable account ID; swapping accounts in one config directory shares a report ID.
  const configHome = env.XDG_CONFIG_HOME ?? path.join(env.HOME ?? homedir(), ".config");
  return hashAccountKey(path.resolve(configHome, "muse"));
}

export class Usage {
  private readonly launches = new Map<string, ProviderLaunch>();
  private readonly sessions = new Map<
    string,
    { account: string; read: () => Promise<Observation> }
  >();
  remember(launch: ProviderLaunch): { account: string } {
    const key = accountKey(launch.env);
    this.launches.set(key, launch);
    return { account: key };
  }
  attach(id: string, launch: ProviderLaunch, read: () => Promise<Observation>): void {
    this.sessions.set(id, { ...this.remember(launch), read });
  }
  detach(id: string): void {
    this.sessions.delete(id);
  }
  private async read(account: string, launch: ProviderLaunch): Promise<Observation> {
    const readers = [...this.sessions.values()].filter((session) => session.account === account);
    if (readers.length > 0) {
      const observations = await Promise.all(readers.map((session) => session.read()));
      const usages = observations.flatMap((observation) =>
        observation.usage ? [observation.usage] : [],
      );
      usages.sort((a, b) => b.observedAtMs - a.observedAtMs);
      return usages.length > 0 ? { usage: usages[0] } : {};
    }
    const host = new MspConnection({ launch });
    try {
      await host.initialize();
      return await host.request("usage/read", {}, usageSchema);
    } finally {
      await host.close();
    }
  }
  registration(): UsageSourceRegistration {
    return {
      id: "muse",
      label: "Muse Code",
      icon: "icon.svg",
      input: inputSchema,
      discover: async (scope) => {
        let accounts: string[];
        if (scope.kind === "global") {
          accounts = [...this.launches.keys()];
        } else {
          if (scope.provider !== "muse") return [];
          const account = accountKey(scope.env);
          accounts = this.launches.has(account) ? [account] : [];
        }
        return accounts.map((account) => ({
          key: account,
          label: "Muse Code",
          input: { account },
        }));
      },
      fetch: async (input) => {
        const { account } = inputSchema.parse(input);
        const launch = this.launches.get(account);
        if (!launch) throw new Error("Muse login no longer exists");
        try {
          const { usage } = await this.read(account, launch);
          if (!usage) return unavailable({ kind: "no_quota", detail: "No usage quota reported" });
          return {
            status: "available",
            planLabel: usage.tier,
            windows: [
              windowFromReportedDuration({
                durationSeconds: usage.window.windowDurationMins * 60,
                unknown: { id: "window", label: "Current window", shortLabel: "" },
                summary: true,
                utilizationPct: usage.window.usedPercent,
                resetsAt: new Date(usage.window.resetsAtMs).toISOString(),
                tone: toneFromUsedPct(usage.window.usedPercent),
              }),
              windowFromUsedPct({
                id: "weekly",
                label: "Weekly",
                shortLabel: "wk",
                summary: true,
                utilizationPct: usage.weekly.usedPercent,
                resetsAt: new Date(usage.weekly.resetsAtMs).toISOString(),
                tone: toneFromUsedPct(usage.weekly.usedPercent),
              }),
            ],
            details: [
              {
                id: "observed",
                label: "Observed",
                value: new Date(usage.observedAtMs).toISOString(),
              },
            ],
          };
        } catch (error) {
          return {
            status: "error",
            error: error instanceof Error ? error.message : String(error),
          };
        }
      },
    };
  }
}