From f579dad768d004f02c2d00775042db1432ec616d Mon Sep 17 00:00:00 2001 From: ulinzeng Date: Mon, 20 Apr 2026 15:15:07 +0800 Subject: [PATCH] feat: add cost-estimate and session-bridge shared libs Extract estimateCost() into scripts/lib/cost-estimate.js for reuse across cost-tracker and ecc-metrics-bridge hooks. Add scripts/lib/session-bridge.js with atomic bridge file I/O, session ID sanitization, and path traversal prevention. --- scripts/lib/cost-estimate.js | 32 ++++++ scripts/lib/session-bridge.js | 81 ++++++++++++++ tests/lib/cost-estimate.test.js | 114 ++++++++++++++++++++ tests/lib/session-bridge.test.js | 174 +++++++++++++++++++++++++++++++ 4 files changed, 401 insertions(+) create mode 100644 scripts/lib/cost-estimate.js create mode 100644 scripts/lib/session-bridge.js create mode 100644 tests/lib/cost-estimate.test.js create mode 100644 tests/lib/session-bridge.test.js diff --git a/scripts/lib/cost-estimate.js b/scripts/lib/cost-estimate.js new file mode 100644 index 00000000..a1651a8c --- /dev/null +++ b/scripts/lib/cost-estimate.js @@ -0,0 +1,32 @@ +'use strict'; + +/** + * Shared cost estimation for ECC hooks. + * + * Approximate per-1M-token blended rates (conservative defaults). + */ + +const RATE_TABLE = { + haiku: { in: 0.8, out: 4.0 }, + sonnet: { in: 3.0, out: 15.0 }, + opus: { in: 15.0, out: 75.0 } +}; + +/** + * Estimate USD cost from token counts. + * @param {string} model - Model name (may contain "haiku", "sonnet", or "opus") + * @param {number} inputTokens + * @param {number} outputTokens + * @returns {number} Estimated cost in USD (rounded to 6 decimal places) + */ +function estimateCost(model, inputTokens, outputTokens) { + const normalized = String(model || '').toLowerCase(); + let rates = RATE_TABLE.sonnet; + if (normalized.includes('haiku')) rates = RATE_TABLE.haiku; + if (normalized.includes('opus')) rates = RATE_TABLE.opus; + + const cost = (inputTokens / 1_000_000) * rates.in + (outputTokens / 1_000_000) * rates.out; + return Math.round(cost * 1e6) / 1e6; +} + +module.exports = { estimateCost, RATE_TABLE }; diff --git a/scripts/lib/session-bridge.js b/scripts/lib/session-bridge.js new file mode 100644 index 00000000..aceae9cb --- /dev/null +++ b/scripts/lib/session-bridge.js @@ -0,0 +1,81 @@ +'use strict'; + +/** + * Shared session bridge utilities for ECC hooks. + * + * The bridge file is a small JSON aggregate in /tmp that allows + * statusline, metrics-bridge, and context-monitor to share state + * without scanning large JSONL logs on every invocation. + */ + +const fs = require('fs'); +const os = require('os'); +const path = require('path'); + +const MAX_SESSION_ID_LENGTH = 64; + +/** + * Sanitize a session ID for safe use in file paths. + * Rejects path traversal, strips unsafe chars, limits length. + * @param {string} raw + * @returns {string|null} Safe session ID or null if invalid + */ +function sanitizeSessionId(raw) { + if (!raw || typeof raw !== 'string') return null; + if (/[/\\]|\.\./.test(raw)) return null; + const safe = raw.replace(/[^a-zA-Z0-9_-]/g, '_').slice(0, MAX_SESSION_ID_LENGTH); + return safe || null; +} + +/** + * Get the bridge file path for a session. + * @param {string} sessionId - Already-sanitized session ID + * @returns {string} + */ +function getBridgePath(sessionId) { + return path.join(os.tmpdir(), `ecc-metrics-${sessionId}.json`); +} + +/** + * Read bridge data. Returns null on any error. + * @param {string} sessionId - Already-sanitized session ID + * @returns {object|null} + */ +function readBridge(sessionId) { + try { + const raw = fs.readFileSync(getBridgePath(sessionId), 'utf8'); + return JSON.parse(raw); + } catch { + return null; + } +} + +/** + * Write bridge data atomically (write .tmp then rename). + * @param {string} sessionId - Already-sanitized session ID + * @param {object} data + */ +function writeBridgeAtomic(sessionId, data) { + const target = getBridgePath(sessionId); + const tmp = `${target}.tmp`; + fs.writeFileSync(tmp, JSON.stringify(data), 'utf8'); + fs.renameSync(tmp, target); +} + +/** + * Resolve session ID from environment variables. + * @returns {string|null} Sanitized session ID or null + */ +function resolveSessionId() { + const raw = process.env.ECC_SESSION_ID || process.env.CLAUDE_SESSION_ID || ''; + return sanitizeSessionId(raw); +} + +module.exports = { + sanitizeSessionId, + getBridgePath, + readBridge, + writeBridgeAtomic, + resolveSessionId, + MAX_SESSION_ID_LENGTH +}; diff --git a/tests/lib/cost-estimate.test.js b/tests/lib/cost-estimate.test.js new file mode 100644 index 00000000..bcb5906b --- /dev/null +++ b/tests/lib/cost-estimate.test.js @@ -0,0 +1,114 @@ +/** + * Tests for scripts/lib/cost-estimate.js + * + * Run with: node tests/lib/cost-estimate.test.js + */ + +const assert = require('assert'); + +const { estimateCost, RATE_TABLE } = require('../../scripts/lib/cost-estimate'); + +// Test helper +function test(name, fn) { + try { + fn(); + console.log(` \u2713 ${name}`); + return true; + } catch (err) { + console.log(` \u2717 ${name}`); + console.log(` Error: ${err.message}`); + return false; + } +} + +function runTests() { + console.log('\n=== Testing cost-estimate.js ===\n'); + + let passed = 0; + let failed = 0; + + // RATE_TABLE structure + console.log('RATE_TABLE:'); + + if ( + test('RATE_TABLE has haiku, sonnet, opus keys', () => { + assert.ok(RATE_TABLE.haiku, 'Missing haiku'); + assert.ok(RATE_TABLE.sonnet, 'Missing sonnet'); + assert.ok(RATE_TABLE.opus, 'Missing opus'); + assert.strictEqual(typeof RATE_TABLE.haiku.in, 'number'); + assert.strictEqual(typeof RATE_TABLE.haiku.out, 'number'); + assert.strictEqual(typeof RATE_TABLE.sonnet.in, 'number'); + assert.strictEqual(typeof RATE_TABLE.sonnet.out, 'number'); + assert.strictEqual(typeof RATE_TABLE.opus.in, 'number'); + assert.strictEqual(typeof RATE_TABLE.opus.out, 'number'); + }) + ) + passed++; + else failed++; + + // estimateCost tests + console.log('\nestimateCost:'); + + if ( + test('opus 1M/1M tokens returns 90', () => { + const cost = estimateCost('opus', 1_000_000, 1_000_000); + assert.strictEqual(cost, 90); + }) + ) + passed++; + else failed++; + + if ( + test('sonnet 1M/1M tokens returns 18', () => { + const cost = estimateCost('sonnet', 1_000_000, 1_000_000); + assert.strictEqual(cost, 18); + }) + ) + passed++; + else failed++; + + if ( + test('haiku 1M/1M tokens returns 4.8', () => { + const cost = estimateCost('haiku', 1_000_000, 1_000_000); + assert.strictEqual(cost, 4.8); + }) + ) + passed++; + else failed++; + + if ( + test('null model with 0 tokens returns 0', () => { + const cost = estimateCost(null, 0, 0); + assert.strictEqual(cost, 0); + }) + ) + passed++; + else failed++; + + if ( + test('full model name claude-opus-4-6 uses opus rates', () => { + const cost = estimateCost('claude-opus-4-6', 500, 200); + // (500 / 1_000_000) * 15 + (200 / 1_000_000) * 75 = 0.0075 + 0.015 = 0.0225 + const expected = Math.round(0.0225 * 1e6) / 1e6; + assert.strictEqual(cost, expected); + }) + ) + passed++; + else failed++; + + if ( + test('unknown model falls back to sonnet rates', () => { + const cost = estimateCost('unknown-model', 1_000_000, 1_000_000); + assert.strictEqual(cost, 18); + }) + ) + passed++; + else failed++; + + // Summary + console.log(`\nResults: ${passed} passed, ${failed} failed\n`); + return { passed, failed }; +} + +const { failed } = runTests(); +process.exit(failed > 0 ? 1 : 0); diff --git a/tests/lib/session-bridge.test.js b/tests/lib/session-bridge.test.js new file mode 100644 index 00000000..60841c9b --- /dev/null +++ b/tests/lib/session-bridge.test.js @@ -0,0 +1,174 @@ +/** + * Tests for scripts/lib/session-bridge.js + * + * Run with: node tests/lib/session-bridge.test.js + */ + +const assert = require('assert'); +const fs = require('fs'); + +const { sanitizeSessionId, getBridgePath, readBridge, writeBridgeAtomic, resolveSessionId, MAX_SESSION_ID_LENGTH } = require('../../scripts/lib/session-bridge'); + +// Test helper +function test(name, fn) { + try { + fn(); + console.log(` \u2713 ${name}`); + return true; + } catch (err) { + console.log(` \u2717 ${name}`); + console.log(` Error: ${err.message}`); + return false; + } +} + +function runTests() { + console.log('\n=== Testing session-bridge.js ===\n'); + + let passed = 0; + let failed = 0; + + // sanitizeSessionId tests + console.log('sanitizeSessionId:'); + + if ( + test('valid ID passes through', () => { + assert.strictEqual(sanitizeSessionId('abc-123'), 'abc-123'); + }) + ) + passed++; + else failed++; + + if ( + test('path traversal returns null', () => { + assert.strictEqual(sanitizeSessionId('../etc/passwd'), null); + }) + ) + passed++; + else failed++; + + if ( + test('forward slash returns null', () => { + assert.strictEqual(sanitizeSessionId('/tmp/evil'), null); + }) + ) + passed++; + else failed++; + + if ( + test('backslash returns null', () => { + assert.strictEqual(sanitizeSessionId('a\\b'), null); + }) + ) + passed++; + else failed++; + + if ( + test('null input returns null', () => { + assert.strictEqual(sanitizeSessionId(null), null); + }) + ) + passed++; + else failed++; + + if ( + test('empty string returns null', () => { + assert.strictEqual(sanitizeSessionId(''), null); + }) + ) + passed++; + else failed++; + + if ( + test('long string is truncated to MAX_SESSION_ID_LENGTH', () => { + const longId = 'a'.repeat(100); + const result = sanitizeSessionId(longId); + assert.ok(result, 'Should not return null for valid chars'); + assert.strictEqual(result.length, MAX_SESSION_ID_LENGTH); + }) + ) + passed++; + else failed++; + + // getBridgePath tests + console.log('\ngetBridgePath:'); + + if ( + test('returns path containing ecc-metrics-', () => { + const p = getBridgePath('test-session'); + assert.ok(p.includes('ecc-metrics-'), `Expected ecc-metrics- in path, got: ${p}`); + }) + ) + passed++; + else failed++; + + // writeBridgeAtomic + readBridge roundtrip + console.log('\nwriteBridgeAtomic / readBridge:'); + + if ( + test('roundtrip write then read returns same data', () => { + const testId = `test-bridge-${Date.now()}`; + const data = { session_id: testId, tool_count: 42 }; + try { + writeBridgeAtomic(testId, data); + const result = readBridge(testId); + assert.deepStrictEqual(result, data); + } finally { + // Clean up + try { + fs.unlinkSync(getBridgePath(testId)); + } catch { + /* ignore */ + } + } + }) + ) + passed++; + else failed++; + + if ( + test('readBridge with nonexistent session returns null', () => { + const result = readBridge('nonexistent-session-id-999'); + assert.strictEqual(result, null); + }) + ) + passed++; + else failed++; + + // resolveSessionId tests + console.log('\nresolveSessionId:'); + + if ( + test('resolveSessionId uses ECC_SESSION_ID env var', () => { + const original = process.env.ECC_SESSION_ID; + try { + process.env.ECC_SESSION_ID = 'env-session-42'; + const result = resolveSessionId(); + assert.strictEqual(result, 'env-session-42'); + } finally { + if (original === undefined) { + delete process.env.ECC_SESSION_ID; + } else { + process.env.ECC_SESSION_ID = original; + } + } + }) + ) + passed++; + else failed++; + + if ( + test('MAX_SESSION_ID_LENGTH is 64', () => { + assert.strictEqual(MAX_SESSION_ID_LENGTH, 64); + }) + ) + passed++; + else failed++; + + // Summary + console.log(`\nResults: ${passed} passed, ${failed} failed\n`); + return { passed, failed }; +} + +const { failed } = runTests(); +process.exit(failed > 0 ? 1 : 0);