@hellcoder/companion 0.96.0

package/server/usage-limits.ts

@@ -0,0 +1,225 @@
+import { execSync, execFileSync } from "node:child_process";
+import { existsSync, readFileSync, writeFileSync } from "node:fs";
+import { join } from "node:path";
+import { homedir } from "node:os";
+
+export interface UsageLimits {
+  five_hour: { utilization: number; resets_at: string | null } | null;
+  seven_day: { utilization: number; resets_at: string | null } | null;
+  extra_usage: {
+    is_enabled: boolean;
+    monthly_limit: number;
+    used_credits: number;
+    utilization: number | null;
+  } | null;
+}
+
+const OAUTH_TOKEN_URL = "https://platform.claude.com/v1/oauth/token";
+const OAUTH_CLIENT_ID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e";
+
+// In-memory cache (60s TTL)
+const CACHE_DURATION_MS = 60 * 1000;
+let cache: { data: UsageLimits; timestamp: number } | null = null;
+
+interface OAuthCredentials {
+  accessToken: string;
+  refreshToken: string;
+  expiresAt: number;
+  [key: string]: unknown;
+}
+
+interface RawCredentials {
+  raw: string;
+  parsed: Record<string, unknown>;
+  oauth: OAuthCredentials;
+  sourcePath?: string;
+}
+
+// Credential file candidates - matches claude-container-auth.ts
+const CREDENTIAL_FILE_NAMES = [
+  ".credentials.json",
+  "auth.json",
+  ".auth.json",
+  "credentials.json",
+];
+
+function readCredentialsFromFile(): RawCredentials | null {
+  const home =
+    process.env.USERPROFILE || process.env.HOME || homedir() || "";
+  const claudeDir = join(home, ".claude");
+
+  for (const fileName of CREDENTIAL_FILE_NAMES) {
+    const credPath = join(claudeDir, fileName);
+    if (!existsSync(credPath)) continue;
+    try {
+      const raw = readFileSync(credPath, "utf-8");
+      const parsed = JSON.parse(raw);
+      if (!parsed?.claudeAiOauth?.accessToken) continue;
+      return { raw, parsed, oauth: parsed.claudeAiOauth, sourcePath: credPath };
+    } catch {
+      continue;
+    }
+  }
+  return null;
+}
+
+function readRawCredentials(): RawCredentials | null {
+  try {
+    // macOS: use Keychain via security command
+    if (process.platform === "darwin") {
+      const raw = execSync(
+        'security find-generic-password -s "Claude Code-credentials" -w',
+        { encoding: "utf-8", timeout: 5000, stdio: ["pipe", "pipe", "pipe"] },
+      ).trim();
+
+      const decoded = raw.startsWith("{")
+        ? raw
+        : Buffer.from(raw, "hex").toString("utf-8");
+
+      const parsed = JSON.parse(decoded);
+      if (!parsed?.claudeAiOauth?.accessToken) return null;
+      return { raw: decoded, parsed, oauth: parsed.claudeAiOauth };
+    }
+
+    // Windows, Linux, Docker, and other platforms: read from credential files
+    return readCredentialsFromFile();
+  } catch {
+    return null;
+  }
+}
+
+function writeCredentials(creds: Record<string, unknown>, sourcePath?: string): void {
+  try {
+    const json = JSON.stringify(creds);
+    if (process.platform === "darwin") {
+      execFileSync(
+        "security",
+        ["add-generic-password", "-U", "-s", "Claude Code-credentials", "-a", "Claude Code", "-w", json],
+        { timeout: 5000, stdio: ["pipe", "pipe", "pipe"] },
+      );
+    } else {
+      // Write back to the same file that was read, or default to .credentials.json
+      const credPath = sourcePath ?? join(
+        process.env.USERPROFILE || process.env.HOME || homedir() || "",
+        ".claude",
+        ".credentials.json",
+      );
+      writeFileSync(credPath, json, "utf-8");
+    }
+  } catch {
+    // best-effort
+  }
+}
+
+async function refreshAccessToken(refreshToken: string): Promise<{
+  accessToken: string;
+  refreshToken: string;
+  expiresIn: number;
+} | null> {
+  try {
+    const res = await fetch(OAUTH_TOKEN_URL, {
+      method: "POST",
+      headers: { "Content-Type": "application/x-www-form-urlencoded" },
+      body: new URLSearchParams({
+        grant_type: "refresh_token",
+        refresh_token: refreshToken,
+        client_id: OAUTH_CLIENT_ID,
+      }),
+    });
+    if (!res.ok) return null;
+    const data = await res.json();
+    return {
+      accessToken: data.access_token,
+      refreshToken: data.refresh_token,
+      expiresIn: data.expires_in,
+    };
+  } catch {
+    return null;
+  }
+}
+
+export function getCredentials(): string | null {
+  const creds = readRawCredentials();
+  return creds?.oauth.accessToken ?? null;
+}
+
+async function getValidAccessToken(): Promise<string | null> {
+  const creds = readRawCredentials();
+  if (!creds) return null;
+
+  const { oauth } = creds;
+
+  // Token still valid (with 5min buffer)
+  if (oauth.expiresAt && Date.now() < oauth.expiresAt - 5 * 60 * 1000) {
+    return oauth.accessToken;
+  }
+
+  // Token expired - try to refresh
+  if (!oauth.refreshToken) return null;
+
+  const refreshed = await refreshAccessToken(oauth.refreshToken);
+  if (!refreshed) return null;
+
+  // Update stored credentials
+  creds.parsed.claudeAiOauth = {
+    ...oauth,
+    accessToken: refreshed.accessToken,
+    refreshToken: refreshed.refreshToken,
+    expiresAt: Date.now() + refreshed.expiresIn * 1000,
+  };
+  writeCredentials(creds.parsed, creds.sourcePath);
+
+  return refreshed.accessToken;
+}
+
+export async function fetchUsageLimits(
+  token: string,
+): Promise<UsageLimits | null> {
+  try {
+    const response = await fetch("https://api.anthropic.com/api/oauth/usage", {
+      method: "GET",
+      headers: {
+        Accept: "application/json",
+        "Content-Type": "application/json",
+        "User-Agent": "claude-code/2.1.39",
+        Authorization: `Bearer ${token}`,
+        "anthropic-beta": "oauth-2025-04-20",
+      },
+    });
+
+    if (!response.ok) return null;
+
+    const data = await response.json();
+    return {
+      five_hour: data.five_hour || null,
+      seven_day: data.seven_day || null,
+      extra_usage: data.extra_usage || null,
+    };
+  } catch {
+    return null;
+  }
+}
+
+export async function getUsageLimits(): Promise<UsageLimits> {
+  const empty: UsageLimits = {
+    five_hour: null,
+    seven_day: null,
+    extra_usage: null,
+  };
+  try {
+    if (cache && Date.now() - cache.timestamp < CACHE_DURATION_MS) {
+      return cache.data;
+    }
+
+    const token = await getValidAccessToken();
+    if (!token) return empty;
+
+    const limits = await fetchUsageLimits(token);
+    if (!limits) return empty;
+
+    cache = { data: limits, timestamp: Date.now() };
+    return limits;
+  } catch {
+    return empty;
+  }
+}

