fix(gastown): add JWT token refresh for persistent agents

cloudflare-gastown/container/plugin/client.test.ts

@@ -1,12 +1,20 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
-import { GastownClient, MayorGastownClient, GastownApiError, createClientFromEnv } from './client';
+import {
+  GastownClient,
+  MayorGastownClient,
+  GastownApiError,
+  createClientFromEnv,
+  registerPluginClient,
+  refreshPluginClientTokens,
+} from './client';
 import type { GastownEnv, MayorGastownEnv } from './types';
 
 const TEST_ENV: GastownEnv = {
   apiUrl: 'https://gastown.example.com',
   sessionToken: 'test-jwt-token',
   agentId: 'agent-111',
   rigId: 'rig-222',
+  townId: 'town-333',
 };
 
 function mockFetch(data: unknown, status = 200) {
@@ -48,7 +56,9 @@ describe('GastownClient', () => {
 
     expect(fetchMock).toHaveBeenCalledTimes(1);
     const [url, init] = fetchMock.mock.calls[0] as [string, RequestInit];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/agents/agent-111/prime');
+    expect(url).toBe(
+      'https://gastown.example.com/api/towns/town-333/rigs/rig-222/agents/agent-111/prime'
+    );
     const headers = new Headers(init.headers);
     expect(headers.get('Authorization')).toBe('Bearer test-jwt-token');
     expect(headers.get('Content-Type')).toBe('application/json');
@@ -81,7 +91,7 @@ describe('GastownClient', () => {
     expect(result).toEqual(bead);
 
     const [url] = (globalThis.fetch as ReturnType<typeof vi.fn>).mock.calls[0] as [string];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/beads/bead-1');
+    expect(url).toBe('https://gastown.example.com/api/towns/town-333/rigs/rig-222/beads/bead-1');
   });
 
   it('closeBead() sends agent_id in body', async () => {
@@ -94,7 +104,9 @@ describe('GastownClient', () => {
       string,
       RequestInit,
     ];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/beads/bead-1/close');
+    expect(url).toBe(
+      'https://gastown.example.com/api/towns/town-333/rigs/rig-222/beads/bead-1/close'
+    );
     expect(init.method).toBe('POST');
     expect(JSON.parse(init.body as string)).toEqual({ agent_id: 'agent-111' });
   });
@@ -112,7 +124,9 @@ describe('GastownClient', () => {
       string,
       RequestInit,
     ];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/agents/agent-111/done');
+    expect(url).toBe(
+      'https://gastown.example.com/api/towns/town-333/rigs/rig-222/agents/agent-111/done'
+    );
     expect(JSON.parse(init.body as string)).toEqual({
       branch: 'feat/test',
       pr_url: 'https://github.com/pr/1',
@@ -145,7 +159,9 @@ describe('GastownClient', () => {
     expect(result).toEqual(mail);
 
     const [url] = (globalThis.fetch as ReturnType<typeof vi.fn>).mock.calls[0] as [string];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/agents/agent-111/mail');
+    expect(url).toBe(
+      'https://gastown.example.com/api/towns/town-333/rigs/rig-222/agents/agent-111/mail'
+    );
   });
 
   it('writeCheckpoint() posts data to checkpoint endpoint', async () => {
@@ -157,7 +173,9 @@ describe('GastownClient', () => {
       string,
       RequestInit,
     ];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/agents/agent-111/checkpoint');
+    expect(url).toBe(
+      'https://gastown.example.com/api/towns/town-333/rigs/rig-222/agents/agent-111/checkpoint'
+    );
     expect(JSON.parse(init.body as string)).toEqual({ data: { step: 3, files: ['a.ts'] } });
   });
 
@@ -172,7 +190,7 @@ describe('GastownClient', () => {
       string,
       RequestInit,
     ];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/escalations');
+    expect(url).toBe('https://gastown.example.com/api/towns/town-333/rigs/rig-222/escalations');
     expect(JSON.parse(init.body as string)).toEqual({ title: 'blocked', priority: 'high' });
   });
 
@@ -246,7 +264,9 @@ describe('GastownClient', () => {
     // Verify no double slashes in the URL by calling prime
     void c.prime();
     const [url] = (globalThis.fetch as ReturnType<typeof vi.fn>).mock.calls[0] as [string];
-    expect(url).toBe('https://gastown.example.com/api/rigs/rig-222/agents/agent-111/prime');
+    expect(url).toBe(
+      'https://gastown.example.com/api/towns/town-333/rigs/rig-222/agents/agent-111/prime'
+    );
   });
 });
 
@@ -262,6 +282,7 @@ describe('createClientFromEnv', () => {
     process.env.GASTOWN_SESSION_TOKEN = 'tok';
     process.env.GASTOWN_AGENT_ID = 'agent-1';
     process.env.GASTOWN_RIG_ID = 'rig-1';
+    process.env.GASTOWN_TOWN_ID = 'town-1';
 
     const client = createClientFromEnv();
     expect(client).toBeInstanceOf(GastownClient);
@@ -358,3 +379,115 @@ describe('MayorGastownClient', () => {
     expect(url).toBe('https://gastown.example.com/api/mayor/town-1/tools/convoys/convoy-1');
   });
 });
+
+// ── Token refresh tests ─────────────────────────────────────────────────
+
+describe('GastownClient.setToken', () => {
+  const originalFetch = globalThis.fetch;
+
+  afterEach(() => {
+    globalThis.fetch = originalFetch;
+  });
+
+  it('uses the new token after setToken()', async () => {
+    const env: GastownEnv = {
+      apiUrl: 'https://gastown.example.com',
+      sessionToken: 'old-token',
+      agentId: 'agent-1',
+      rigId: 'rig-1',
+      townId: 'town-1',
+    };
+    const client = new GastownClient(env);
+
+    // Replace token
+    client.setToken('new-token');
+
+    const fetchMock = mockFetch({
+      agent: {},
+      hooked_bead: null,
+      undelivered_mail: [],
+      open_beads: [],
+    });
+    globalThis.fetch = fetchMock;
+
+    await client.prime();
+
+    const [, init] = fetchMock.mock.calls[0] as [string, RequestInit];
+    const headers = new Headers(init.headers);
+    expect(headers.get('Authorization')).toBe('Bearer new-token');
+  });
+});
+
+describe('MayorGastownClient.setToken', () => {
+  const originalFetch = globalThis.fetch;
+
+  afterEach(() => {
+    globalThis.fetch = originalFetch;
+  });
+
+  it('uses the new token after setToken()', async () => {
+    const env: MayorGastownEnv = {
+      apiUrl: 'https://gastown.example.com',
+      sessionToken: 'old-mayor-token',
+      agentId: 'mayor-1',
+      townId: 'town-1',
+    };
+    const client = new MayorGastownClient(env);
+
+    client.setToken('fresh-mayor-token');
+
+    const fetchMock = mockFetch([]);
+    globalThis.fetch = fetchMock;
+
+    await client.listRigs();
+
+    const [, init] = fetchMock.mock.calls[0] as [string, RequestInit];
+    const headers = new Headers(init.headers);
+    expect(headers.get('Authorization')).toBe('Bearer fresh-mayor-token');
+  });
+});
+
+describe('plugin client registry', () => {
+  const REGISTRY_KEY = Symbol.for('gastown.pluginClientRegistry');
+
+  beforeEach(() => {
+    // Clear the global registry before each test
+    (globalThis as Record<symbol, unknown>)[REGISTRY_KEY] = [];
+  });
+
+  it('registerPluginClient + refreshPluginClientTokens updates all clients', () => {
+    const client1 = new GastownClient({
+      apiUrl: 'https://example.com',
+      sessionToken: 'old-1',
+      agentId: 'a1',
+      rigId: 'r1',
+      townId: 't1',
+    });
+    const client2 = new MayorGastownClient({
+      apiUrl: 'https://example.com',
+      sessionToken: 'old-2',
+      agentId: 'a2',
+      townId: 't2',
+    });
+
+    registerPluginClient(client1);
+    registerPluginClient(client2);
+
+    refreshPluginClientTokens('refreshed-token');
+
+    // Verify by making a request and checking the auth header
+    const fetchMock = mockFetch([]);
+    const originalFetch = globalThis.fetch;
+    globalThis.fetch = fetchMock;
+
+    void client1.checkMail();
+    const [, init1] = fetchMock.mock.calls[0] as [string, RequestInit];
+    expect(new Headers(init1.headers).get('Authorization')).toBe('Bearer refreshed-token');
+
+    void client2.listRigs();
+    const [, init2] = fetchMock.mock.calls[1] as [string, RequestInit];
+    expect(new Headers(init2.headers).get('Authorization')).toBe('Bearer refreshed-token');
+
+    globalThis.fetch = originalFetch;
+  });
+});

cloudflare-gastown/container/plugin/client.ts

@@ -38,6 +38,11 @@ export class GastownClient {
     this.townId = env.townId;
   }
 
+  /** Hot-swap the session token for subsequent API calls. */
+  setToken(token: string): void {
+    this.token = token;
+  }
+
   private rigPath(path: string): string {
     return `${this.baseUrl}/api/towns/${this.townId}/rigs/${this.rigId}${path}`;
   }
@@ -46,11 +51,16 @@ export class GastownClient {
     return this.rigPath(`/agents/${this.agentId}${path}`);
   }
 
+  /** Returns the most up-to-date token: explicit setToken() value first, then this.token. */
+  private currentToken(): string {
+    return this.token;
+  }
+
   private async request<T>(url: string, init?: RequestInit): Promise<T> {
     // Normalize headers so callers can pass plain objects, Headers instances, or tuples
     const headers = new Headers(init?.headers);
     headers.set('Content-Type', 'application/json');
-    headers.set('Authorization', `Bearer ${this.token}`);
+    headers.set('Authorization', `Bearer ${this.currentToken()}`);
 
     let response: Response;
     try {
@@ -208,14 +218,24 @@ export class MayorGastownClient {
     this.townId = env.townId;
   }
 
+  /** Hot-swap the session token for subsequent API calls. */
+  setToken(token: string): void {
+    this.token = token;
+  }
+
   private mayorPath(path: string): string {
     return `${this.baseUrl}/api/mayor/${this.townId}/tools${path}`;
   }
 
+  /** Returns the most up-to-date token. */
+  private currentToken(): string {
+    return this.token;
+  }
+
   private async request<T>(url: string, init?: RequestInit): Promise<T> {
     const headers = new Headers(init?.headers);
     headers.set('Content-Type', 'application/json');
-    headers.set('Authorization', `Bearer ${this.token}`);
+    headers.set('Authorization', `Bearer ${this.currentToken()}`);
 
     let response: Response;
     try {
@@ -322,6 +342,40 @@ export class MayorGastownClient {
   }
 }
 
+// ── Plugin client token registry ─────────────────────────────────────────
+// The plugin client runs inside the kilo SDK plugin (loaded by createKilo()),
+// while the control server runs in the container's main module. They share
+// the same Bun process but different TypeScript project roots, so they can't
+// cross-import. We use a globalThis registry keyed by a well-known symbol so
+// the control server can push fresh tokens into plugin client instances.
+
+type TokenRefreshable = { setToken(token: string): void };
+
+const REGISTRY_KEY = Symbol.for('gastown.pluginClientRegistry');
+
+function getRegistry(): TokenRefreshable[] {
+  const g = globalThis as Record<symbol, unknown>;
+  if (!Array.isArray(g[REGISTRY_KEY])) {
+    g[REGISTRY_KEY] = [];
+  }
+  return g[REGISTRY_KEY] as TokenRefreshable[];
+}
+
+/** Register a plugin client so its token can be refreshed externally. */
+export function registerPluginClient(client: TokenRefreshable): void {
+  getRegistry().push(client);
+}
+
+/**
+ * Update the session token on all registered plugin clients.
+ * Accessible from any module in the same process via the global symbol.
+ */
+export function refreshPluginClientTokens(token: string): void {
+  for (const client of getRegistry()) {
+    client.setToken(token);
+  }
+}
+
 export class GastownApiError extends Error {
   readonly status: number;
 

cloudflare-gastown/container/plugin/index.ts

@@ -1,5 +1,10 @@
 import type { Plugin } from '@kilocode/plugin';
-import { createClientFromEnv, createMayorClientFromEnv, GastownApiError } from './client';
+import {
+  createClientFromEnv,
+  createMayorClientFromEnv,
+  GastownApiError,
+  registerPluginClient,
+} from './client';
 import { createTools } from './tools';
 import { createMayorTools } from './mayor-tools';
 
@@ -29,6 +34,7 @@ export const GastownPlugin: Plugin = async ({ client }) => {
   if (isMayor) {
     try {
       mayorClient = createMayorClientFromEnv();
+      registerPluginClient(mayorClient);
     } catch (err) {
       const message = err instanceof Error ? err.message : String(err);
       console.error(
@@ -41,6 +47,7 @@ export const GastownPlugin: Plugin = async ({ client }) => {
   } else {
     try {
       gastownClient = createClientFromEnv();
+      registerPluginClient(gastownClient);
     } catch (err) {
       const message = err instanceof Error ? err.message : String(err);
       console.error(