Merge pull request #2699 from Yeachan-Heo/issue-2695-hud-rate-limits-stdin

Prefer stdin rate limits over cold-start HUD API fetches
This commit is contained in:
Bellman
2026-04-17 22:25:47 +09:00
committed by GitHub
5 changed files with 224 additions and 44 deletions

View File

@@ -1,7 +1,7 @@
import { describe, expect, it } from 'vitest';
import type { StatuslineStdin } from '../../hud/types.js';
import { getContextPercent, getModelName, stabilizeContextPercent } from '../../hud/stdin.js';
import { getContextPercent, getModelName, getRateLimitsFromStdin, stabilizeContextPercent } from '../../hud/stdin.js';
function makeStdin(overrides: Partial<StatuslineStdin> = {}): StatuslineStdin {
return {
@@ -176,3 +176,49 @@ describe('HUD stdin model display', () => {
expect(getModelName(makeStdin({ model: undefined }))).toBe('Unknown');
});
});
describe('HUD stdin rate limits', () => {
it('parses stdin rate_limits into the existing RateLimits shape', () => {
const result = getRateLimitsFromStdin(makeStdin({
rate_limits: {
five_hour: {
used_percentage: 11,
resets_at: 1776348000,
},
seven_day: {
used_percentage: 2,
resets_at: '2026-04-22T00:00:00.000Z',
},
},
}));
expect(result).toEqual({
fiveHourPercent: 11,
weeklyPercent: 2,
fiveHourResetsAt: new Date(1776348000 * 1000),
weeklyResetsAt: new Date('2026-04-22T00:00:00.000Z'),
});
});
it('returns null when stdin omits rate limits', () => {
expect(getRateLimitsFromStdin(makeStdin())).toBeNull();
});
it('tolerates invalid reset values without breaking the result', () => {
const result = getRateLimitsFromStdin(makeStdin({
rate_limits: {
five_hour: {
used_percentage: 140,
resets_at: 'not-a-date',
},
},
}));
expect(result).toEqual({
fiveHourPercent: 100,
weeklyPercent: undefined,
fiveHourResetsAt: null,
weeklyResetsAt: null,
});
});
});

View File

