From f0460762b88dc043f54232f274a8d420958be02a Mon Sep 17 00:00:00 2001 From: Neil Date: Mon, 14 Sep 2026 12:06:53 -0700 Subject: [PATCH] fix(omp): validate remote hooks before reviving retired panes Co-authored-by: euclodius <244837876+euclodius@users.noreply.github.com> --- .../server-omp-retired-ingress.test.ts | 68 +++++++++++++++++++ .../server/server-ingest-remote.ts | 65 +++++++++++------- ...ay-session-agent-hooks.integration.test.ts | 25 +++++++ src/main/ssh/ssh-relay-session.ts | 6 +- 4 files changed, 139 insertions(+), 25 deletions(-) create mode 100644 src/main/agent-hooks/server-omp-retired-ingress.test.ts diff --git a/src/main/agent-hooks/server-omp-retired-ingress.test.ts b/src/main/agent-hooks/server-omp-retired-ingress.test.ts new file mode 100644 index 00000000000..a9844a6e1c2 --- /dev/null +++ b/src/main/agent-hooks/server-omp-retired-ingress.test.ts @@ -0,0 +1,68 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { AgentHookServer } from './server' +import { PANE } from './server.test-fixtures' + +vi.mock('../telemetry/client', () => ({ track: vi.fn() })) +vi.mock('../telemetry/cohort-classifier', () => ({ getCohortAtEmit: vi.fn() })) +afterEach(() => vi.restoreAllMocks()) + +const working = { agentType: 'omp', state: 'working', prompt: 'new turn' } + +describe('OMP retired-pane remote ingress', () => { + it('recovers a legacy source-less OMP new turn using validated payload identity', () => { + const server = new AgentHookServer() + server.retirePaneAuthority(PANE) + server.ingestRemote( + { paneKey: PANE, hookEventName: 'before_agent_start', payload: working }, + 'ssh' + ) + expect(server.getStatusSnapshot()).toEqual([ + expect.objectContaining({ + paneKey: PANE, + observation: expect.objectContaining({ boundary: true }) + }) + ]) + }) + + it.each([ + { payload: { ...working, state: 'invalid' } }, + { payload: { ...working, agentType: 'claude' } }, + { isReplay: 'true' }, + { isReplay: null }, + { launchToken: 42 }, + { providerSessionOnly: true } + ])('rejected metadata leaves retirement intact: %j', (invalid) => { + const server = new AgentHookServer() + server.retirePaneAuthority(PANE) + const input = { + paneKey: PANE, + source: 'omp', + hookEventName: 'before_agent_start', + payload: working, + ...invalid + } + // Exercise raw JSON ingress, including malformed fields a typed caller cannot create. + server.ingestRemote(JSON.parse(JSON.stringify(input)), 'ssh') + expect(server.getStatusSnapshot()).toEqual([]) + server.ingestRemote( + { + paneKey: PANE, + source: 'omp', + hookEventName: 'agent_end', + payload: { ...working, state: 'done' } + }, + 'ssh' + ) + expect(server.getStatusSnapshot()).toEqual([]) + }) + + it('a replay cannot restore a retired OMP pane', () => { + const server = new AgentHookServer() + server.retirePaneAuthority(PANE) + server.ingestRemote( + { paneKey: PANE, hookEventName: 'before_agent_start', isReplay: true, payload: working }, + 'ssh' + ) + expect(server.getStatusSnapshot()).toEqual([]) + }) +}) diff --git a/src/main/agent-hooks/server/server-ingest-remote.ts b/src/main/agent-hooks/server/server-ingest-remote.ts index d14714ae09b..7cd9e8779a0 100644 --- a/src/main/agent-hooks/server/server-ingest-remote.ts +++ b/src/main/agent-hooks/server/server-ingest-remote.ts @@ -91,6 +91,12 @@ export abstract class AgentHookServerIngestRemote extends AgentHookServerIngestS if (paneKey.length > MAX_PANE_KEY_LEN || !parsedPaneKey) { return } + if ( + (envelope.isReplay !== undefined && typeof envelope.isReplay !== 'boolean') || + (envelope.launchToken !== undefined && typeof envelope.launchToken !== 'string') + ) { + return + } // Why: fence relay spool replay at main so stale generations cannot overwrite hydrated state. if (envelope.isReplay === true) { const expectedLaunchTokenHash = this.hydratedLaunchTokenHashByPaneKey.get(paneKey) @@ -136,23 +142,6 @@ export abstract class AgentHookServerIngestRemote extends AgentHookServerIngestS (envelope.compactTrigger === 'manual' || envelope.compactTrigger === 'auto') ? envelope.compactTrigger : undefined - const statusDisposition = this.getAgentStatusDisposition(paneKey, { - source, - rawSource: envelope.source, - hookEventName, - isReplay: envelope.isReplay === true, - hasExplicitPrompt: envelope.hasExplicitPrompt === true, - launchToken: envelope.launchToken - }) - if (statusDisposition === 'suppress') { - return - } - if (statusDisposition === 'restart') { - // Why: same rebind as the HTTP path — a retired pane taking a new turn is a new session. - // Why paneKey, not envelope.paneKey: alias resolution already mapped it to the - // stable pane, so the rebind cannot land on a legacy key. - this.observations.rebind(paneKey) - } const worktreeId = envelope.worktreeId !== undefined && envelope.worktreeId.trim().length > 0 ? envelope.worktreeId.trim() @@ -184,12 +173,46 @@ export abstract class AgentHookServerIngestRemote extends AgentHookServerIngestS if (!validatedPayload) { return } + if ( + envelope.source !== undefined && + (source === 'omp' || validatedPayload.agentType === 'omp') && + envelope.source !== validatedPayload.agentType + ) { + return + } // Why: restore a shed roster only when its digest and turn identity still match the cache. let normalizedPayload = restoreShedStatusFields( validatedPayload, envelope.shedFields, this.state.lastStatusByPaneKey.get(paneKey)?.payload ) + if ( + envelope.providerSessionOnly === true && + !isValidPiProviderSessionOnly(providerSession, normalizedPayload.agentType) + ) { + return + } + // Older relays omit source; canonical OMP identity preserves boundary provenance. + const effectiveSource = + source ?? + (envelope.source === undefined && validatedPayload.agentType === 'omp' ? 'omp' : undefined) + const statusDisposition = this.getAgentStatusDisposition(paneKey, { + source: effectiveSource, + rawSource: envelope.source, + hookEventName, + isReplay: envelope.isReplay === true, + hasExplicitPrompt: envelope.hasExplicitPrompt === true, + launchToken: envelope.launchToken + }) + if (statusDisposition === 'suppress') { + return + } + if (statusDisposition === 'restart') { + // Why: same rebind as the HTTP path — a retired pane taking a new turn is a new session. + // Why paneKey, not envelope.paneKey: alias resolution already mapped it to the + // stable pane, so the rebind cannot land on a legacy key. + this.observations.rebind(paneKey) + } const previousStatus = this.state.lastStatusByPaneKey.get(paneKey) let acceptedCompactCompletion = false if (hookEventName === 'PreCompact' || hookEventName === 'PostCompact') { @@ -249,12 +272,6 @@ export abstract class AgentHookServerIngestRemote extends AgentHookServerIngestS ) { normalizedPayload = { ...normalizedPayload, prompt: previousStatus.payload.prompt } } - if ( - envelope.providerSessionOnly === true && - !isValidPiProviderSessionOnly(providerSession, normalizedPayload.agentType) - ) { - return - } const applyClaudeBackgroundWork = normalizedPayload.agentType === 'claude' && typeof envelope.claudeRunningNonAgentTask === 'boolean' && @@ -268,7 +285,7 @@ export abstract class AgentHookServerIngestRemote extends AgentHookServerIngestS }) const event: AgentHookEventPayload = { paneKey, - source, + source: effectiveSource, launchToken: statusDisposition === 'restart' ? undefined : envelope.launchToken, tabId, worktreeId, diff --git a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts index 87d5c77f205..0ac77cfd32c 100644 --- a/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts +++ b/src/main/ssh/ssh-relay-session-agent-hooks.integration.test.ts @@ -684,6 +684,31 @@ describe('SshRelaySession agent hooks over a fake relay transport', () => { expect(trackMock).not.toHaveBeenCalledWith('agent_prompt_sent', expect.anything()) }) + it.each([{ isReplay: 'true' }, { isReplay: null }, { launchToken: 42 }])( + 'rejects malformed OMP authority metadata before forwarding: %j', + async (invalid) => { + relay = createFakeRelay() + vi.mocked(deployAndLaunchRelay).mockResolvedValue({ + transport: relay.transport, + serverBuildId: 'test-relay-build', + platform: 'linux-x64' + }) + session = createSession('conn-omp-invalid') + await session.establish({} as SshConnection) + const ingestSpy = vi.spyOn(agentHookServer, 'ingestRemote') + const envelope = makeEnvelope({ + source: 'omp', + hookEventName: 'before_agent_start', + payload: { agentType: 'omp', state: 'working', prompt: 'new turn' } + }) + relay.notifyAgentHook(JSON.parse(JSON.stringify({ ...envelope, ...invalid }))) + await new Promise((resolve) => setImmediate(resolve)) + await new Promise((resolve) => setImmediate(resolve)) + expect(ingestSpy).not.toHaveBeenCalled() + ingestSpy.mockRestore() + } + ) + it('preserves replay metadata from remote hook notifications', async () => { relay = createFakeRelay() vi.mocked(deployAndLaunchRelay).mockResolvedValue({ diff --git a/src/main/ssh/ssh-relay-session.ts b/src/main/ssh/ssh-relay-session.ts index 6108f990d26..3bc02dd734a 100644 --- a/src/main/ssh/ssh-relay-session.ts +++ b/src/main/ssh/ssh-relay-session.ts @@ -1574,7 +1574,11 @@ export class SshRelaySession { return } const envelope = params - if (typeof envelope.paneKey !== 'string') { + if ( + typeof envelope.paneKey !== 'string' || + (envelope.isReplay !== undefined && typeof envelope.isReplay !== 'boolean') || + (envelope.launchToken !== undefined && typeof envelope.launchToken !== 'string') + ) { return } // Why: forward the agent CLI's env/version verbatim (not the relay's) so warn-once protocol-mismatch diagnostics fire for remote events too.