Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion services/gastown/container/src/agent-runner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,8 @@ GASTOWN_AGENT_ID="${env.GASTOWN_AGENT_ID}"
GASTOWN_RIG_ID="${env.GASTOWN_RIG_ID}"
GASTOWN_TOWN_ID="${env.GASTOWN_TOWN_ID}"`);

// Fall back to X-Town-Config for KILOCODE_TOKEN if not in request or process.env
// Fall back to cached config for KILOCODE_TOKEN if not in request or process.env.
// New requests should provide this through envVars rather than X-Town-Config.
if (!env.KILOCODE_TOKEN) {
const townConfig = getCurrentTownConfig();
const tokenFromConfig =
Expand Down
80 changes: 80 additions & 0 deletions services/gastown/container/src/control-server.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
import { describe, it, expect, vi } from 'vitest';

vi.mock('./agent-runner', () => ({
runAgent: vi.fn(),
resolveGitCredentials: vi.fn(),
writeMayorSystemPromptToAgentsMd: vi.fn(),
}));

vi.mock('./process-manager', () => ({
stopAgent: vi.fn(),
sendMessage: vi.fn(),
updateAgentModel: vi.fn(),
getAgentStatus: vi.fn(),
activeAgentCount: vi.fn(() => 0),
activeServerCount: vi.fn(() => 0),
getUptime: vi.fn(() => 0),
getStartTime: vi.fn(() => '2026-01-01T00:00:00.000Z'),
getMayorReadyAt: vi.fn(() => null),
stopAll: vi.fn(),
drainAll: vi.fn(),
isDraining: vi.fn(() => false),
getAgentEvents: vi.fn(),
registerEventSink: vi.fn(),
refreshTokenForAllAgents: vi.fn(),
listAgents: vi.fn(() => []),
awaitHydration: vi.fn(),
}));

vi.mock('./logger', () => ({
log: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
}));

vi.mock('./heartbeat', () => ({
startHeartbeat: vi.fn(),
stopHeartbeat: vi.fn(),
notifyContainerReady: vi.fn(),
}));

vi.mock('./dashboard-context', () => ({
pushContext: vi.fn(),
}));

vi.mock('./git-manager', () => ({
mergeBranch: vi.fn(),
setupRigBrowseWorktree: vi.fn(),
}));

const { app } = await import('./control-server');

describe('/sync-config', () => {
it('preserves startup GASTOWN_API_URL when the sync payload omits it', async () => {
const previousApiUrl = process.env.GASTOWN_API_URL;
const previousToken = process.env.KILOCODE_TOKEN;
const previousGitToken = process.env.GIT_TOKEN;
process.env.GASTOWN_API_URL = 'https://startup.example.com';
process.env.KILOCODE_TOKEN = 'old-token';

try {
const response = await app.request('/sync-config', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ envVars: { GIT_TOKEN: 'new-token' } }),
});

expect(response.status).toBe(200);
expect(process.env.GASTOWN_API_URL).toBe('https://startup.example.com');
expect(process.env.KILOCODE_TOKEN).toBeUndefined();
expect(process.env.GIT_TOKEN).toBe('new-token');
} finally {
if (previousApiUrl !== undefined) process.env.GASTOWN_API_URL = previousApiUrl;
else delete process.env.GASTOWN_API_URL;

if (previousToken !== undefined) process.env.KILOCODE_TOKEN = previousToken;
else delete process.env.KILOCODE_TOKEN;

if (previousGitToken !== undefined) process.env.GIT_TOKEN = previousGitToken;
else delete process.env.GIT_TOKEN;
}
});
});
136 changes: 57 additions & 79 deletions services/gastown/container/src/control-server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import { log } from './logger';
import { startHeartbeat, stopHeartbeat, notifyContainerReady } from './heartbeat';
import { pushContext as pushDashboardContext } from './dashboard-context';
import { mergeBranch, setupRigBrowseWorktree } from './git-manager';
import { CONFIG_SYNC_ENV_KEYS, RESERVED_ENV_KEYS } from './env-keys';
import {
StartAgentRequest,
SendMessageRequest,
Expand All @@ -41,42 +42,27 @@ import type {
const MAX_TICKETS = 1000;
const streamTickets = new Map<string, { agentId: string; expiresAt: number }>();

// Minimal Zod schema for the town config delivered via X-Town-Config header.
export { RESERVED_ENV_KEYS };

// Minimal Zod schema for the non-secret town config delivered via X-Town-Config header.
// Uses z.record() so any string-keyed object is accepted and future keys are preserved.
const TownConfigHeader = z.record(z.string(), z.unknown());

const SyncConfigRequest = z.object({
envVars: z.record(z.string(), z.string()).optional(),
});

const CONFIG_SYNC_ENV_KEYS_PRESERVE_WHEN_ABSENT = new Set(['GASTOWN_API_URL']);

// Last-known-good town config. Updated on every request that carries the header.
// Used as a fallback by code that runs outside a request context (e.g. background tasks).
let lastKnownTownConfig: Record<string, unknown> | null = null;
let lastBodySyncedEnvVars: Record<string, string> | null = null;

// Track which custom env var keys were applied last sync so removed keys can be cleared.
let lastAppliedEnvVarKeys = new Set<string>();

// Env keys managed by the control plane that custom env_vars must never override.
// If a custom key collides with a reserved key, the infra value wins and the
// custom value is silently ignored — matching the !(key in env) guard in buildAgentEnv.
export const RESERVED_ENV_KEYS = new Set([
'KILOCODE_TOKEN',
'GIT_TOKEN',
'GITHUB_TOKEN',
'GITLAB_TOKEN',
'GITLAB_INSTANCE_URL',
'GITHUB_CLI_PAT',
'GH_TOKEN',
'GASTOWN_GIT_AUTHOR_NAME',
'GASTOWN_GIT_AUTHOR_EMAIL',
'GASTOWN_DISABLE_AI_COAUTHOR',
'GASTOWN_ORGANIZATION_ID',
'GASTOWN_CONTAINER_TOKEN',
'GASTOWN_SESSION_TOKEN',
'GASTOWN_API_URL',
// Runtime routing vars read by pending-nudge routes and plugin clients —
// must never be overwritten by user-supplied env_vars.
'GASTOWN_TOWN_ID',
'GASTOWN_RIG_ID',
]);

/** Get the latest town config delivered via X-Town-Config header. */
/** Get the latest non-secret town config delivered via X-Town-Config header. */
export function getCurrentTownConfig(): Record<string, unknown> | null {
return lastKnownTownConfig;
}
Expand All @@ -87,18 +73,41 @@ export function getLastAppliedEnvVarKeys(): Set<string> {
}

/**
* Sync config-derived env vars from the last-known town config into
* process.env. Safe to call at any time — no-ops when no config is cached.
* Sync config-derived env vars into process.env. Secret-bearing values arrive
* in the JSON body because Cloudflare event logs can record request headers.
*/
function syncTownConfigToProcessEnv(): void {
function syncTownConfigToProcessEnv(envVars?: Record<string, string>): void {
if (envVars) {
for (const key of CONFIG_SYNC_ENV_KEYS) {
Comment thread
evanjacobson marked this conversation as resolved.
if (CONFIG_SYNC_ENV_KEYS_PRESERVE_WHEN_ABSENT.has(key)) continue;
if (!(key in envVars)) delete process.env[key];
}
for (const [key, value] of Object.entries(envVars)) {
if (RESERVED_ENV_KEYS.has(key) && !CONFIG_SYNC_ENV_KEYS.has(key)) continue;
Comment thread
evanjacobson marked this conversation as resolved.
process.env[key] = value;
}

const customKeys = Object.keys(envVars).filter(key => !RESERVED_ENV_KEYS.has(key));
const newCustomKeys = new Set(customKeys);
const customEnvVars = Object.fromEntries(customKeys.map(key => [key, envVars[key]]));
for (const key of lastAppliedEnvVarKeys) {
if (!newCustomKeys.has(key) && !RESERVED_ENV_KEYS.has(key)) delete process.env[key];
}
lastAppliedEnvVarKeys = newCustomKeys;
lastBodySyncedEnvVars = customEnvVars;
lastKnownTownConfig = {
...(lastKnownTownConfig ?? {}),
env_vars: customEnvVars,
};
return;
}

const cfg = getCurrentTownConfig();
if (!cfg) return;

const CONFIG_ENV_MAP: Array<[string, string]> = [
['github_cli_pat', 'GITHUB_CLI_PAT'],
['git_author_name', 'GASTOWN_GIT_AUTHOR_NAME'],
['git_author_email', 'GASTOWN_GIT_AUTHOR_EMAIL'],
['kilocode_token', 'KILOCODE_TOKEN'],
];
for (const [cfgKey, envKey] of CONFIG_ENV_MAP) {
const val = cfg[cfgKey];
Expand All @@ -109,23 +118,6 @@ function syncTownConfigToProcessEnv(): void {
}
}

const gitAuth = cfg.git_auth;
if (typeof gitAuth === 'object' && gitAuth !== null) {
const auth = gitAuth as Record<string, unknown>;
for (const [authKey, envKey] of [
['github_token', 'GIT_TOKEN'],
['gitlab_token', 'GITLAB_TOKEN'],
['gitlab_instance_url', 'GITLAB_INSTANCE_URL'],
] as const) {
const val = auth[authKey];
if (typeof val === 'string' && val) {
process.env[envKey] = val;
} else {
delete process.env[envKey];
}
}
}

if (cfg.disable_ai_coauthor) {
process.env.GASTOWN_DISABLE_AI_COAUTHOR = '1';
} else {
Expand All @@ -140,32 +132,11 @@ function syncTownConfigToProcessEnv(): void {
} else {
delete process.env.GASTOWN_ORGANIZATION_ID;
}

// Apply custom env_vars from the town config. Reserved infra keys are
// skipped so the control-plane values always take precedence — matching the
// !(key in env) guard in buildAgentEnv.
const rawEnvVars = cfg.env_vars;
const customEnvVars: Record<string, string> =
rawEnvVars !== null && typeof rawEnvVars === 'object' && !Array.isArray(rawEnvVars)
? (rawEnvVars as Record<string, string>)
: {};
const newCustomKeys = new Set(Object.keys(customEnvVars));
// Remove keys that were present in the previous sync but are gone now.
// Skip reserved keys — deleting those would wipe a control-plane value.
for (const key of lastAppliedEnvVarKeys) {
if (!newCustomKeys.has(key) && !RESERVED_ENV_KEYS.has(key)) delete process.env[key];
}
// Apply current custom env vars, skipping reserved keys.
for (const [key, value] of Object.entries(customEnvVars)) {
if (RESERVED_ENV_KEYS.has(key)) continue;
process.env[key] = String(value);
}
lastAppliedEnvVarKeys = newCustomKeys;
}

export const app = new Hono();

// Parse and validate town config from X-Town-Config header (sent by TownDO on
// Parse and validate non-secret town config from X-Town-Config header (sent by TownDO on
// every request). The validated config is stored in a module-level cache
// accessible via getCurrentTownConfig().
app.use('*', async (c, next) => {
Expand All @@ -175,11 +146,14 @@ app.use('*', async (c, next) => {
const raw: unknown = JSON.parse(configHeader);
const result = TownConfigHeader.safeParse(raw);
if (result.success) {
lastKnownTownConfig = result.data;
const hasToken =
typeof result.data.kilocode_token === 'string' && result.data.kilocode_token.length > 0;
const headerConfig = { ...result.data };
delete headerConfig.env_vars;
lastKnownTownConfig = {
...headerConfig,
...(lastBodySyncedEnvVars ? { env_vars: lastBodySyncedEnvVars } : {}),
};
console.log(
`[control-server] X-Town-Config received: hasKilocodeToken=${hasToken} keys=${Object.keys(result.data).join(',')}`
`[control-server] X-Town-Config received: keys=${Object.keys(result.data).join(',')}`
);
} else {
console.warn(
Expand Down Expand Up @@ -302,12 +276,17 @@ app.post('/refresh-token', async c => {
});

// POST /sync-config
// Push config-derived env vars from X-Town-Config into process.env on
// Push config-derived env vars from the JSON body into process.env on
// the running container. Called by TownDO.syncConfigToContainer() after
// persisting env vars to DO storage, so the live process picks up
// changes (e.g. refreshed KILOCODE_TOKEN) without a container restart.
app.post('/sync-config', async c => {
syncTownConfigToProcessEnv();
const body: unknown = await c.req.json().catch(() => ({}));
const parsed = SyncConfigRequest.safeParse(body);
if (!parsed.success) {
return c.json({ error: 'Invalid request body', issues: parsed.error.issues }, 400);
}
syncTownConfigToProcessEnv(parsed.data.envVars);
return c.json({ synced: true });
});

Expand Down Expand Up @@ -422,9 +401,8 @@ app.patch('/agents/:agentId/model', async c => {
process.env.GASTOWN_ORGANIZATION_ID = parsed.data.organizationId;
}

// Sync config-derived env vars from X-Town-Config into process.env so
// the SDK server restart picks up fresh tokens and git identity.
// The middleware already parsed the header into lastKnownTownConfig.
// Sync non-secret config fields from X-Town-Config into process.env.
// Secret-bearing values are refreshed through /sync-config or request envVars.
syncTownConfigToProcessEnv();

await updateAgentModel(
Expand Down
35 changes: 35 additions & 0 deletions services/gastown/container/src/env-keys.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
// Env keys managed by the control plane that custom env_vars must never override.
// If a custom key collides with a reserved key, the infra value wins.
export const RESERVED_ENV_KEYS = new Set([
'KILOCODE_TOKEN',
'GIT_TOKEN',
'GITHUB_TOKEN',
'GITLAB_TOKEN',
'GITLAB_INSTANCE_URL',
'GITHUB_CLI_PAT',
'GH_TOKEN',
'GASTOWN_GIT_AUTHOR_NAME',
'GASTOWN_GIT_AUTHOR_EMAIL',
'GASTOWN_DISABLE_AI_COAUTHOR',
'GASTOWN_ORGANIZATION_ID',
'GASTOWN_CONTAINER_TOKEN',
'GASTOWN_SESSION_TOKEN',
'GASTOWN_API_URL',
// Runtime routing vars read by pending-nudge routes and plugin clients —
// must never be overwritten by user-supplied env_vars.
'GASTOWN_TOWN_ID',
'GASTOWN_RIG_ID',
]);

export const CONFIG_SYNC_ENV_KEYS = new Set([
'KILOCODE_TOKEN',
'GIT_TOKEN',
'GITLAB_TOKEN',
'GITLAB_INSTANCE_URL',
'GITHUB_CLI_PAT',
'GASTOWN_GIT_AUTHOR_NAME',
'GASTOWN_GIT_AUTHOR_EMAIL',
'GASTOWN_DISABLE_AI_COAUTHOR',
'GASTOWN_ORGANIZATION_ID',
'GASTOWN_API_URL',
]);
2 changes: 1 addition & 1 deletion services/gastown/container/src/process-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1530,7 +1530,7 @@ export async function sendMessage(agentId: string, prompt: string): Promise<void
* 1. Agent's `organizationId` field — set at startup from StartAgentRequest,
* survives process.env restores and model hot-swaps.
* 2. GASTOWN_ORGANIZATION_ID env var — set by control-server on /agents/start
* and updated on every PATCH /model via X-Town-Config.
* and updated by config sync or PATCH /model.
* 3. KILO_CONFIG_CONTENT — legacy fallback, may be absent after env restore.
*/
function extractOrganizationId(agent?: ManagedAgent): string | undefined {
Expand Down
Loading