package/server/worktree-tracker.test.ts

@@ -0,0 +1,243 @@
+import {
+  mkdtempSync,
+  rmSync,
+  writeFileSync,
+  mkdirSync,
+  readFileSync,
+} from "node:fs";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import type { WorktreeMapping } from "./worktree-tracker.js";
+
+let tempDir: string;
+let WorktreeTracker: typeof import("./worktree-tracker.js").WorktreeTracker;
+
+const mockHomedir = vi.hoisted(() => {
+  let dir = "";
+  return {
+    get: () => dir,
+    set: (d: string) => {
+      dir = d;
+    },
+  };
+});
+
+vi.mock("node:os", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("node:os")>();
+  return {
+    ...actual,
+    homedir: () => mockHomedir.get(),
+  };
+});
+
+beforeEach(async () => {
+  tempDir = mkdtempSync(join(tmpdir(), "wt-test-"));
+  mockHomedir.set(tempDir);
+  vi.resetModules();
+  const mod = await import("./worktree-tracker.js");
+  WorktreeTracker = mod.WorktreeTracker;
+});
+
+afterEach(() => {
+  rmSync(tempDir, { recursive: true, force: true });
+});
+
+function makeMapping(overrides: Partial<WorktreeMapping> = {}): WorktreeMapping {
+  return {
+    sessionId: "session-1",
+    repoRoot: "/repo",
+    branch: "feat-1",
+    worktreePath: "/worktrees/feat-1",
+    createdAt: Date.now(),
+    ...overrides,
+  };
+}
+
+function trackerFilePath(): string {
+  return join(tempDir, ".companion", "worktrees.json");
+}
+
+function readTrackerFile(): WorktreeMapping[] {
+  return JSON.parse(readFileSync(trackerFilePath(), "utf-8"));
+}
+
+// ─── Constructor ─────────────────────────────────────────────────────────────
+
+describe("WorktreeTracker", () => {
+  describe("constructor", () => {
+    it("initializes with empty mappings when no file exists", () => {
+      const tracker = new WorktreeTracker();
+      expect(tracker.getBySession("anything")).toBeNull();
+    });
+
+    it("loads existing mappings from disk", () => {
+      const mapping = makeMapping();
+      mkdirSync(join(tempDir, ".companion"), { recursive: true });
+      writeFileSync(trackerFilePath(), JSON.stringify([mapping]));
+
+      const tracker = new WorktreeTracker();
+      expect(tracker.getBySession("session-1")).toEqual(mapping);
+    });
+
+    it("handles corrupt JSON gracefully", () => {
+      mkdirSync(join(tempDir, ".companion"), { recursive: true });
+      writeFileSync(trackerFilePath(), "NOT VALID JSON {{{");
+
+      const tracker = new WorktreeTracker();
+      expect(tracker.getBySession("anything")).toBeNull();
+    });
+  });
+
+  // ─── addMapping ──────────────────────────────────────────────────────────────
+
+  describe("addMapping", () => {
+    it("persists mapping to disk", () => {
+      const tracker = new WorktreeTracker();
+      const mapping = makeMapping();
+      tracker.addMapping(mapping);
+
+      const onDisk = readTrackerFile();
+      expect(onDisk).toHaveLength(1);
+      expect(onDisk[0]).toEqual(mapping);
+    });
+
+    it("replaces existing mapping for same sessionId", () => {
+      const tracker = new WorktreeTracker();
+      const original = makeMapping({ branch: "feat-1" });
+      tracker.addMapping(original);
+
+      const updated = makeMapping({ branch: "feat-2", worktreePath: "/worktrees/feat-2" });
+      tracker.addMapping(updated);
+
+      const onDisk = readTrackerFile();
+      expect(onDisk).toHaveLength(1);
+      expect(onDisk[0].branch).toBe("feat-2");
+      expect(onDisk[0].worktreePath).toBe("/worktrees/feat-2");
+    });
+
+    it("allows multiple mappings with different sessionIds", () => {
+      const tracker = new WorktreeTracker();
+      tracker.addMapping(makeMapping({ sessionId: "s1" }));
+      tracker.addMapping(makeMapping({ sessionId: "s2" }));
+      tracker.addMapping(makeMapping({ sessionId: "s3" }));
+
+      const onDisk = readTrackerFile();
+      expect(onDisk).toHaveLength(3);
+    });
+  });
+
+  // ─── removeBySession ────────────────────────────────────────────────────────
+
+  describe("removeBySession", () => {
+    it("returns removed mapping and persists deletion", () => {
+      const tracker = new WorktreeTracker();
+      const mapping = makeMapping();
+      tracker.addMapping(mapping);
+
+      const removed = tracker.removeBySession("session-1");
+      expect(removed).toEqual(mapping);
+
+      const onDisk = readTrackerFile();
+      expect(onDisk).toHaveLength(0);
+    });
+
+    it("returns null for unknown sessionId", () => {
+      const tracker = new WorktreeTracker();
+      const result = tracker.removeBySession("nonexistent");
+      expect(result).toBeNull();
+    });
+  });
+
+  // ─── getBySession ───────────────────────────────────────────────────────────
+
+  describe("getBySession", () => {
+    it("returns mapping when found", () => {
+      const tracker = new WorktreeTracker();
+      const mapping = makeMapping();
+      tracker.addMapping(mapping);
+
+      expect(tracker.getBySession("session-1")).toEqual(mapping);
+    });
+
+    it("returns null when not found", () => {
+      const tracker = new WorktreeTracker();
+      expect(tracker.getBySession("nonexistent")).toBeNull();
+    });
+  });
+
+  // ─── getSessionsForWorktree ─────────────────────────────────────────────────
+
+  describe("getSessionsForWorktree", () => {
+    it("returns all mappings sharing a worktree path", () => {
+      const tracker = new WorktreeTracker();
+      const sharedPath = "/worktrees/shared";
+      tracker.addMapping(makeMapping({ sessionId: "s1", worktreePath: sharedPath }));
+      tracker.addMapping(makeMapping({ sessionId: "s2", worktreePath: sharedPath }));
+      tracker.addMapping(makeMapping({ sessionId: "s3", worktreePath: "/worktrees/other" }));
+
+      const results = tracker.getSessionsForWorktree(sharedPath);
+      expect(results).toHaveLength(2);
+      expect(results.map((r) => r.sessionId).sort()).toEqual(["s1", "s2"]);
+    });
+
+    it("returns empty array when no sessions use the worktree", () => {
+      const tracker = new WorktreeTracker();
+      tracker.addMapping(makeMapping({ worktreePath: "/worktrees/other" }));
+
+      const results = tracker.getSessionsForWorktree("/worktrees/nonexistent");
+      expect(results).toEqual([]);
+    });
+  });
+
+  // ─── getSessionsForRepo ─────────────────────────────────────────────────────
+
+  describe("getSessionsForRepo", () => {
+    it("returns all mappings for a repo root", () => {
+      const tracker = new WorktreeTracker();
+      tracker.addMapping(makeMapping({ sessionId: "s1", repoRoot: "/repo-a", branch: "feat-1" }));
+      tracker.addMapping(makeMapping({ sessionId: "s2", repoRoot: "/repo-a", branch: "feat-2" }));
+      tracker.addMapping(makeMapping({ sessionId: "s3", repoRoot: "/repo-b", branch: "feat-1" }));
+
+      const results = tracker.getSessionsForRepo("/repo-a");
+      expect(results).toHaveLength(2);
+      expect(results.map((r) => r.sessionId).sort()).toEqual(["s1", "s2"]);
+    });
+
+    it("returns empty array when no sessions belong to the repo", () => {
+      const tracker = new WorktreeTracker();
+      const results = tracker.getSessionsForRepo("/nonexistent-repo");
+      expect(results).toEqual([]);
+    });
+  });
+
+  // ─── isWorktreeInUse ────────────────────────────────────────────────────────
+
+  describe("isWorktreeInUse", () => {
+    it("returns true when another session uses the worktree", () => {
+      const tracker = new WorktreeTracker();
+      tracker.addMapping(makeMapping({ sessionId: "s1", worktreePath: "/worktrees/feat" }));
+
+      expect(tracker.isWorktreeInUse("/worktrees/feat")).toBe(true);
+    });
+
+    it("returns false when no session uses the worktree", () => {
+      const tracker = new WorktreeTracker();
+      expect(tracker.isWorktreeInUse("/worktrees/feat")).toBe(false);
+    });
+
+    it("excludes the specified session from the check", () => {
+      const tracker = new WorktreeTracker();
+      tracker.addMapping(makeMapping({ sessionId: "s1", worktreePath: "/worktrees/feat" }));
+
+      expect(tracker.isWorktreeInUse("/worktrees/feat", "s1")).toBe(false);
+    });
+
+    it("returns true when other sessions use it despite excludeSessionId", () => {
+      const tracker = new WorktreeTracker();
+      tracker.addMapping(makeMapping({ sessionId: "s1", worktreePath: "/worktrees/feat" }));
+      tracker.addMapping(makeMapping({ sessionId: "s2", worktreePath: "/worktrees/feat" }));
+
+      expect(tracker.isWorktreeInUse("/worktrees/feat", "s1")).toBe(true);
+    });
+  });
+});