@@ -1,38 +1,50 @@
import { beforeEach, afterEach, describe, expect, it, vi } from 'vitest';
const fakeStdin = {
cwd: '/tmp/worktree',
transcript_path: '/tmp/worktree/transcript.jsonl',
model: { id: 'claude-test' },
context_window: {
used_percentage: 12,
current_usage: { input_tokens: 10, cache_creation_input_tokens: 0, cache_read_input_tokens: 0 },
context_window_size: 100,
},
};
function makeStdin(withRateLimits = false) {
return {
cwd: '/tmp/worktree',
transcript_path: '/tmp/worktree/transcript.jsonl',
model: { id: 'claude-test' },
context_window: {
used_percentage: 12,
current_usage: { input_tokens: 10, cache_creation_input_tokens: 0, cache_read_input_tokens: 0 },
context_window_size: 100,
},
...(withRateLimits
? {
rate_limits: {
five_hour: { used_percentage: 11, resets_at: 1776348000 },
seven_day: { used_percentage: 2, resets_at: 1776916800 },
},
}
: {}),
};
}
const fakeConfig = {
preset: 'focused',
elements: {
rateLimits: false,
apiKeySource: false,
safeMode: false,
missionBoard: false,
},
thresholds: {
contextWarning: 70,
contextCritical: 85,
},
staleTaskThresholdMinutes: 30,
contextLimitWarning: {
autoCompact: false,
threshold: 90,
},
missionBoard: {
enabled: false,
},
usageApiPollIntervalMs: 300000,
} as const;
function makeConfig(rateLimits = false) {
return {
preset: 'focused',
elements: {
rateLimits,
apiKeySource: false,
safeMode: false,
missionBoard: false,
},
thresholds: {
contextWarning: 70,
contextCritical: 85,
},
staleTaskThresholdMinutes: 30,
contextLimitWarning: {
autoCompact: false,
threshold: 90,
},
missionBoard: {
enabled: false,
},
usageApiPollIntervalMs: 300000,
} as const;
}
describe('HUD watch mode initialization', () => {
const originalIsTTY = Object.getOwnPropertyDescriptor(process.stdin, 'isTTY');
@@ -40,23 +52,47 @@ describe('HUD watch mode initialization', () => {
let readRalphStateForHud: ReturnType<typeof vi.fn>;
let readUltraworkStateForHud: ReturnType<typeof vi.fn>;
let readAutopilotStateForHud: ReturnType<typeof vi.fn>;
let getUsage: ReturnType<typeof vi.fn>;
let render: ReturnType<typeof vi.fn>;
let consoleLogSpy: ReturnType<typeof vi.spyOn>;
let consoleErrorSpy: ReturnType<typeof vi.spyOn>;
async function importHudModule() {
async function importHudModule(overrides: {
config?: ReturnType<typeof makeConfig>;
stdin?: ReturnType<typeof makeStdin>;
getUsageResult?: unknown;
} = {}) {
vi.resetModules();
const stdin = overrides.stdin ?? makeStdin();
const config = overrides.config ?? makeConfig();
initializeHUDState = vi.fn(async () => {});
readRalphStateForHud = vi.fn(() => null);
readUltraworkStateForHud = vi.fn(() => null);
readAutopilotStateForHud = vi.fn(() => null);
getUsage = vi.fn(async () => overrides.getUsageResult ?? null);
render = vi.fn(async () => '[HUD] ok');
vi.doMock('../../hud/stdin.js', () => ({
readStdin: vi.fn(async () => null),
writeStdinCache: vi.fn(),
readStdinCache: vi.fn(() => fakeStdin),
readStdinCache: vi.fn(() => stdin),
getContextPercent: vi.fn(() => 12),
getModelName: vi.fn(() => 'claude-test'),
getRateLimitsFromStdin: vi.fn((value) => {
const fiveHour = value.rate_limits?.five_hour?.used_percentage;
const sevenDay = value.rate_limits?.seven_day?.used_percentage;
if (fiveHour == null && sevenDay == null) {
return null;
}
return {
fiveHourPercent: fiveHour ?? 0,
weeklyPercent: sevenDay,
fiveHourResetsAt: fiveHour == null ? null : new Date(1776348000 * 1000),
weeklyResetsAt: sevenDay == null ? null : new Date(1776916800 * 1000),
};
}),
stabilizeContextPercent: vi.fn((value) => value),
}));
vi.doMock('../../hud/transcript.js', () => ({
@@ -75,7 +111,7 @@ describe('HUD watch mode initialization', () => {
vi.doMock('../../hud/state.js', () => ({
initializeHUDState,
readHudConfig: vi.fn(() => fakeConfig),
readHudConfig: vi.fn(() => config),
readHudState: vi.fn(() => null),
getRunningTasks: vi.fn(() => []),
writeHudState: vi.fn(() => true),
@@ -88,9 +124,9 @@ describe('HUD watch mode initialization', () => {
readAutopilotStateForHud,
}));
vi.doMock('../../hud/usage-api.js', () => ({ getUsage: vi.fn(async () => null) }));
vi.doMock('../../hud/usage-api.js', () => ({ getUsage }));
vi.doMock('../../hud/custom-rate-provider.js', () => ({ executeCustomProvider: vi.fn(async () => null) }));
vi.doMock('../../hud/render.js', () => ({ render: vi.fn(async () => '[HUD] ok') }));
vi.doMock('../../hud/render.js', () => ({ render }));
vi.doMock('../../hud/elements/api-key-source.js', () => ({ detectApiKeySource: vi.fn(() => null) }));
vi.doMock('../../hud/mission-board.js', () => ({ refreshMissionBoardState: vi.fn(async () => null) }));
vi.doMock('../../hud/sanitize.js', () => ({ sanitizeOutput: vi.fn((value: string) => value) }));
@@ -115,7 +151,6 @@ describe('HUD watch mode initialization', () => {
});
afterEach(() => {
fakeStdin.transcript_path = '/tmp/worktree/transcript.jsonl';
vi.resetModules();
vi.clearAllMocks();
vi.doUnmock('../../hud/stdin.js');
@@ -167,8 +202,9 @@ describe('HUD watch mode initialization', () => {
});
it('passes the current session id to OMC state readers', async () => {
const hud = await importHudModule();
fakeStdin.transcript_path = '/tmp/worktree/transcripts/123e4567-e89b-12d3-a456-426614174000.jsonl';
const stdin = makeStdin();
stdin.transcript_path = '/tmp/worktree/transcripts/123e4567-e89b-12d3-a456-426614174000.jsonl';
const hud = await importHudModule({ stdin });
await hud.main(true, false);
@@ -176,4 +212,36 @@ describe('HUD watch mode initialization', () => {
expect(readUltraworkStateForHud).toHaveBeenCalledWith('/tmp/worktree', '123e4567-e89b-12d3-a456-426614174000');
expect(readAutopilotStateForHud).toHaveBeenCalledWith('/tmp/worktree', '123e4567-e89b-12d3-a456-426614174000');
});
it('prefers stdin rate limits over the usage API when available', async () => {
const hud = await importHudModule({
config: makeConfig(true),
stdin: makeStdin(true),
});
await hud.main(true, false);
expect(getUsage).not.toHaveBeenCalled();
expect(render).toHaveBeenCalledWith(expect.objectContaining({
rateLimitsResult: {
rateLimits: {
fiveHourPercent: 11,
weeklyPercent: 2,
fiveHourResetsAt: new Date(1776348000 * 1000),
weeklyResetsAt: new Date(1776916800 * 1000),
},
},
}), expect.anything());
});
it('falls back to the usage API when stdin omits rate limits', async () => {
const hud = await importHudModule({
config: makeConfig(true),
getUsageResult: { rateLimits: { fiveHourPercent: 55, weeklyPercent: 10 } },
});
await hud.main(true, false);
expect(getUsage).toHaveBeenCalledTimes(1);
});
});

View File

@@ -12,6 +12,7 @@ import {
readStdinCache,
getContextPercent,
getModelName,
getRateLimitsFromStdin,
stabilizeContextPercent,
} from "./stdin.js";
import { parseTranscript } from "./transcript.js";
@@ -339,9 +340,14 @@ async function main(watchMode = false, skipInit = false): Promise<void> {
writeHudState(stateToWrite, cwd, currentSessionId ?? undefined);
}
// Fetch rate limits from OAuth API (if available)
// Prefer Claude Code stdin rate limits when available to avoid cold-start API fetches.
const stdinRateLimits = getRateLimitsFromStdin(stdin);
const rateLimitsResult =
config.elements.rateLimits !== false ? await getUsage() : null;
config.elements.rateLimits === false
? null
: stdinRateLimits
? { rateLimits: stdinRateLimits }
: await getUsage();
// Fetch custom rate limit buckets (if configured)
const customBuckets =

View File

@@ -8,7 +8,7 @@
import { existsSync, readFileSync, writeFileSync, mkdirSync } from 'fs';
import { join } from 'path';
import { getWorktreeRoot } from '../lib/worktree-paths.js';
import type { StatuslineStdin } from './types.js';
import type { RateLimits, StatuslineStdin } from './types.js';
const TRANSIENT_CONTEXT_PERCENT_TOLERANCE = 3;
@@ -92,6 +92,35 @@ function getCurrentUsage(stdin: StatuslineStdin) {
return stdin.context_window?.current_usage;
}
function clampPercent(value: number | undefined): number {
if (value == null || !isFinite(value)) {
return 0;
}
return Math.max(0, Math.min(100, value));
}
function parseResetDate(value: number | string | undefined): Date | null {
if (value == null) {
return null;
}
const numericValue = typeof value === 'number'
? value
: (typeof value === 'string' && value.trim() !== '' ? Number(value) : Number.NaN);
if (Number.isFinite(numericValue)) {
const millis = Math.abs(numericValue) < 1e12 ? numericValue * 1000 : numericValue;
const date = new Date(millis);
return Number.isNaN(date.getTime()) ? null : date;
}
if (typeof value === 'string') {
const date = new Date(value);
return Number.isNaN(date.getTime()) ? null : date;
}
return null;
}
/**
* Get total tokens from stdin context_window.current_usage
*/
@@ -179,6 +208,25 @@ export function getContextPercent(stdin: StatuslineStdin): number {
return getManualContextPercent(stdin) ?? 0;
}
/**
* Convert Claude Code stdin rate_limits into the existing HUD RateLimits shape.
*/
export function getRateLimitsFromStdin(stdin: StatuslineStdin): RateLimits | null {
const fiveHour = stdin.rate_limits?.five_hour?.used_percentage;
const sevenDay = stdin.rate_limits?.seven_day?.used_percentage;
if (fiveHour == null && sevenDay == null) {
return null;
}
return {
fiveHourPercent: clampPercent(fiveHour),
weeklyPercent: sevenDay == null ? undefined : clampPercent(sevenDay),
fiveHourResetsAt: parseResetDate(stdin.rate_limits?.five_hour?.resets_at),
weeklyResetsAt: parseResetDate(stdin.rate_limits?.seven_day?.resets_at),
};
}
/**
* Get model display name from stdin.
* Prefer the official display name field, then fall back to the raw model id.

View File

@@ -66,6 +66,18 @@ export interface StatuslineStdin {
cache_read_input_tokens?: number;
};
};
/** Rate limits from Claude Code statusline stdin */
rate_limits?: {
five_hour?: {
used_percentage?: number;
resets_at?: number | string;
};
seven_day?: {
used_percentage?: number;
resets_at?: number | string;
};
};
}
// ============================================================================