diff --git a/apps/daemon/src/server.ts b/apps/daemon/src/server.ts index b52cbed144..f129805eab 100644 --- a/apps/daemon/src/server.ts +++ b/apps/daemon/src/server.ts @@ -6606,6 +6606,7 @@ export async function startServer({ // E-lite: stamp the last-activity clock BEFORE the disabled-watchdog bail // so `last_progress_age_ms` is recorded even when the watchdog is off. run.lastAgentActivityAt = Date.now(); + toolTokenRegistry.refreshRun(toolTokenGrant?.runId ?? runId); const delay = activeInactivityTimeoutMs(); if (delay <= 0) return; clearInactivityWatchdog(); diff --git a/apps/daemon/src/tool-tokens.ts b/apps/daemon/src/tool-tokens.ts index 6376ea3c4d..04110b152f 100644 --- a/apps/daemon/src/tool-tokens.ts +++ b/apps/daemon/src/tool-tokens.ts @@ -76,6 +76,7 @@ export type ToolTokenValidationResult = interface StoredToolTokenGrant extends ToolTokenGrant { tokenHash: string; expiresAtMs: number; + ttlMs: number; timer: NodeJS.Timeout; } @@ -88,7 +89,13 @@ function createOpaqueToolToken(): string { } function asPublicGrant(stored: StoredToolTokenGrant): ToolTokenGrant { - const { tokenHash: _tokenHash, expiresAtMs: _expiresAtMs, timer: _timer, ...grant } = stored; + const { + tokenHash: _tokenHash, + expiresAtMs: _expiresAtMs, + ttlMs: _ttlMs, + timer: _timer, + ...grant + } = stored; return grant; } @@ -150,6 +157,7 @@ export class ToolTokenRegistry { issuedAt: new Date(nowMs).toISOString(), expiresAt: new Date(expiresAtMs).toISOString(), expiresAtMs, + ttlMs, timer, ...(options.pluginSnapshotId ? { pluginSnapshotId: options.pluginSnapshotId } : {}), ...(options.pluginTrust ? { pluginTrust: options.pluginTrust } : {}), @@ -195,6 +203,31 @@ export class ToolTokenRegistry { return { ok: true, grant: asPublicGrant(stored) }; } + refreshRun(runId: string, nowMs = Date.now()): number { + const runTokens = this.#tokenHashesByRunId.get(runId); + if (!runTokens) return 0; + + let refreshed = 0; + for (const hash of [...runTokens]) { + const stored = this.#byTokenHash.get(hash); + if (!stored) continue; + if (nowMs >= stored.expiresAtMs) { + this.revokeToken(stored.token, 'ttl_expired'); + continue; + } + + clearTimeout(stored.timer); + stored.expiresAtMs = nowMs + stored.ttlMs; + stored.expiresAt = new Date(stored.expiresAtMs).toISOString(); + stored.timer = setTimeout(() => { + this.revokeToken(stored.token, 'ttl_expired'); + }, stored.ttlMs); + stored.timer.unref?.(); + refreshed += 1; + } + return refreshed; + } + revokeToken(token: string | null | undefined, _reason: ToolTokenRevocationReason = 'manual'): boolean { if (!token) return false; const hash = tokenHash(token); diff --git a/apps/daemon/tests/tool-token-active-run-lease.test.ts b/apps/daemon/tests/tool-token-active-run-lease.test.ts new file mode 100644 index 0000000000..7fab4f7e74 --- /dev/null +++ b/apps/daemon/tests/tool-token-active-run-lease.test.ts @@ -0,0 +1,129 @@ +import type { Server } from 'node:http'; +import { randomUUID } from 'node:crypto'; +import { chmod, mkdtemp, rm, writeFile } from 'node:fs/promises'; +import os from 'node:os'; +import path from 'node:path'; +import { afterEach, expect, it, vi } from 'vitest'; + +import { startServer } from '../src/server.js'; +import { toolTokenRegistry } from '../src/tool-tokens.js'; + +type StartedServer = { url: string; server: Server; shutdown?: () => Promise | void }; + +const originalInactivityTimeout = process.env.OD_CHAT_RUN_INACTIVITY_TIMEOUT_MS; +let started: StartedServer | null = null; +let binDir: string | null = null; + +afterEach(async () => { + vi.restoreAllMocks(); + toolTokenRegistry.clear(); + await Promise.resolve(started?.shutdown?.()); + if (started?.server) { + await new Promise((resolve) => started?.server.close(() => resolve())); + } + started = null; + if (binDir) await rm(binDir, { recursive: true, force: true }); + binDir = null; + if (originalInactivityTimeout === undefined) { + delete process.env.OD_CHAT_RUN_INACTIVITY_TIMEOUT_MS; + } else { + process.env.OD_CHAT_RUN_INACTIVITY_TIMEOUT_MS = originalInactivityTimeout; + } +}); + +async function writeActiveFakeClaude(dir: string): Promise { + const bin = path.join(dir, 'claude'); + await writeFile( + bin, + `#!/usr/bin/env node +if (process.argv.includes('--version')) { console.log('claude-code 1.0.0-token-lease'); process.exit(0); } +if (process.argv.includes('--help')) { console.log('Usage: claude -p [--include-partial-messages] [--add-dir DIR]'); process.exit(0); } +setTimeout(() => { + console.log(JSON.stringify({ type: 'system', subtype: 'init', model: 'claude-token-lease' })); +}, 100); +setTimeout(() => { + console.log(JSON.stringify({ + type: 'assistant', + message: { + id: 'msg-token-lease', + content: [{ type: 'text', text: 'Still working.' }], + stop_reason: 'end_turn' + } + })); +}, 200); +setTimeout(() => process.exit(0), 300); +`, + 'utf8', + ); + await chmod(bin, 0o755); + return bin; +} + +it('refreshes the tool token on agent activity with the inactivity watchdog disabled, then revokes it on exit', async () => { + process.env.OD_CHAT_RUN_INACTIVITY_TIMEOUT_MS = '0'; + binDir = await mkdtemp(path.join(os.tmpdir(), 'od-token-lease-bin-')); + const fakeClaude = await writeActiveFakeClaude(binDir); + const refreshSpy = vi.spyOn(toolTokenRegistry, 'refreshRun'); + + started = (await startServer({ port: 0, returnServer: true })) as StartedServer; + const configResponse = await fetch(`${started.url}/api/app-config`, { + method: 'PUT', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + agentId: 'claude', + agentCliEnv: { claude: { CLAUDE_BIN: fakeClaude } }, + telemetry: { metrics: false, content: false, artifactManifest: false }, + privacyDecisionAt: Date.now(), + }), + }); + expect(configResponse.status).toBe(200); + + const projectId = `token_lease_${randomUUID()}`; + const projectResponse = await fetch(`${started.url}/api/projects`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + id: projectId, + name: 'Token lease smoke', + metadata: { kind: 'prototype' }, + skipDiscoveryBrief: true, + }), + }); + expect(projectResponse.status).toBe(200); + const { conversationId } = (await projectResponse.json()) as { conversationId: string }; + + const runResponse = await fetch(`${started.url}/api/runs`, { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-od-analytics-device-id': 'token-lease-test', + 'x-od-analytics-session-id': 'token-lease-session', + 'x-od-analytics-client-type': 'web', + }, + body: JSON.stringify({ + projectId, + conversationId, + assistantMessageId: `assistant_${randomUUID()}`, + clientRequestId: `client_${randomUUID()}`, + agentId: 'claude', + message: 'emit activity before exiting', + currentPrompt: 'emit activity before exiting', + }), + }); + expect(runResponse.status).toBe(202); + const { runId } = (await runResponse.json()) as { runId: string }; + + let status = ''; + for (let i = 0; i < 100; i++) { + const response = await fetch(`${started.url}/api/runs/${runId}`); + expect(response.status).toBe(200); + status = ((await response.json()) as { status: string }).status; + if (status === 'failed' || status === 'succeeded' || status === 'canceled') break; + await new Promise((resolve) => setTimeout(resolve, 20)); + } + + expect(status).toBe('succeeded'); + const runRefreshes = refreshSpy.mock.calls.filter(([refreshedRunId]) => refreshedRunId === runId); + expect(runRefreshes.length).toBeGreaterThanOrEqual(2); + expect(toolTokenRegistry.activeRunTokenCount(runId)).toBe(0); +}); diff --git a/apps/daemon/tests/tool-tokens.test.ts b/apps/daemon/tests/tool-tokens.test.ts index 5ddf0f8add..74036f740f 100644 --- a/apps/daemon/tests/tool-tokens.test.ts +++ b/apps/daemon/tests/tool-tokens.test.ts @@ -78,6 +78,110 @@ describe('run-scoped tool tokens', () => { expect(registry.activeTokenCount()).toBe(0); }); + it('refreshes an active run token as an inactivity lease without changing its scope', () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + const registry = new ToolTokenRegistry(); + const grant = registry.mint({ + runId: 'run-refresh', + projectId: 'project-a', + allowedEndpoints: ['/api/tools/media/generate'], + allowedOperations: ['media:generate'], + ttlMs: 10, + }); + + vi.advanceTimersByTime(9); + expect(registry.refreshRun('run-refresh')).toBe(1); + vi.advanceTimersByTime(1); + + expect(registry.validate(grant.token)).toMatchObject({ + ok: true, + grant: { + token: grant.token, + runId: 'run-refresh', + projectId: 'project-a', + allowedEndpoints: ['/api/tools/media/generate'], + allowedOperations: ['media:generate'], + issuedAt: grant.issuedAt, + expiresAt: new Date(1_019).toISOString(), + }, + }); + + vi.advanceTimersByTime(9); + expect(registry.validate(grant.token)).toMatchObject({ + ok: false, + code: 'TOOL_TOKEN_INVALID', + }); + }); + + it('does not recreate unknown, expired, or manually revoked tokens during a run refresh', () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + const registry = new ToolTokenRegistry(); + + expect(registry.refreshRun('run-unknown')).toBe(0); + + const expired = registry.mint({ + runId: 'run-expired-refresh', + projectId: 'project-a', + ttlMs: 10, + }); + vi.advanceTimersByTime(10); + expect(registry.refreshRun('run-expired-refresh')).toBe(0); + expect(registry.validate(expired.token)).toMatchObject({ + ok: false, + code: 'TOOL_TOKEN_INVALID', + }); + + const revoked = registry.mint({ + runId: 'run-revoked-refresh', + projectId: 'project-a', + ttlMs: 10, + }); + expect(registry.revokeToken(revoked.token)).toBe(true); + expect(registry.refreshRun('run-revoked-refresh')).toBe(0); + expect(registry.validate(revoked.token)).toMatchObject({ + ok: false, + code: 'TOOL_TOKEN_INVALID', + }); + expect(registry.activeTokenCount()).toBe(0); + }); + + it('refreshes only the active run among concurrent runs for the same project', () => { + vi.useFakeTimers(); + vi.setSystemTime(1_000); + const registry = new ToolTokenRegistry(); + const refreshed = registry.mint({ + runId: 'run-refreshed', + projectId: 'project-a', + ttlMs: 10, + }); + const alsoRefreshed = registry.mint({ + runId: 'run-refreshed', + projectId: 'project-a', + ttlMs: 10, + }); + const untouched = registry.mint({ + runId: 'run-untouched', + projectId: 'project-a', + ttlMs: 10, + }); + + vi.advanceTimersByTime(9); + expect(registry.refreshRun('run-refreshed')).toBe(2); + vi.advanceTimersByTime(1); + + expect(registry.validate(refreshed.token).ok).toBe(true); + expect(registry.validate(alsoRefreshed.token).ok).toBe(true); + expect(registry.validate(untouched.token)).toMatchObject({ + ok: false, + code: 'TOOL_TOKEN_INVALID', + }); + expect(registry.activeRunTokenCount('run-refreshed')).toBe(2); + expect(registry.activeRunTokenCount('run-untouched')).toBe(0); + registry.clear(); + }); + it('uses the chat tool endpoint and operation allowlists by default', () => { const registry = new ToolTokenRegistry(); const grant = registry.mint({ runId: 'run-defaults', projectId: 'project-a', nowMs: 1_000 });