package/server/worktree-tracker.ts

@@ -0,0 +1,84 @@
+import {
+  mkdirSync,
+  readFileSync,
+  writeFileSync,
+  existsSync,
+} from "node:fs";
+import { join, dirname } from "node:path";
+import { COMPANION_HOME } from "./paths.js";
+
+// ─── Types ──────────────────────────────────────────────────────────────────
+
+export interface WorktreeMapping {
+  sessionId: string;
+  repoRoot: string;
+  branch: string;
+  /** Actual git branch in the worktree (may differ from `branch` for -wt-N branches) */
+  actualBranch?: string;
+  worktreePath: string;
+  createdAt: number;
+}
+
+// ─── Paths ──────────────────────────────────────────────────────────────────
+
+const TRACKER_PATH = join(COMPANION_HOME, "worktrees.json");
+
+// ─── Tracker ────────────────────────────────────────────────────────────────
+
+export class WorktreeTracker {
+  private mappings: WorktreeMapping[] = [];
+
+  constructor() {
+    this.load();
+  }
+
+  load(): WorktreeMapping[] {
+    try {
+      if (existsSync(TRACKER_PATH)) {
+        const raw = readFileSync(TRACKER_PATH, "utf-8");
+        this.mappings = JSON.parse(raw) as WorktreeMapping[];
+      }
+    } catch {
+      this.mappings = [];
+    }
+    return this.mappings;
+  }
+
+  private save(): void {
+    mkdirSync(dirname(TRACKER_PATH), { recursive: true });
+    writeFileSync(TRACKER_PATH, JSON.stringify(this.mappings, null, 2), "utf-8");
+  }
+
+  addMapping(mapping: WorktreeMapping): void {
+    // Remove any existing mapping for this session
+    this.mappings = this.mappings.filter((m) => m.sessionId !== mapping.sessionId);
+    this.mappings.push(mapping);
+    this.save();
+  }
+
+  removeBySession(sessionId: string): WorktreeMapping | null {
+    const idx = this.mappings.findIndex((m) => m.sessionId === sessionId);
+    if (idx === -1) return null;
+    const [removed] = this.mappings.splice(idx, 1);
+    this.save();
+    return removed;
+  }
+
+  getBySession(sessionId: string): WorktreeMapping | null {
+    return this.mappings.find((m) => m.sessionId === sessionId) || null;
+  }
+
+  getSessionsForWorktree(worktreePath: string): WorktreeMapping[] {
+    return this.mappings.filter((m) => m.worktreePath === worktreePath);
+  }
+
+  getSessionsForRepo(repoRoot: string): WorktreeMapping[] {
+    return this.mappings.filter((m) => m.repoRoot === repoRoot);
+  }
+
+  isWorktreeInUse(worktreePath: string, excludeSessionId?: string): boolean {
+    return this.mappings.some(
+      (m) => m.worktreePath === worktreePath && m.sessionId !== excludeSessionId,
+    );
+  }
+}

