332 lines
13 KiB
JavaScript
332 lines
13 KiB
JavaScript
import { test } from 'node:test';
|
|
import assert from 'node:assert/strict';
|
|
import { mkdtemp, readdir, rm, stat, writeFile } from 'node:fs/promises';
|
|
import { tmpdir } from 'node:os';
|
|
import path from 'node:path';
|
|
import { getOutputSpeed } from '../dist/speed-tracker.js';
|
|
import { existsSync } from 'node:fs';
|
|
|
|
function restoreEnvVar(name, value) {
|
|
if (value === undefined) {
|
|
delete process.env[name];
|
|
return;
|
|
}
|
|
process.env[name] = value;
|
|
}
|
|
|
|
async function createTempHome() {
|
|
return await mkdtemp(path.join(tmpdir(), 'claude-hud-speed-'));
|
|
}
|
|
|
|
async function createTranscript(tempHome, name = 'session.jsonl') {
|
|
const transcriptPath = path.join(tempHome, name);
|
|
await writeFile(transcriptPath, '', 'utf8');
|
|
return transcriptPath;
|
|
}
|
|
|
|
function stdinWith(transcriptPath, outputTokens) {
|
|
return {
|
|
transcript_path: transcriptPath,
|
|
context_window: { current_usage: { output_tokens: outputTokens } },
|
|
};
|
|
}
|
|
|
|
test('getOutputSpeed returns null when output tokens are missing', () => {
|
|
const speed = getOutputSpeed({
|
|
transcript_path: '/tmp/claude-hud-speed-missing.jsonl',
|
|
context_window: { current_usage: { input_tokens: 10 } },
|
|
});
|
|
assert.equal(speed, null);
|
|
});
|
|
|
|
test('getOutputSpeed estimates speed from transcript growth when output tokens are missing', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
const first = getOutputSpeed(
|
|
{ transcript_path: transcriptPath, context_window: { current_usage: { input_tokens: 10 } } },
|
|
{ ...base, now: () => 1000 }
|
|
);
|
|
assert.equal(first, null);
|
|
|
|
await writeFile(transcriptPath, 'x'.repeat(40), 'utf8');
|
|
const second = getOutputSpeed(
|
|
{ transcript_path: transcriptPath, context_window: { current_usage: { input_tokens: 10 } } },
|
|
{ ...base, now: () => 1500 }
|
|
);
|
|
assert.ok(second !== null);
|
|
assert.ok(Math.abs(second - 20) < 0.01);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed accumulates short transcript-growth windows until the fallback sample matures', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
const stdin = { transcript_path: transcriptPath, context_window: { current_usage: { input_tokens: 10 } } };
|
|
|
|
getOutputSpeed(stdin, { ...base, now: () => 1000 });
|
|
|
|
await writeFile(transcriptPath, 'x'.repeat(40), 'utf8');
|
|
assert.equal(getOutputSpeed(stdin, { ...base, now: () => 1200 }), null);
|
|
|
|
await writeFile(transcriptPath, 'x'.repeat(80), 'utf8');
|
|
assert.equal(getOutputSpeed(stdin, { ...base, now: () => 1400 }), null);
|
|
|
|
await writeFile(transcriptPath, 'x'.repeat(120), 'utf8');
|
|
const matured = getOutputSpeed(stdin, { ...base, now: () => 1600 });
|
|
assert.ok(matured !== null);
|
|
assert.ok(Math.abs(matured - 50) < 0.01);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed ignores fallback samples without transcript growth', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
const stdin = { transcript_path: transcriptPath, context_window: { current_usage: { input_tokens: 10 } } };
|
|
|
|
getOutputSpeed(stdin, { ...base, now: () => 1000 });
|
|
const idle = getOutputSpeed(stdin, { ...base, now: () => 1500 });
|
|
assert.equal(idle, null);
|
|
|
|
await writeFile(transcriptPath, 'x'.repeat(40), 'utf8');
|
|
const afterIdle = getOutputSpeed(stdin, { ...base, now: () => 2000 });
|
|
assert.ok(afterIdle !== null);
|
|
assert.ok(Math.abs(afterIdle - 20) < 0.01);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed ignores stale fallback windows', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
const stdin = { transcript_path: transcriptPath, context_window: { current_usage: { input_tokens: 10 } } };
|
|
|
|
getOutputSpeed(stdin, { ...base, now: () => 1000 });
|
|
|
|
await writeFile(transcriptPath, 'x'.repeat(40), 'utf8');
|
|
const stale = getOutputSpeed(stdin, { ...base, now: () => 8000 });
|
|
assert.equal(stale, null);
|
|
|
|
await writeFile(transcriptPath, 'x'.repeat(80), 'utf8');
|
|
const fresh = getOutputSpeed(stdin, { ...base, now: () => 8500 });
|
|
assert.ok(fresh !== null);
|
|
assert.ok(Math.abs(fresh - 20) < 0.01);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed creates fallback cache under CLAUDE_CONFIG_DIR by default', async () => {
|
|
const tempHome = await createTempHome();
|
|
const customConfigDir = path.join(tempHome, '.claude-alt');
|
|
const originalHome = process.env.HOME;
|
|
const originalConfigDir = process.env.CLAUDE_CONFIG_DIR;
|
|
process.env.HOME = tempHome;
|
|
process.env.CLAUDE_CONFIG_DIR = customConfigDir;
|
|
|
|
try {
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
const speed = getOutputSpeed(
|
|
{ transcript_path: transcriptPath, context_window: { current_usage: { input_tokens: 10 } } },
|
|
{ now: () => 1000 }
|
|
);
|
|
assert.equal(speed, null);
|
|
|
|
const customCacheDir = path.join(customConfigDir, 'plugins', 'claude-hud', 'speed-cache');
|
|
const defaultCacheDir = path.join(tempHome, '.claude', 'plugins', 'claude-hud', 'speed-cache');
|
|
assert.equal(existsSync(customCacheDir), true);
|
|
assert.equal(existsSync(defaultCacheDir), false);
|
|
assert.equal((await stat(customCacheDir)).mode & 0o777, 0o700);
|
|
|
|
const cacheFiles = await readdir(customCacheDir);
|
|
const fileSizeCache = cacheFiles.find((name) => name.endsWith('.json.fs'));
|
|
assert.ok(fileSizeCache, `expected file-size cache in ${cacheFiles.join(', ')}`);
|
|
assert.equal((await stat(path.join(customCacheDir, fileSizeCache))).mode & 0o777, 0o600);
|
|
} finally {
|
|
restoreEnvVar('HOME', originalHome);
|
|
restoreEnvVar('CLAUDE_CONFIG_DIR', originalConfigDir);
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed ignores fallback tracking for non-file transcript paths', async () => {
|
|
const tempHome = await createTempHome();
|
|
|
|
try {
|
|
const speed = getOutputSpeed(
|
|
{ transcript_path: tempHome, context_window: { current_usage: { input_tokens: 10 } } },
|
|
{ homeDir: () => tempHome, now: () => 1000 }
|
|
);
|
|
assert.equal(speed, null);
|
|
|
|
const cacheDir = path.join(tempHome, '.claude', 'plugins', 'claude-hud', 'speed-cache');
|
|
assert.equal(existsSync(cacheDir), false);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed returns null when transcript_path is missing', async () => {
|
|
const tempHome = await createTempHome();
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
const speed = getOutputSpeed(
|
|
{ context_window: { current_usage: { output_tokens: 10 } } },
|
|
{ ...base, now: () => 1000 }
|
|
);
|
|
assert.equal(speed, null);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed computes tokens per second within window', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
const first = getOutputSpeed(stdinWith(transcriptPath, 10), { ...base, now: () => 1000 });
|
|
assert.equal(first, null);
|
|
|
|
const second = getOutputSpeed(stdinWith(transcriptPath, 20), { ...base, now: () => 1500 });
|
|
assert.ok(second !== null);
|
|
assert.ok(Math.abs(second - 20) < 0.01);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed ignores sub-window bursts to avoid inflated rates', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
getOutputSpeed(stdinWith(transcriptPath, 10), { ...base, now: () => 1000 });
|
|
|
|
// Status line re-renders ~50ms later with 60 more tokens. A naive rate
|
|
// calculation would report 1200 tok/s; we expect null instead (#481).
|
|
const speed = getOutputSpeed(stdinWith(transcriptPath, 70), { ...base, now: () => 1050 });
|
|
assert.equal(speed, null);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed accumulates repeated short windows until the sample matures', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
getOutputSpeed(stdinWith(transcriptPath, 10), { ...base, now: () => 1000 });
|
|
|
|
const firstBurst = getOutputSpeed(stdinWith(transcriptPath, 40), { ...base, now: () => 1200 });
|
|
assert.equal(firstBurst, null);
|
|
|
|
const secondBurst = getOutputSpeed(stdinWith(transcriptPath, 70), { ...base, now: () => 1400 });
|
|
assert.equal(secondBurst, null);
|
|
|
|
const matured = getOutputSpeed(stdinWith(transcriptPath, 100), { ...base, now: () => 1600 });
|
|
assert.ok(matured !== null);
|
|
assert.ok(Math.abs(matured - 150) < 0.01);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed ignores stale windows', async () => {
|
|
const tempHome = await createTempHome();
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
getOutputSpeed(stdinWith(transcriptPath, 10), { ...base, now: () => 1000 });
|
|
|
|
const speed = getOutputSpeed(stdinWith(transcriptPath, 30), { ...base, now: () => 8000 });
|
|
assert.equal(speed, null);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed isolates cache across concurrent sessions', async () => {
|
|
const tempHome = await createTempHome();
|
|
const sessionA = await createTranscript(tempHome, 'session-a.jsonl');
|
|
const sessionB = await createTranscript(tempHome, 'session-b.jsonl');
|
|
|
|
try {
|
|
const base = { homeDir: () => tempHome };
|
|
|
|
// Session A streams: seeds its cache, then reports a speed on the next tick.
|
|
getOutputSpeed(stdinWith(sessionA, 100), { ...base, now: () => 1000 });
|
|
const aSpeed = getOutputSpeed(stdinWith(sessionA, 200), { ...base, now: () => 1500 });
|
|
assert.ok(aSpeed !== null);
|
|
|
|
// Session B is idle with a much smaller counter. Before the fix, B would
|
|
// read A's cache entry as its `previous`, compute a bogus speed, or reset
|
|
// A's cache to B's value and poison subsequent A readings. With per-session
|
|
// caches the first B tick must seed a fresh cache and return null.
|
|
const bFirst = getOutputSpeed(stdinWith(sessionB, 5), { ...base, now: () => 1600 });
|
|
assert.equal(bFirst, null);
|
|
|
|
// Session A's cache must survive B's tick and keep producing stable speeds.
|
|
const aContinued = getOutputSpeed(stdinWith(sessionA, 300), { ...base, now: () => 2000 });
|
|
assert.ok(aContinued !== null);
|
|
assert.ok(Math.abs(aContinued - 200) < 0.01);
|
|
} finally {
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|
|
|
|
test('getOutputSpeed writes cache under CLAUDE_CONFIG_DIR by default', async () => {
|
|
const tempHome = await createTempHome();
|
|
const customConfigDir = path.join(tempHome, '.claude-alt');
|
|
const originalHome = process.env.HOME;
|
|
const originalConfigDir = process.env.CLAUDE_CONFIG_DIR;
|
|
process.env.HOME = tempHome;
|
|
process.env.CLAUDE_CONFIG_DIR = customConfigDir;
|
|
|
|
try {
|
|
const transcriptPath = await createTranscript(tempHome);
|
|
const first = getOutputSpeed(stdinWith(transcriptPath, 10), { now: () => 1000 });
|
|
assert.equal(first, null);
|
|
|
|
const second = getOutputSpeed(stdinWith(transcriptPath, 20), { now: () => 1500 });
|
|
assert.ok(second !== null);
|
|
|
|
const customCacheDir = path.join(customConfigDir, 'plugins', 'claude-hud', 'speed-cache');
|
|
const defaultCacheDir = path.join(tempHome, '.claude', 'plugins', 'claude-hud', 'speed-cache');
|
|
assert.equal(existsSync(customCacheDir), true);
|
|
assert.equal(existsSync(defaultCacheDir), false);
|
|
assert.equal((await stat(customCacheDir)).mode & 0o777, 0o700);
|
|
|
|
const cacheFiles = await readdir(customCacheDir);
|
|
const speedCache = cacheFiles.find((name) => name.endsWith('.json'));
|
|
assert.ok(speedCache, `expected speed cache in ${cacheFiles.join(', ')}`);
|
|
assert.equal((await stat(path.join(customCacheDir, speedCache))).mode & 0o777, 0o600);
|
|
} finally {
|
|
restoreEnvVar('HOME', originalHome);
|
|
restoreEnvVar('CLAUDE_CONFIG_DIR', originalConfigDir);
|
|
await rm(tempHome, { recursive: true, force: true });
|
|
}
|
|
});
|