@traice/codex-collector 0.1.0

package/install.mjs

@@ -0,0 +1,469 @@
+#!/usr/bin/env node
+import { createHash, randomBytes } from 'node:crypto';
+import { execFileSync, spawnSync } from 'node:child_process';
+import { existsSync, mkdirSync, readFileSync, writeFileSync, chmodSync, copyFileSync } from 'node:fs';
+import { homedir, hostname, userInfo } from 'node:os';
+import { resolve } from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+const DEFAULT_OUT_DIR = '~/.traice/internal-spend-codex';
+const LEGACY_INSTALL_CONFIG = '.traice-internal-spend/install.json';
+const DEFAULT_WORKSPACE_SLUG = 'internal-spend-poc';
+const DEFAULT_WORKSPACE_NAME = 'Internal Spend POC';
+const DEFAULT_KEY_NAME = 'internal-spend-codex-device';
+const BEGIN_MARKER = '# BEGIN trAIce Internal Spend Codex OTel';
+const END_MARKER = '# END trAIce Internal Spend Codex OTel';
+const LEGACY_BEGIN_MARKER = '# BEGIN trAIce Codex OTel POC';
+const LEGACY_END_MARKER = '# END trAIce Codex OTel POC';
+const LAUNCH_AGENT_LABEL = 'com.traice.internal-spend-codex-collector';
+
+const args = parseArgs(process.argv.slice(2));
+loadEnvFile(resolve('apps/web/.env'));
+
+const outDir = resolveHome(args.get('out') ?? DEFAULT_OUT_DIR);
+const installPath = resolve(outDir, 'install.json');
+const serverUrl = normalizeUrl(
+  args.get('server-url') ??
+    process.env.TRAICE_API_URL ??
+    process.env.NEXT_PUBLIC_APP_URL ??
+    'http://127.0.0.1:3003',
+);
+const workspaceSlug = args.get('workspace-slug') ?? DEFAULT_WORKSPACE_SLUG;
+const workspaceName = args.get('workspace-name') ?? DEFAULT_WORKSPACE_NAME;
+const keyName = args.get('key-name') ?? DEFAULT_KEY_NAME;
+const codexHome = resolveHome(args.get('codex-home') ?? process.env.CODEX_HOME ?? '~/.codex');
+const listenHost = args.get('listen-host') ?? '127.0.0.1';
+const listenPort = parsePort(args.get('listen-port') ?? '4318');
+const patchCodexConfig = args.get('patch-codex-config') !== 'false';
+const installLaunchAgent = args.get('launch-agent') === 'true' || args.get('start-service') === 'true';
+const employeeEmail = args.get('employee-email') ?? process.env.TRAICE_EMPLOYEE_EMAIL ?? gitConfig('user.email') ?? `${userInfo().username}@local.dev`;
+const employeeName = args.get('employee-name') ?? process.env.TRAICE_EMPLOYEE_NAME ?? gitConfig('user.name') ?? userInfo().username;
+const teamName = args.get('team-name') ?? process.env.TRAICE_TEAM_NAME ?? 'Engineering';
+const sourcePrincipal = args.get('source-principal') ?? `${hostname()}:${userInfo().username}`;
+const seatMonthlyUsd = parseOptionalMoney(args.get('seat-monthly-usd') ?? process.env.TRAICE_SEAT_MONTHLY_USD);
+const providedApiKey = readApiKey();
+
+mkdirSync(outDir, { recursive: true });
+
+const previousInstall = readJsonIfExists(installPath) ?? readLegacyInstallIfDefaultOut();
+const previousApiKey = typeof previousInstall?.apiKey === 'string' ? previousInstall.apiKey : null;
+
+let prisma = null;
+
+try {
+  const setup = providedApiKey ? providedApiKeySetup(providedApiKey) : await localDevSetup();
+
+  const installConfig = {
+    version: 1,
+    installedAt: new Date().toISOString(),
+    serverUrl,
+    apiKey: setup.key.raw,
+    apiKeyPrefix: setup.key.prefix,
+    workspaceId: setup.workspace.id,
+    workspaceSlug: setup.workspace.slug,
+    sourceKey: 'codex-local',
+    sourceName: 'Codex local collector',
+    sourceKind: 'codex_otel',
+    tool: 'codex',
+    category: 'coding_agent',
+    employeeEmail,
+    employeeName,
+    teamName,
+    sourcePrincipal,
+    seatMonthlyUsd,
+    codexHome,
+    listenHost,
+    listenPort,
+    pollIntervalMs: 3000,
+    backfillHours: 6,
+    includeZeroOutput: false,
+  };
+
+  writeFileSync(installPath, `${JSON.stringify(installConfig, null, 2)}\n`, { mode: 0o600 });
+  try {
+    chmodSync(installPath, 0o600);
+  } catch {
+    // Best effort on non-POSIX filesystems.
+  }
+
+  const configResult = patchCodexConfig
+    ? patchUserCodexConfig({ codexHome, listenHost, listenPort })
+    : { status: 'skipped', path: resolve(codexHome, 'config.toml') };
+  const launchAgentResult = installLaunchAgent
+    ? installUserLaunchAgent({ installPath, outDir })
+    : { status: 'skipped' };
+
+  console.log(
+    JSON.stringify(
+      {
+        ok: true,
+        mode: setup.mode,
+        workspace: setup.workspace,
+        apiKey: { prefix: setup.key.prefix, reused: setup.key.reused },
+        identity: { employeeEmail, employeeName, teamName, sourcePrincipal },
+        cost: { seatMonthlyUsd },
+        installConfig: installPath,
+        codexConfig: configResult,
+        launchAgent: launchAgentResult,
+        collectorCommand: `npx @traice/codex-collector@latest collect --config ${installPath}`,
+      },
+      null,
+      2,
+    ),
+  );
+} finally {
+  await prisma?.$disconnect();
+}
+
+function providedApiKeySetup(rawApiKey) {
+  return {
+    mode: 'provided_api_key',
+    workspace: { id: null, slug: workspaceSlug, name: workspaceName, plan: null },
+    key: {
+      raw: rawApiKey,
+      prefix: rawApiKey.slice(0, 12),
+      reused: true,
+    },
+  };
+}
+
+function readLegacyInstallIfDefaultOut() {
+  if (args.get('out')) return null;
+  return readJsonIfExists(resolve(LEGACY_INSTALL_CONFIG));
+}
+
+async function localDevSetup() {
+  if (!process.env.POSTGRES_PRISMA_URL || !process.env.POSTGRES_URL_NON_POOLING) {
+    console.error([
+      'Missing API key and local database environment.',
+      'For a normal install, set TRAICE_API_KEY or pass --api-key-stdin.',
+      'For local repo development, check apps/web/.env for POSTGRES_PRISMA_URL and POSTGRES_URL_NON_POOLING.',
+    ].join('\n'));
+    process.exit(2);
+  }
+
+  const { PrismaClient } = await import('@prisma/client');
+  prisma = new PrismaClient();
+  const workspace = await prisma.workspace.upsert({
+    where: { slug: workspaceSlug },
+    create: {
+      name: workspaceName,
+      slug: workspaceSlug,
+      plan: 'TEAM',
+      captureSamplesOptIn: false,
+      internalSpendEnabled: true,
+    },
+    update: {
+      name: workspaceName,
+      internalSpendEnabled: true,
+    },
+    select: { id: true, slug: true, name: true, plan: true },
+  });
+
+  const key = await ensureApiKey(workspace.id, keyName, previousApiKey);
+  await ensureInternalSource(workspace.id, seatMonthlyUsd);
+  if (process.env.DEV_AUTH_ENABLED === 'true') {
+    await ensureDevMembership(workspace.id, args.get('dev-user') ?? 'admin@local.dev');
+  }
+
+  return {
+    mode: 'local_dev_bootstrap',
+    workspace: { id: workspace.id, slug: workspace.slug, name: workspace.name, plan: workspace.plan },
+    key,
+  };
+}
+
+async function ensureDevMembership(workspaceId, email) {
+  const user = await prisma.user.upsert({
+    where: { email },
+    create: {
+      email,
+      name: email.split('@')[0],
+      emailVerified: new Date(),
+    },
+    update: {
+      emailVerified: new Date(),
+    },
+  });
+
+  await prisma.workspaceUser.upsert({
+    where: { workspaceId_userId: { workspaceId, userId: user.id } },
+    create: { workspaceId, userId: user.id, role: 'OWNER' },
+    update: { role: 'OWNER' },
+  });
+}
+
+async function ensureInternalSource(workspaceId, monthlySeatUsd) {
+  const metadata = monthlySeatUsd > 0
+    ? {
+        subscription: {
+          monthlySeatUsd,
+          allocation: 'monthly_seat_commitment',
+          note: 'Subscription commitment is shown separately from per-request usage cost.',
+        },
+      }
+    : undefined;
+
+  await prisma.internalSource.upsert({
+    where: { workspaceId_key: { workspaceId, key: 'codex-local' } },
+    create: {
+      workspaceId,
+      key: 'codex-local',
+      name: 'Codex local collector',
+      kind: 'codex_otel',
+      ...(metadata ? { metadata } : {}),
+    },
+    update: {
+      name: 'Codex local collector',
+      kind: 'codex_otel',
+      status: 'active',
+      ...(metadata ? { metadata } : {}),
+    },
+  });
+}
+
+async function ensureApiKey(workspaceId, name, previousRawKey) {
+  if (previousRawKey) {
+    const existing = await prisma.apiKey.findUnique({
+      where: { hash: hashApiKey(previousRawKey) },
+      select: { id: true, prefix: true, revokedAt: true, workspaceId: true },
+    });
+    if (existing && !existing.revokedAt && existing.workspaceId === workspaceId) {
+      return { raw: previousRawKey, prefix: existing.prefix, reused: true };
+    }
+  }
+
+  const generated = generateApiKey();
+  await prisma.apiKey.create({
+    data: {
+      workspaceId,
+      name,
+      prefix: generated.prefix,
+      hash: generated.hash,
+    },
+  });
+  return { raw: generated.raw, prefix: generated.prefix, reused: false };
+}
+
+function patchUserCodexConfig({ codexHome, listenHost, listenPort }) {
+  mkdirSync(codexHome, { recursive: true });
+  const configPath = resolve(codexHome, 'config.toml');
+  const current = existsSync(configPath) ? readFileSync(configPath, 'utf8') : '';
+  const block = [
+    BEGIN_MARKER,
+    '[otel]',
+    'environment = "traice-device"',
+    'log_user_prompt = false',
+    `exporter = { otlp-http = { endpoint = "http://${listenHost}:${listenPort}/v1/logs", protocol = "json" } }`,
+    'trace_exporter = "none"',
+    'metrics_exporter = "none"',
+    END_MARKER,
+    '',
+  ].join('\n');
+
+  if (current.includes(BEGIN_MARKER) || current.includes(LEGACY_BEGIN_MARKER)) {
+    const begin = current.includes(BEGIN_MARKER) ? BEGIN_MARKER : LEGACY_BEGIN_MARKER;
+    const end = current.includes(BEGIN_MARKER) ? END_MARKER : LEGACY_END_MARKER;
+    const pattern = new RegExp(`${escapeRegExp(begin)}[\\s\\S]*?${escapeRegExp(end)}\\n?`, 'm');
+    writeFileSync(configPath, current.replace(pattern, block));
+    return { status: 'updated', path: configPath };
+  }
+
+  if (/^\s*\[otel\]\s*$/m.test(current)) {
+    return {
+      status: 'blocked_existing_otel',
+      path: configPath,
+      message: 'Existing [otel] table found; left it unchanged.',
+    };
+  }
+
+  const next = current.trimEnd().length > 0 ? `${current.trimEnd()}\n\n${block}` : block;
+  writeFileSync(configPath, next);
+  return { status: 'added', path: configPath };
+}
+
+function installUserLaunchAgent({ installPath, outDir }) {
+  const launchAgentsDir = resolveHome('~/Library/LaunchAgents');
+  mkdirSync(launchAgentsDir, { recursive: true });
+  const plistPath = resolve(launchAgentsDir, `${LAUNCH_AGENT_LABEL}.plist`);
+  const nodePath = process.execPath;
+  const collectorPath = copyCollectorForService(outDir);
+  const logPath = resolve(outDir, 'collector.log');
+  const errPath = resolve(outDir, 'collector.err.log');
+  const plist = `<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+  <key>Label</key>
+  <string>${escapeXml(LAUNCH_AGENT_LABEL)}</string>
+  <key>ProgramArguments</key>
+  <array>
+    <string>${escapeXml(nodePath)}</string>
+    <string>${escapeXml(collectorPath)}</string>
+    <string>--config</string>
+    <string>${escapeXml(installPath)}</string>
+    <string>--out</string>
+    <string>${escapeXml(outDir)}</string>
+  </array>
+  <key>WorkingDirectory</key>
+  <string>${escapeXml(outDir)}</string>
+  <key>RunAtLoad</key>
+  <true/>
+  <key>KeepAlive</key>
+  <true/>
+  <key>StandardOutPath</key>
+  <string>${escapeXml(logPath)}</string>
+  <key>StandardErrorPath</key>
+  <string>${escapeXml(errPath)}</string>
+</dict>
+</plist>
+`;
+  writeFileSync(plistPath, plist);
+
+  const guiTarget = `gui/${process.getuid()}`;
+  spawnSync('launchctl', ['bootout', guiTarget, plistPath], { stdio: 'ignore' });
+  const result = spawnSync('launchctl', ['bootstrap', guiTarget, plistPath], { encoding: 'utf8' });
+  if (result.status !== 0) {
+    return {
+      status: 'written_not_loaded',
+      path: plistPath,
+      error: (result.stderr || result.stdout || '').trim(),
+    };
+  }
+  return { status: 'loaded', path: plistPath, logPath, errPath };
+}
+
+function copyCollectorForService(outDir) {
+  const sourcePath = fileURLToPath(new URL('./collector.mjs', import.meta.url));
+  const targetPath = resolve(outDir, 'collector.mjs');
+  copyFileSync(sourcePath, targetPath);
+  try {
+    chmodSync(targetPath, 0o755);
+  } catch {
+    // Best effort on non-POSIX filesystems.
+  }
+  return targetPath;
+}
+
+function readApiKey() {
+  if (args.get('api-key-stdin') === 'true') {
+    return readFileSync(0, 'utf8').trim();
+  }
+  return args.get('api-key') ?? process.env.TRAICE_API_KEY;
+}
+
+function parseArgs(argv) {
+  const parsed = new Map();
+  for (let i = 0; i < argv.length; i += 1) {
+    const arg = argv[i];
+    if (!arg.startsWith('--')) continue;
+    const key = arg.slice(2);
+    const next = argv[i + 1];
+    if (next && !next.startsWith('--')) {
+      parsed.set(key, next);
+      i += 1;
+    } else {
+      parsed.set(key, 'true');
+    }
+  }
+  return parsed;
+}
+
+function gitConfig(key) {
+  try {
+    const value = execFileSync('git', ['config', '--global', key], {
+      encoding: 'utf8',
+      stdio: ['ignore', 'pipe', 'ignore'],
+    }).trim();
+    return value || null;
+  } catch {
+    return null;
+  }
+}
+
+function loadEnvFile(path) {
+  if (!existsSync(path)) return;
+  for (const line of readFileSync(path, 'utf8').split('\n')) {
+    const trimmed = line.trim();
+    if (!trimmed || trimmed.startsWith('#')) continue;
+    const match = trimmed.match(/^([A-Za-z_][A-Za-z0-9_]*)=(.*)$/);
+    if (!match) continue;
+    const [, key, rawValue] = match;
+    if (process.env[key] != null) continue;
+    process.env[key] = unquoteEnvValue(rawValue.trim());
+  }
+}
+
+function unquoteEnvValue(value) {
+  if (
+    (value.startsWith('"') && value.endsWith('"')) ||
+    (value.startsWith("'") && value.endsWith("'"))
+  ) {
+    return value.slice(1, -1);
+  }
+  return value;
+}
+
+function readJsonIfExists(path) {
+  if (!existsSync(path)) return null;
+  try {
+    return JSON.parse(readFileSync(path, 'utf8'));
+  } catch {
+    return null;
+  }
+}
+
+function normalizeUrl(value) {
+  return String(value).replace(/\/+$/, '');
+}
+
+function resolveHome(value) {
+  if (value === '~') return homedir();
+  if (value.startsWith('~/')) return resolve(homedir(), value.slice(2));
+  return resolve(value);
+}
+
+function parsePort(value) {
+  const port = Number(value);
+  if (!Number.isInteger(port) || port <= 0 || port > 65535) {
+    console.error(`Invalid listen port: ${value}`);
+    process.exit(2);
+  }
+  return port;
+}
+
+function parseOptionalMoney(value) {
+  if (value == null || value === '') return 0;
+  const parsed = Number(String(value).replace(/[$,]/g, ''));
+  if (!Number.isFinite(parsed) || parsed < 0) {
+    console.error(`Invalid USD amount: ${value}`);
+    process.exit(2);
+  }
+  return parsed;
+}
+
+function generateApiKey() {
+  const raw = `lm_live_${randomBytes(24).toString('hex')}`;
+  return {
+    raw,
+    prefix: raw.slice(0, 12),
+    hash: hashApiKey(raw),
+  };
+}
+
+function hashApiKey(raw) {
+  return createHash('sha256').update(raw).digest('hex');
+}
+
+function escapeRegExp(value) {
+  return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
+function escapeXml(value) {
+  return String(value)
+    .replace(/&/g, '&amp;')
+    .replace(/</g, '&lt;')
+    .replace(/>/g, '&gt;')
+    .replace(/"/g, '&quot;')
+    .replace(/'/g, '&apos;');
+}

package/package.json

@@ -0,0 +1,43 @@
+{
+  "name": "@traice/codex-collector",
+  "version": "0.1.0",
+  "description": "Local Codex usage collector for trAIce Internal Spend.",
+  "license": "UNLICENSED",
+  "type": "module",
+  "bin": {
+    "traice-codex-collector": "cli.mjs"
+  },
+  "scripts": {
+    "check": "node --check cli.mjs && node --check install.mjs && node --check collector.mjs",
+    "prepublishOnly": "npm run check && npm pack --dry-run"
+  },
+  "engines": {
+    "node": ">=20"
+  },
+  "repository": {
+    "type": "git",
+    "url": "git+ssh://git@ssh.github.com/shmulikdav/trAIce.git",
+    "directory": "tools/internal-spend-codex"
+  },
+  "bugs": {
+    "url": "https://github.com/shmulikdav/trAIce/issues"
+  },
+  "homepage": "https://runtraice.com",
+  "keywords": [
+    "traice",
+    "codex",
+    "otel",
+    "usage",
+    "tokens",
+    "llm"
+  ],
+  "publishConfig": {
+    "access": "public"
+  },
+  "files": [
+    "cli.mjs",
+    "collector.mjs",
+    "install.mjs",
+    "README.md"
+  ]
+}