package/server/ws-auth.test.ts

@@ -0,0 +1,59 @@
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { createToken } from "./middleware/managed-auth.js";
+import { authenticateManagedWebSocket } from "./ws-auth.js";
+
+const TEST_SECRET = "test-secret-key-for-ws-auth";
+
+describe("authenticateManagedWebSocket", () => {
+  const savedSecret = process.env.COMPANION_AUTH_SECRET;
+
+  beforeEach(() => {
+    process.env.COMPANION_AUTH_SECRET = TEST_SECRET;
+  });
+
+  afterEach(() => {
+    if (savedSecret === undefined) delete process.env.COMPANION_AUTH_SECRET;
+    else process.env.COMPANION_AUTH_SECRET = savedSecret;
+  });
+
+  it("returns 401 when no token is provided", async () => {
+    const req = new Request("https://example.com/ws/browser/abc");
+    const result = await authenticateManagedWebSocket(req);
+    expect(result.ok).toBe(false);
+    expect(result.status).toBe(401);
+  });
+
+  it("accepts a valid query token", async () => {
+    const token = await createToken(TEST_SECRET, 60);
+    const req = new Request(`https://example.com/ws/browser/abc?token=${token}`);
+    const result = await authenticateManagedWebSocket(req);
+    expect(result.ok).toBe(true);
+  });
+
+  it("accepts a valid cookie token", async () => {
+    const token = await createToken(TEST_SECRET, 60);
+    const req = new Request("https://example.com/ws/browser/abc", {
+      headers: { cookie: `companion_token=${token}` },
+    });
+    const result = await authenticateManagedWebSocket(req);
+    expect(result.ok).toBe(true);
+  });
+
+  it("prefers query token over cookie", async () => {
+    const token = await createToken(TEST_SECRET, 60);
+    const req = new Request(`https://example.com/ws/browser/abc?token=${token}`, {
+      headers: { cookie: "companion_token=bad.token" },
+    });
+    const result = await authenticateManagedWebSocket(req);
+    expect(result.ok).toBe(true);
+  });
+
+  it("returns 500 when secret is missing", async () => {
+    delete process.env.COMPANION_AUTH_SECRET;
+    const req = new Request("https://example.com/ws/browser/abc");
+    const result = await authenticateManagedWebSocket(req);
+    expect(result.ok).toBe(false);
+    expect(result.status).toBe(500);
+  });
+});
+

