From e895987527ebb776b877be44a88f52bf91eaf251 Mon Sep 17 00:00:00 2001 From: Brennan Benson <79079362+brennanb2025@users.noreply.github.com> Date: Mon, 14 Sep 2026 00:05:06 -0700 Subject: [PATCH] fix(orchestration): require registered structured worker pane key --- .../structured-worker-identity.test.ts | 75 +++++++++++++++++-- .../runtime/structured-worker-identity.ts | 31 ++++++-- 2 files changed, 94 insertions(+), 12 deletions(-) diff --git a/src/main/runtime/structured-worker-identity.test.ts b/src/main/runtime/structured-worker-identity.test.ts index 9b678ebb0eb..2a66536291a 100644 --- a/src/main/runtime/structured-worker-identity.test.ts +++ b/src/main/runtime/structured-worker-identity.test.ts @@ -1,6 +1,9 @@ import { describe, expect, it, beforeEach } from 'vitest' import { isTerminalLeafId, parsePaneKey } from '../../shared/stable-pane-id' -import { structuredAgentSessionPaneKey } from '../../shared/structured-agent-session-projection' +import { + structuredAgentSessionPaneKey, + structuredAgentSessionTabId +} from '../../shared/structured-agent-session-projection' import { selectExactWorkerProviderSession } from './orchestration/worker-provider-session' import { structuredWorkerChildIdentityEnv } from './structured-worker-child-identity-env' import { @@ -85,12 +88,57 @@ describe('structured worker identity', () => { ) }) - it("accepts a persisted pane key for its own session and rejects another session's", () => { + it('accepts only the registered pane key for its session', () => { + const handle = mintStructuredWorkerHandle() const paneKey = mintStructuredWorkerPaneKey(SESSION_ID) - expect(structuredWorkerPaneKeyBelongsToSession(paneKey, SESSION_ID)).toBe(true) - expect(structuredWorkerPaneKeyBelongsToSession(paneKey, 'another-session-id')).toBe(false) - expect(structuredWorkerPaneKeyBelongsToSession('not-a-pane-key', SESSION_ID)).toBe(false) - expect(structuredWorkerPaneKeyBelongsToSession(null, SESSION_ID)).toBe(false) + structuredWorkerIdentities.register({ + handle, + sessionId: SESSION_ID, + agent: 'claude', + paneKey, + processIncarnation: structuredWorkerProcessIncarnation(SESSION_ID), + worktreeId: 'wt_1', + hostScope: { kind: 'local', hostId: 'local' } + }) + try { + expect(structuredWorkerPaneKeyBelongsToSession(paneKey, SESSION_ID)).toBe(true) + expect( + structuredWorkerPaneKeyBelongsToSession(mintStructuredWorkerPaneKey(SESSION_ID), SESSION_ID) + ).toBe(false) + expect(structuredWorkerPaneKeyBelongsToSession(paneKey, 'another-session-id')).toBe(false) + expect(structuredWorkerPaneKeyBelongsToSession('not-a-pane-key', SESSION_ID)).toBe(false) + expect(structuredWorkerPaneKeyBelongsToSession(null, SESSION_ID)).toBe(false) + } finally { + structuredWorkerIdentities.forget(handle) + } + }) + + it('rejects the deterministic public status key even for a registered worker', () => { + const handle = mintStructuredWorkerHandle() + const paneKey = mintStructuredWorkerPaneKey(SESSION_ID) + structuredWorkerIdentities.register({ + handle, + sessionId: SESSION_ID, + agent: 'claude', + paneKey, + processIncarnation: structuredWorkerProcessIncarnation(SESSION_ID), + worktreeId: 'wt_1', + hostScope: { kind: 'local', hostId: 'local' } + }) + try { + const statusPaneKey = structuredAgentSessionPaneKey( + structuredAgentSessionTabId(SESSION_ID), + SESSION_ID + ) + expect(structuredWorkerPaneKeyBelongsToSession(statusPaneKey, SESSION_ID)).toBe(false) + } finally { + structuredWorkerIdentities.forget(handle) + } + }) + + it('fails closed when the session has no registry record', () => { + const paneKey = mintStructuredWorkerPaneKey(SESSION_ID) + expect(structuredWorkerPaneKeyBelongsToSession(paneKey, SESSION_ID)).toBe(false) }) it('derives a pane key whose leaf passes the terminal leaf check', () => { @@ -169,6 +217,21 @@ describe('structured worker identity registry', () => { ).toBeNull() }) + it('refuses to rehydrate the deterministic public status key as a worker credential', () => { + expect( + registry.rehydrate({ + terminal_handle: mintStructuredWorkerHandle(), + pane_key: structuredAgentSessionPaneKey( + structuredAgentSessionTabId(SESSION_ID), + SESSION_ID + ), + process_incarnation: structuredWorkerProcessIncarnation(SESSION_ID), + worktree_id: 'wt_1', + host_scope: JSON.stringify({ kind: 'local', hostId: 'local' }) + }) + ).toBeNull() + }) + it('forgets both indexes', () => { const handle = mintStructuredWorkerHandle() registry.register({ diff --git a/src/main/runtime/structured-worker-identity.ts b/src/main/runtime/structured-worker-identity.ts index 161ae55dd5d..b29d68c297a 100644 --- a/src/main/runtime/structured-worker-identity.ts +++ b/src/main/runtime/structured-worker-identity.ts @@ -20,7 +20,10 @@ import type { AgentSessionRecord } from '../../shared/agent-session-record' import { LOCAL_EXECUTION_HOST_ID } from '../../shared/execution-host' -import { structuredAgentSessionTabId } from '../../shared/structured-agent-session-projection' +import { + structuredAgentSessionPaneKey, + structuredAgentSessionTabId +} from '../../shared/structured-agent-session-projection' import { isTerminalLeafId, makePaneKey, parsePaneKey } from '../../shared/stable-pane-id' import { parseWorkerTerminalHostScope, @@ -67,13 +70,30 @@ export function mintStructuredWorkerPaneKey(sessionId: string): string { return makePaneKey(structuredAgentSessionTabId(sessionId), randomUUID()) } -/** Integrity check for a persisted pane key: same session's tab, and a real terminal leaf. */ +/** Credential check: only the pane key registered for this session can prove its identity. */ export function structuredWorkerPaneKeyBelongsToSession( paneKey: string | null | undefined, sessionId: string ): boolean { + const registered = structuredWorkerIdentities.getBySessionId(sessionId) const parsed = paneKey ? parsePaneKey(paneKey) : null return Boolean( + registered && + registered.paneKey === paneKey && + parsed && + parsed.tabId === structuredAgentSessionTabId(sessionId) + ) +} + +/** Bootstrap validation for a durable row before its key can enter the registry. */ +function persistedStructuredWorkerPaneKeyIsValid( + paneKey: string | null | undefined, + sessionId: string +): paneKey is string { + const parsed = paneKey ? parsePaneKey(paneKey) : null + return Boolean( + paneKey && + paneKey !== structuredAgentSessionPaneKey(structuredAgentSessionTabId(sessionId), sessionId) && parsed && parsed.tabId === structuredAgentSessionTabId(sessionId) && isTerminalLeafId(parsed.leafId) @@ -176,9 +196,8 @@ export class StructuredWorkerIdentityRegistry { !hostScope || !row.worktree_id || !isStructuredWorkerHandle(row.terminal_handle) || - // The leaf is random, so the row IS the only source for it; verify only that it is a real - // leaf under this session's tab rather than trying to re-derive it. - !structuredWorkerPaneKeyBelongsToSession(row.pane_key, sessionId) + // The durable row bootstraps the registry after restart, so validate it before registration. + !persistedStructuredWorkerPaneKeyIsValid(row.pane_key, sessionId) ) { return null } @@ -187,7 +206,7 @@ export class StructuredWorkerIdentityRegistry { sessionId, // The row does not carry the provider; callers that need it read the durable record. agent: null, - paneKey: row.pane_key as string, + paneKey: row.pane_key, processIncarnation: structuredWorkerProcessIncarnation(sessionId), worktreeId: row.worktree_id, hostScope