aimux

1import { atom, read, update } from 'claude-code';
2import type { EngineInterface, Register, SessionRateLimit } from 'claude-code';
3
4import type { Probe, Usage, View } from '../types';
5
6/**
7 * aimux inside Claude Code: every subscription's 5-hour and weekly usage in a band
8 * above the prompt, a warning before the one this session runs on is spent, and —
9 * once it is — `/exit` in the prompt, so one Enter moves the conversation to the
10 * freest one. `/aimux` prints the whole table with reset times.
11 *
12 * This session's own windows arrive with every response (`session.measure`). The
13 * other subscriptions are asked of aimux (`aimux status --json`), at the start and
14 * after a turn once the figures are five minutes old — an idle session asks nothing,
15 * and aimux's on-disk cache lets every open session share one probe.
16 *
17 * Moving is aimux's job, after claude exits: it carries the conversation over when it
18 * finds the hit in the transcript. The mod only tells it the person already agreed —
19 * by running the `/exit` it offered — so aimux does not ask a second time.
20 *
21 * A band of its own rather than `$.ui.status`: the engine draws that line as one of
22 * its pinned warnings, yellow with a ⚠, which reads as something being wrong.
23 *
24 * Inert in a session aimux did not launch: nothing names the profile or aimux.
25 */
26
27const REFRESH_MS = 5 * 60_000;
28/** Where the warning comes, and where a window counts as spent after a refused turn. */
29const WARN_AT = 90;
30const SPENT_AT = 95;
31/** Where the band paints a window red and says when it resets. */
32const NEAR = 80;
33const WINDOW_NAME: Record<string, string> = { five_hour: '5-hour', seven_day: 'weekly' };
34
35const pct = (p: number | null) => (p === null ? '–' : String(p));
36const windowName = (w: SessionRateLimit) => WINDOW_NAME[w.kind] ?? w.kind;
37
38/** A window's reset as epoch milliseconds, from the ISO time the engine reports. */
39const parseTime = (iso: string | undefined) => (iso ? Date.parse(iso) : undefined);
40
41/** The band's colors for a window's use, as `aimux status` paints them. */
42const levelColor = (p: number) => (p >= NEAR ? 'red' : p >= 60 ? 'yellow' : 'green');
43
44/**
45 * The freest other claude subscription: its tightest window lowest, and lower than
46 * this one's (`own`) — moving is only worth it to somewhere with more room left.
47 */
48function freest(others: Record<string, Probe>, current: string, own: number): [string, Usage] | undefined {
49  let best: [string, Usage, number] | undefined;
50
51  for (const [name, probe] of Object.entries(others)) {
52    if (name === current || probe.cli !== 'claude' || !probe.status) continue;
53
54    const windows = [probe.status.fiveHourPct, probe.status.weeklyPct].filter((p): p is number => p !== null);
55    if (windows.length === 0) continue;
56
57    const tightest = Math.max(...windows);
58    if (tightest < Math.min(own, 100) && (!best || tightest < best[2])) best = [name, probe.status, tightest];
59  }
60
61  return best && [best[0], best[1]];
62}
63
64/** ` (resets in 2h 10m)`, or nothing when the reset time is unknown or past. Counted
65 *  from `now` rather than printed as a clock time: the mod's sandbox has no time zone. */
66export function resetsIn(resetsAt: number | undefined, now: number): string {
67  if (resetsAt === undefined || !Number.isFinite(resetsAt) || resetsAt <= now) return '';
68
69  const minutes = Math.ceil((resetsAt - now) / 60_000);
70  const days = Math.floor(minutes / 1440);
71  const hours = Math.floor((minutes % 1440) / 60);
72  const text = days > 0 ? `${days}d ${hours}h` : hours > 0 ? `${hours}h ${minutes % 60}m` : `${minutes}m`;
73
74  return ` (resets in ${text})`;
75}
76
77/** One subscription: its windows, each named as claude's own status line names them;
78 *  one it does not have (codex's 5-hour one) is left out. */
79export type Cell = { name: string; isExpired: boolean; windows: { label: string; pct: number; resetsAt?: number }[] };
80
81/**
82 * The session's own subscription first, wherever it sits in the profile list, then
83 * the rest. Its own login is never in doubt: the session running is proof enough, and
84 * an older reading that said otherwise predates it.
85 */
86export function cells(view: View): Cell[] {
87  const cell = (name: string, usage: Usage | null | undefined, isExpired: boolean): Cell => ({
88    name,
89    isExpired,
90    windows: [
91      ...(usage?.fiveHourPct != null ? [{ label: '5h', pct: usage.fiveHourPct, resetsAt: usage.fiveHourResetsAt }] : []),
92      ...(usage?.weeklyPct != null ? [{ label: '7d', pct: usage.weeklyPct, resetsAt: usage.weeklyResetsAt }] : []),
93    ],
94  });
95
96  const rest = Object.entries(view.others)
97    .filter(([name]) => name !== view.current)
98    .map(([name, probe]) => cell(name, probe.status, probe.error === 'auth'));
99
100  return [cell(view.current, view.live ?? view.others[view.current]?.status, false), ...rest];
101}
102
103/** What `/aimux` prints: every subscription on its own line, each window with its reset. */
104export function table(view: View): string {
105  const all = cells(view);
106  const width = Math.max(...all.map((c) => c.name.length));
107
108  const lines = all.map((c, i) => {
109    const name = `${i === 0 ? '▸' : ' '} ${c.name.padEnd(width)}`;
110    if (c.isExpired) return `${name}  login expired — run: aimux run ${c.name}`;
111    if (c.windows.length === 0) return `${name}  no figures right now`;
112
113    const windows = c.windows.map((w) => `${w.label} ${w.pct}%${resetsIn(w.resetsAt, view.now)}`).join(' · ');
114    return `${name}  ${windows}${i === 0 ? '  ← this session' : ''}`;
115  });
116
117  return ['How much of each subscription is used', ...lines].join('\n');
118}
119
120export function warning(current: string, window: SessionRateLimit, others: Record<string, Probe>, own: number, now: number): string {
121  const head = `${current} has used ${Math.round(window.percentUsed)}% of its ${windowName(window)} window${resetsIn(parseTime(window.resetsAt), now)}.`;
122  if (Object.keys(others).length === 0) return `${head} aimux could not read the other subscriptions.`;
123
124  const next = freest(others, current, own);
125  if (!next) return `${head} No other claude subscription has more room right now.`;
126
127  const [name, usage] = next;
128  return `${head} Freest now: ${name} (5h ${pct(usage.fiveHourPct)}%, 7d ${pct(usage.weeklyPct)}%). `
129    + 'When this one runs out, aimux will offer to move this conversation there.';
130}
131
132// What the band draws: host state, so a write redraws it.
133const view = atom({ plugin: 'aimux', key: 'view' } as const, null);
134
135// The session's own figures. Module state: a hot reload starts it over, which only
136// means one more probe.
137// ponytail: an `aimux split` of N panes starts N sessions at once, each probing every
138// subscription before the shared reading exists — N×M one-token requests, once. A lock
139// around the probe would make it one round; not worth it yet.
140let current: string | undefined;
141let aimux: string[] | undefined;
142let handoff: string | undefined;
143let others: Record<string, Probe> = {};
144let live: Usage | undefined;
145let fetchedAt: number | undefined;
146let reading: Promise<boolean> | undefined;
147let offered: string | undefined;
148const warned = new Set<string>();
149
150async function snapshot($: EngineInterface): Promise<View | undefined> {
151  return current ? { current, others, live: live ?? null, now: await $.clock.now() } : undefined;
152}
153
154async function show($: EngineInterface) {
155  const next = await snapshot($);
156  if (next) await update($, view, () => next);
157}
158
159/** One question to aimux; whether it answered. */
160async function readOthers($: EngineInterface, command: string[]): Promise<boolean> {
161  try {
162    const { exitCode, stdout } = await $.process.run(command, { timeoutMs: 60_000 });
163    if (exitCode !== 0) return false;
164
165    others = JSON.parse(stdout).profiles;
166    return true;
167  } catch {
168    // aimux unreachable this time: keep the figures we have
169    return false;
170  } finally {
171    // A failed attempt waits its turn too: this runs after every response, and a
172    // broken aimux must not be started again each time.
173    fetchedAt = await $.clock.now();
174  }
175}
176
177/**
178 * Ask aimux for every subscription's figures, no older than `maxAgeSeconds`; whether
179 * it answered. One question at a time: a second caller waits for the one under way.
180 */
181function fetchOthers($: EngineInterface, maxAgeSeconds: number): Promise<boolean> {
182  if (!aimux) return Promise.resolve(false);
183
184  reading ??= readOthers($, [...aimux, 'status', '--json', '--max-age', String(maxAgeSeconds)])
185    .finally(() => { reading = undefined; });
186
187  return reading;
188}
189
190async function refresh($: EngineInterface) {
191  const now = await $.clock.now();
192  if (fetchedAt !== undefined && now - fetchedAt < REFRESH_MS) return;
193
194  await fetchOthers($, 240);
195  await show($);
196}
197
198/** `/aimux`: the whole table, read fresh — or saying so when it could not be. */
199async function report($: EngineInterface): Promise<string> {
200  // A background reading under way may be an older one than asked for here: let it
201  // land, then ask again.
202  await reading;
203  const isFresh = await fetchOthers($, 30);
204  await show($);
205
206  const now = await snapshot($);
207  if (!now) return 'aimux did not start this session, so it has nothing to show here.';
208
209  return isFresh ? table(now) : `${table(now)}\n(aimux could not be reached just now — these are the last figures it gave)`;
210}
211
212/** The window is spent: name where to go, and put `/exit` where one Enter runs it. */
213async function offerMove($: EngineInterface, window: SessionRateLimit, now: number) {
214  if (!current) return;
215
216  const head = `${current} is out of its ${windowName(window)} window${resetsIn(parseTime(window.resetsAt), now)}.`;
217  const next = freest(others, current, 100);
218  if (!next) {
219    $.ui.toast(`${head} No other claude subscription has room right now.`, { timeoutMs: 30_000 });
220    return;
221  }
222
223  const [name, usage] = next;
224  offered = name;
225
226  // Never over a draft: the person's own words stay, and they type /exit themselves.
227  const box = await $.prompt.read();
228  const filled = box.text.trim() === '' && (await $.prompt.fill({ text: '/exit' })).isFilled;
229  const how = filled ? 'Press Enter' : 'Run /exit';
230
231  $.ui.toast(
232    `${head} ${how} to carry this conversation over to ${name} (5h ${pct(usage.fiveHourPct)}%, 7d ${pct(usage.weeklyPct)}%).`,
233    { timeoutMs: 30_000 },
234  );
235}
236
237/**
238 * A spent window, by whichever sign came first — its use reaching 100%, or a turn the
239 * API refused over the limit. One offer per window, and not before aimux has answered
240 * once: the offer is about where to go next.
241 */
242async function spent($: EngineInterface, window: SessionRateLimit, now: number) {
243  const key = `${window.kind}:${window.resetsAt ?? ''}:out`;
244  if (fetchedAt === undefined || warned.has(key)) return;
245
246  warned.add(key);
247  await offerMove($, window, now);
248}
249
250/** A turn the API refused over a rate limit: spent, if a window is in fact nearly full. */
251async function refusedOverLimit($: EngineInterface) {
252  const { rateLimits } = await $.session.usage();
253  const tightest = rateLimits
254    .filter((w) => w.kind in WINDOW_NAME)
255    .sort((a, b) => b.percentUsed - a.percentUsed)[0];
256
257  // A refusal with room left in every window is the API being busy, not a spent
258  // subscription: nothing to move for.
259  if (tightest && tightest.percentUsed >= SPENT_AT) await spent($, tightest, await $.clock.now());
260}
261
262/** The person ran /exit after the offer: tell aimux, so it moves without asking again. */
263async function agreeToMove($: EngineInterface) {
264  if (offered && handoff) await $.fs.write(handoff, offered);
265}
266
267export const register: Register = (on) => {
268  on('session.start', async ($, e, next) => {
269    current = await $.env.get('AIMUX_RUN_PROFILE');
270    const self = await $.env.get('AIMUX_SELF');
271    aimux = current && self ? JSON.parse(self) : undefined;
272    handoff = await $.env.get('AIMUX_HANDOFF');
273
274    if (aimux) {
275      await $.command.register({ name: 'aimux', description: 'How much of every subscription is used, and when each resets' });
276
277      // On a timer, not awaited: the first prompt must not wait on a network probe.
278      $.clock.after(0, () => void refresh($));
279    }
280
281    return next(e);
282  });
283
284  on('session.measure', async ($, e, next) => {
285    if (!current) return next(e);
286
287    const now = await $.clock.now();
288    const five = e.rateLimits.find((w) => w.kind === 'five_hour');
289    const week = e.rateLimits.find((w) => w.kind === 'seven_day');
290    if (five || week) {
291      live = {
292        fiveHourPct: five ? Math.round(five.percentUsed) : null,
293        weeklyPct: week ? Math.round(week.percentUsed) : null,
294        fiveHourResetsAt: parseTime(five?.resetsAt),
295        weeklyResetsAt: parseTime(week?.resetsAt),
296      };
297    }
298
299    const own = Math.max(...[five, week].map((w) => w?.percentUsed ?? 0));
300    for (const window of [five, week]) {
301      if (!window) continue;
302
303      // One notice per window and stage: a new reset time is a new window, and one
304      // that eased back under the line may warn again.
305      const at = `${window.kind}:${window.resetsAt ?? ''}`;
306      if (window.percentUsed < WARN_AT) {
307        warned.delete(`${at}:near`);
308        warned.delete(`${at}:out`);
309        continue;
310      }
311
312      if (window.percentUsed >= 100) {
313        await spent($, window, now);
314        continue;
315      }
316
317      // Not before aimux has answered once: the notice names where to go next.
318      if (fetchedAt === undefined || warned.has(`${at}:near`)) continue;
319
320      warned.add(`${at}:near`);
321      $.ui.toast(warning(current, window, others, own, now), { timeoutMs: 15_000 });
322    }
323
324    await show($);
325    $.clock.after(0, () => void refresh($));
326    return next(e);
327  });
328
329  on('classic.StopFailure', async ($, e, next) => {
330    if (current && e.error === 'rate_limit') await refusedOverLimit($);
331    return next(e);
332  });
333
334  on('command.run', { command: 'exit' }, async ($, e, next) => {
335    await agreeToMove($);
336    return next(e);
337  });
338
339  on('command.run', { command: 'aimux' }, async ($) => ({ text: await report($) }));
340
341  // `aimux  ▸ dt (this session) 5h:27% 7d:43%  │  main 5h:13% 7d:12% · cx 7d:36%`
342  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
343    const drawn = await read($, view);
344    if (drawn === null || e.props.hasSurvey) return next(e);
345
346    const { Box, Text } = $.ui.resolve(e);
347    const [own, ...rest] = cells(drawn);
348    if (!own) return next(e);
349
350    const usage = (cell: Cell) => {
351      if (cell.isExpired) return [<Text color="yellow">login expired</Text>];
352      if (cell.windows.length === 0) return [<Text dimColor>–</Text>];
353
354      // A window close to its limit also says when it frees up.
355      return cell.windows.flatMap((w, i) => [
356        <Text dimColor>{`${i > 0 ? ' ' : ''}${w.label}:`}</Text>,
357        <Text color={levelColor(w.pct)}>{`${w.pct}%`}</Text>,
358        <Text dimColor>{w.pct >= NEAR ? resetsIn(w.resetsAt, drawn.now) : ''}</Text>,
359      ]);
360    };
361
362    // One Text of colored spans, not a row of boxes: a row squeezes each box on a
363    // narrow terminal and wraps them one by one; a single text wraps as a line does.
364    return (
365      <Box>
366        <Text>
367          <Text dimColor>aimux  </Text>
368          <Text bold>{`▸ ${own.name} `}</Text>
369          <Text dimColor>(this session) </Text>
370          {usage(own)}
371          {rest.length > 0 ? <Text dimColor>{'  │  '}</Text> : null}
372          {rest.flatMap((cell, i) => [
373            <Text dimColor>{`${i > 0 ? ' · ' : ''}`}</Text>,
374            <Text>{`${cell.name} `}</Text>,
375            ...usage(cell),
376          ])}
377        </Text>
378      </Box>
379    );
380  });
381};
382

1export type Usage = {
2  fiveHourPct: number | null;
3  weeklyPct: number | null;
4  /** When each window resets, in epoch milliseconds, where known. */
5  fiveHourResetsAt?: number;
6  weeklyResetsAt?: number;
7};
8export type Probe = { cli: string; status: Usage | null; error?: string };
9
10/** What the band above the prompt draws: this session's subscription and the rest,
11 *  as of `now` (which the reset times are counted from). */
12export type View = { current: string; others: Record<string, Probe>; live: Usage | null; now: number };
13
14declare module 'claude-code' {
15  interface PluginState {
16    aimux: { view: View | null };
17  }
18}
19