mirror of
https://fastgit.cc/github.com/Yeachan-Heo/oh-my-claudecode
synced 2026-04-20 21:00:50 +08:00
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:
@@ -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,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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 =
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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;
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
|
||||
Reference in New Issue
Block a user