package/server/ws-auth.ts

@@ -0,0 +1,41 @@
+import { verifyToken } from "./middleware/managed-auth.js";
+
+export interface WsAuthResult {
+  ok: boolean;
+  status: number;
+  body?: string;
+}
+
+function getCookie(header: string | null, name: string): string | undefined {
+  if (!header) return undefined;
+  const match = header.match(new RegExp(`(?:^|;\\s*)${name}=([^;]*)`));
+  return match?.[1];
+}
+
+/**
+ * Authenticate browser/terminal WebSocket upgrade requests in managed mode.
+ * Accepts token in query param or companion_token cookie (query takes precedence).
+ */
+export async function authenticateManagedWebSocket(req: Request): Promise<WsAuthResult> {
+  const secret = process.env.COMPANION_AUTH_SECRET?.trim();
+  if (!secret) {
+    return { ok: false, status: 500, body: "Server misconfigured" };
+  }
+
+  const url = new URL(req.url);
+  const queryToken = url.searchParams.get("token");
+  const cookieToken = getCookie(req.headers.get("cookie"), "companion_token");
+  const token = queryToken || cookieToken;
+
+  if (!token) {
+    return { ok: false, status: 401, body: "Unauthorized" };
+  }
+
+  const valid = await verifyToken(token, secret);
+  if (!valid) {
+    return { ok: false, status: 401, body: "Unauthorized" };
+  }
+
+  return { ok: true, status: 200 };
+}
+