From 6b228864deb6c00d2a81ac4099a8bf231b8861c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?=
 <zhangyanghaha0407@outlook.com>
Date: Fri, 24 Apr 2026 03:35:20 -0400
Subject: [PATCH] fix(auth): persist agent token signing key

---
 src/main/agent-auth.ts        | 30 ++++++++++++++++-------
 tests/unit/agent-auth.test.ts | 45 +++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 9 deletions(-)
 create mode 100644 tests/unit/agent-auth.test.ts

diff --git a/src/main/agent-auth.ts b/src/main/agent-auth.ts
index 6a715222..4bf5b7d9 100644
--- a/src/main/agent-auth.ts
+++ b/src/main/agent-auth.ts
@@ -1,9 +1,11 @@
 import { createHmac, randomBytes, timingSafeEqual } from "node:crypto";
-import { mkdirSync, writeFileSync } from "node:fs";
+import { chmodSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
 import { join } from "node:path";
+import { FatalError } from "./errors/index.js";
 
 let signingKey: Buffer | null = null;
 const TOKEN_TTL_MS = 48 * 60 * 60 * 1000; // 48 hours
+const SIGNING_KEY_BYTES = 32;
 
 // Once-per-process warning when a malformed token shape is seen. Helps
 // diagnose schema breaks after upgrades — without leaking the full token.
@@ -12,7 +14,7 @@ function warnMalformedTokenOnce(token: string, reason: string) {
   if (loggedMalformedTokenWarning) return;
   loggedMalformedTokenWarning = true;
   const fingerprint = token.slice(0, 8);
-   
+
   console.warn(`[agent-auth] token shape invalid (${reason}); first8="${fingerprint}..."`);
 }
 
@@ -69,15 +71,25 @@ function getSigningKey() {
 }
 
 export function initializeAgentTokenSigningKey(profileDir: string) {
-  // Rotate the signing key on every launch. Any token issued by a prior
-  // process becomes unverifiable — agent sessions re-authenticate when
-  // they resume, so the cost of invalidating in-flight tokens is bounded
-  // and the blast radius of a leaked key is "until next restart".
   mkdirSync(profileDir, { recursive: true });
   const keyPath = join(profileDir, "agent-api-signing.key");
-  const nextKey = randomBytes(32);
-  writeFileSync(keyPath, nextKey, { mode: 0o600 });
-  signingKey = nextKey;
+  let key: Buffer;
+  try {
+    key = readFileSync(keyPath);
+  } catch (error) {
+    if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
+    key = randomBytes(SIGNING_KEY_BYTES);
+    writeFileSync(keyPath, key, { mode: 0o600, flag: "wx" });
+  }
+  if (key.length !== SIGNING_KEY_BYTES) {
+    throw new FatalError({
+      code: "agent_auth.signing_key.invalid",
+      message: `Agent token signing key must be ${SIGNING_KEY_BYTES} bytes.`,
+      context: { keyPath, actualBytes: key.length },
+    });
+  }
+  chmodSync(keyPath, 0o600);
+  signingKey = key;
 }
 
 export function createAgentToken(agentId: string, companyId: string, runId: string): string {
diff --git a/tests/unit/agent-auth.test.ts b/tests/unit/agent-auth.test.ts
new file mode 100644
index 00000000..c186fe97
--- /dev/null
+++ b/tests/unit/agent-auth.test.ts
@@ -0,0 +1,45 @@
+import { mkdtemp, rm, writeFile } from "node:fs/promises";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import { afterEach, describe, expect, it } from "vitest";
+import {
+  createAgentToken,
+  initializeAgentTokenSigningKey,
+  validateAgentToken,
+} from "../../src/main/agent-auth";
+import { FatalError } from "../../src/main/errors";
+
+const tempDirs: string[] = [];
+
+async function createTempDir() {
+  const dir = await mkdtemp(join(tmpdir(), "agentcompany-agent-auth-"));
+  tempDirs.push(dir);
+  return dir;
+}
+
+afterEach(async () => {
+  await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
+});
+
+describe("agent token signing key", () => {
+  it("keeps existing run tokens valid across profile reinitialization", async () => {
+    const profileDir = await createTempDir();
+    initializeAgentTokenSigningKey(profileDir);
+    const token = createAgentToken("agent-1", "company-1", "run-1");
+
+    initializeAgentTokenSigningKey(profileDir);
+
+    expect(validateAgentToken(token)).toMatchObject({
+      agentId: "agent-1",
+      companyId: "company-1",
+      runId: "run-1",
+    });
+  });
+
+  it("treats a corrupt signing key as fatal state", async () => {
+    const profileDir = await createTempDir();
+    await writeFile(join(profileDir, "agent-api-signing.key"), Buffer.from("too-short"));
+
+    expect(() => initializeAgentTokenSigningKey(profileDir)).toThrow(FatalError);
+  });
+});