From 4468dc1f1bc7b3ecfdd7fc5e4090b36e56be66d8 Mon Sep 17 00:00:00 2001 From: Jinwoo-H Date: Mon, 17 Aug 2026 00:49:24 -0400 Subject: [PATCH] fix(orchestration): bind coordinator mutations to dispatch --- .../coordinator-decision-gates.test.ts | 18 ++++- .../coordinator-decision-gates.ts | 12 ++-- .../coordinator-escalation-triage.test.ts | 10 ++- .../coordinator-escalation-triage.ts | 10 ++- .../runtime/orchestration/coordinator.test.ts | 6 +- src/main/runtime/orchestration/db.ts | 3 +- .../orchestration/dispatch-message-binding.ts | 23 +++++++ .../orchestration/federation-sync.test.ts | 68 +++++++++++++++++++ .../runtime/orchestration/federation-sync.ts | 3 +- ...hestration-send-dispatch-authority.test.ts | 35 ++++++++++ src/main/runtime/rpc/methods/orchestration.ts | 20 ++++-- 11 files changed, 186 insertions(+), 22 deletions(-) create mode 100644 src/main/runtime/orchestration/dispatch-message-binding.ts diff --git a/src/main/runtime/orchestration/coordinator-decision-gates.test.ts b/src/main/runtime/orchestration/coordinator-decision-gates.test.ts index 74326daa079..e5a76ab6a22 100644 --- a/src/main/runtime/orchestration/coordinator-decision-gates.test.ts +++ b/src/main/runtime/orchestration/coordinator-decision-gates.test.ts @@ -23,7 +23,11 @@ describe('coordinator decision-gate authority', () => { subject: 'Need approval', type: 'decision_gate', senderPaneKey: 'tab_owner:leaf_owner', - payload: JSON.stringify({ taskId: task.id, question: 'Proceed?' }) + payload: JSON.stringify({ + taskId: task.id, + dispatchId: dispatch.id, + question: 'Proceed?' + }) }), (message) => logs.push(message) ) @@ -54,7 +58,11 @@ describe('coordinator decision-gate authority', () => { subject: 'Block the victim', type: 'decision_gate', senderPaneKey: 'tab_attacker:leaf_attacker', - payload: JSON.stringify({ taskId: victimTask.id, question: 'Stop?' }) + payload: JSON.stringify({ + taskId: victimTask.id, + dispatchId: attacker.id, + question: 'Stop?' + }) }), (message) => logs.push(message) ) @@ -79,7 +87,11 @@ describe('coordinator decision-gate authority', () => { to: 'term_coordinator', subject: 'Remote approval required', type: 'decision_gate', - payload: JSON.stringify({ taskId: task.id, question: 'Proceed remotely?' }) + payload: JSON.stringify({ + taskId: task.id, + dispatchId: dispatch.id, + question: 'Proceed remotely?' + }) }), () => {} ) diff --git a/src/main/runtime/orchestration/coordinator-decision-gates.ts b/src/main/runtime/orchestration/coordinator-decision-gates.ts index 4d3170296d5..3da386fc2ed 100644 --- a/src/main/runtime/orchestration/coordinator-decision-gates.ts +++ b/src/main/runtime/orchestration/coordinator-decision-gates.ts @@ -10,7 +10,7 @@ export function openDecisionGateFromMessage( ): void { onLog(`Decision gate from ${msg.from_handle}: ${msg.subject}`) - let payload: { taskId?: string; question?: string; options?: string[] } = {} + let payload: { taskId?: string; dispatchId?: string; question?: string; options?: string[] } = {} if (msg.payload) { try { payload = JSON.parse(msg.payload) @@ -19,8 +19,8 @@ export function openDecisionGateFromMessage( } } - if (!payload.taskId || !payload.question) { - onLog(`Warning: decision_gate missing taskId or question`) + if (!payload.taskId || !payload.dispatchId || !payload.question) { + onLog(`Warning: decision_gate missing taskId, dispatchId, or question`) return } @@ -29,7 +29,11 @@ export function openDecisionGateFromMessage( taskId: payload.taskId, question: payload.question, options: payload.options, - requester: { handle: msg.from_handle, paneKey: msg.sender_pane_key } + requester: { + handle: msg.from_handle, + paneKey: msg.sender_pane_key, + dispatchId: payload.dispatchId + } }) } catch (error) { if ( diff --git a/src/main/runtime/orchestration/coordinator-escalation-triage.test.ts b/src/main/runtime/orchestration/coordinator-escalation-triage.test.ts index 33172e3a699..af973100336 100644 --- a/src/main/runtime/orchestration/coordinator-escalation-triage.test.ts +++ b/src/main/runtime/orchestration/coordinator-escalation-triage.test.ts @@ -12,7 +12,11 @@ describe('coordinator escalation authority', () => { it('rejects an escalation targeting another active Dispatch', () => { db = new OrchestrationDb(':memory:') const attackerTask = db.createTask({ spec: 'attacker assignment' }) - db.createDispatchContext(attackerTask.id, 'term_attacker', 'tab_attacker:leaf_attacker') + const attacker = db.createDispatchContext( + attackerTask.id, + 'term_attacker', + 'tab_attacker:leaf_attacker' + ) const victimTask = db.createTask({ spec: 'victim assignment' }) const victim = db.createDispatchContext(victimTask.id, 'term_victim') const logs: string[] = [] @@ -25,7 +29,7 @@ describe('coordinator escalation authority', () => { subject: 'Fail the victim', type: 'escalation', senderPaneKey: 'tab_attacker:leaf_attacker', - payload: JSON.stringify({ taskId: victimTask.id }) + payload: JSON.stringify({ taskId: victimTask.id, dispatchId: attacker.id }) }), (message) => logs.push(message) ) @@ -47,7 +51,7 @@ describe('coordinator escalation authority', () => { to: 'term_coordinator', subject: 'Remote worker failed', type: 'escalation', - payload: JSON.stringify({ taskId: task.id }) + payload: JSON.stringify({ taskId: task.id, dispatchId: dispatch.id }) }), () => {} ) diff --git a/src/main/runtime/orchestration/coordinator-escalation-triage.ts b/src/main/runtime/orchestration/coordinator-escalation-triage.ts index 5cbfad7f2c9..4c4a195ec92 100644 --- a/src/main/runtime/orchestration/coordinator-escalation-triage.ts +++ b/src/main/runtime/orchestration/coordinator-escalation-triage.ts @@ -9,16 +9,19 @@ export function applyEscalationToDispatch( onLog: (msg: string) => void ): string | null { let taskId: string | undefined + let dispatchId: string | undefined if (msg.payload) { try { const payload = JSON.parse(msg.payload) taskId = payload.taskId + dispatchId = payload.dispatchId } catch { // Escalation without structured payload — log subject as context } } - if (!taskId) { + if (!taskId || !dispatchId) { + onLog(`Rejected escalation from ${msg.from_handle}: missing exact Dispatch binding`) return null } @@ -27,8 +30,9 @@ export function applyEscalationToDispatch( return null } - const dispatch = db.getDispatchContext(taskId) - if (!dispatch) { + const dispatch = db.getDispatchContextById(dispatchId) + if (!dispatch || dispatch.task_id !== taskId) { + onLog(`Rejected escalation from ${msg.from_handle}: Dispatch does not own Task ${taskId}`) return null } if ( diff --git a/src/main/runtime/orchestration/coordinator.test.ts b/src/main/runtime/orchestration/coordinator.test.ts index 5bcbaf28ad4..7fcc3483be1 100644 --- a/src/main/runtime/orchestration/coordinator.test.ts +++ b/src/main/runtime/orchestration/coordinator.test.ts @@ -336,12 +336,13 @@ describe('Coordinator', () => { setTimeout(r, 100) }) const dispatch = db.getDispatchContext(task.id) + expect(dispatch).toBeDefined() db.insertMessage({ from: dispatch?.assignee_handle ?? 'missing-worker', to: 'coord', subject: `Failed attempt ${i + 1}`, type: 'escalation', - payload: JSON.stringify({ taskId: task.id }) + payload: JSON.stringify({ taskId: task.id, dispatchId: dispatch!.id }) }) } @@ -393,6 +394,8 @@ describe('Coordinator', () => { }) // Worker sends decision gate + const dispatch = db.getDispatchContext(task.id) + expect(dispatch).toBeDefined() db.insertMessage({ from: 'term_a', to: 'coord', @@ -400,6 +403,7 @@ describe('Coordinator', () => { type: 'decision_gate', payload: JSON.stringify({ taskId: task.id, + dispatchId: dispatch!.id, question: 'Proceed with destructive migration?', options: ['yes', 'no'] }) diff --git a/src/main/runtime/orchestration/db.ts b/src/main/runtime/orchestration/db.ts index b6a8f3e3dab..68a3c64f173 100644 --- a/src/main/runtime/orchestration/db.ts +++ b/src/main/runtime/orchestration/db.ts @@ -7831,7 +7831,7 @@ export class OrchestrationDb { taskId: string question: string options?: string[] - requester?: { handle: string; paneKey?: string | null } + requester?: { handle: string; paneKey?: string | null; dispatchId: string } }): DecisionGateRow { this.db.exec('SAVEPOINT create_gate') try { @@ -7845,6 +7845,7 @@ export class OrchestrationDb { if ( gate.requester && (!active || + active.id !== gate.requester.dispatchId || !this.isDispatchMessageSender({ dispatchId: active.id, handle: gate.requester.handle, diff --git a/src/main/runtime/orchestration/dispatch-message-binding.ts b/src/main/runtime/orchestration/dispatch-message-binding.ts new file mode 100644 index 00000000000..2646e346af7 --- /dev/null +++ b/src/main/runtime/orchestration/dispatch-message-binding.ts @@ -0,0 +1,23 @@ +import type { MessageType } from './types' + +export function bindCoordinatorMutationPayload( + type: MessageType, + payload: string | null | undefined, + dispatchId: string +): string | undefined { + if (type !== 'escalation' && type !== 'decision_gate') { + return payload ?? undefined + } + if (!payload) { + return JSON.stringify({ dispatchId }) + } + try { + const parsed: unknown = JSON.parse(payload) + if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) { + return payload + } + return JSON.stringify({ ...parsed, dispatchId }) + } catch { + return payload + } +} diff --git a/src/main/runtime/orchestration/federation-sync.test.ts b/src/main/runtime/orchestration/federation-sync.test.ts index b5a57e0394f..55372924d57 100644 --- a/src/main/runtime/orchestration/federation-sync.test.ts +++ b/src/main/runtime/orchestration/federation-sync.test.ts @@ -116,6 +116,74 @@ describe('federation relay parsing', () => { parseRelayedMessage(JSON.stringify({ subject: 'bad', body: 'Blocked', type: 'invented' })) ).toThrowError('Federated relay message type invented is not supported.') }) + + it.each(['escalation', 'decision_gate'] as const)( + 'binds an old remote %s payload to the imported Dispatch', + async (type) => { + const db = new OrchestrationDb(':memory:') + const run = db.createRun({ + objective: 'Federated mutation binding', + coordinatorHandle: 'term_coordinator', + coordinatorPaneKey: 'tab_coordinator:11111111-1111-4111-8111-111111111111' + }) + const task = db.createTask({ spec: 'Remote work', runId: run.id }) + const { dispatch } = db.createStartingWorkerDispatch({ + taskId: task.id, + startOptions: {}, + federation: { + environmentId: 'environment_windows', + environmentName: 'windows', + peerFingerprint: 'windows_peer_fingerprint', + protocolVersion: 3 + } + }) + db.recordWorkerStage({ dispatchId: dispatch.id, stage: 'ready', state: 'ready' }) + const runtime = new OrcaRuntimeService() + runtime.setOrchestrationDb(db) + vi.spyOn(runtime, 'resolveOrchestrationWorkerServer').mockReturnValue({ + peerFingerprint: 'windows_peer_fingerprint' + } as never) + vi.spyOn(runtime, 'callOrchestrationWorkerServer').mockImplementation( + async (_environmentId, method) => { + if (method === 'orchestration.federationPull') { + return { + runtimeEpoch: 'remote_epoch_1', + items: [ + { + dispatch_id: dispatch.id, + direction: 'to_home', + sequence: 1, + message_id: `msg_remote_${type}`, + kind: type, + payload: JSON.stringify({ + subject: 'Remote control mutation', + body: '', + type, + payload: JSON.stringify({ + taskId: task.id, + ...(type === 'decision_gate' ? { question: 'Proceed?' } : {}) + }) + }) + } + ] + } + } + if (method === 'orchestration.federationAck') { + return { acknowledgedThrough: 1 } + } + throw new Error(`Unexpected method ${method}`) + } + ) + + await syncFederatedDispatch(runtime, dispatch.id) + + expect(JSON.parse(db.getMessageById(`msg_remote_${type}`)!.payload!)).toMatchObject({ + taskId: task.id, + dispatchId: dispatch.id + }) + db.close() + } + ) }) describe('federation relay acknowledgments', () => { diff --git a/src/main/runtime/orchestration/federation-sync.ts b/src/main/runtime/orchestration/federation-sync.ts index 3bc7f0f5180..673d8c62002 100644 --- a/src/main/runtime/orchestration/federation-sync.ts +++ b/src/main/runtime/orchestration/federation-sync.ts @@ -14,6 +14,7 @@ import { recordFederationAckCheckpoint } from './federation-ack-checkpoints' import { parseFederatedWorkerReportPayload } from './federation-worker-report-payload' +import { bindCoordinatorMutationPayload } from './dispatch-message-binding' const MESSAGE_TYPE_SET = new Set(MESSAGE_TYPES) const FEDERATION_PULL_PAGE_SIZE = 50 @@ -112,7 +113,7 @@ async function syncFederatedDispatchPages( type: message.type, priority: message.priority, threadId: message.threadId ?? undefined, - payload: message.payload ?? undefined + payload: bindCoordinatorMutationPayload(message.type, message.payload, dispatchId) }, lifecycle: parseFederatedLifecycle(message, item.message_id, dispatchId, dispatch.task_id) }) diff --git a/src/main/runtime/rpc/methods/orchestration-send-dispatch-authority.test.ts b/src/main/runtime/rpc/methods/orchestration-send-dispatch-authority.test.ts index f08ce857a23..cbba56c1c2a 100644 --- a/src/main/runtime/rpc/methods/orchestration-send-dispatch-authority.test.ts +++ b/src/main/runtime/rpc/methods/orchestration-send-dispatch-authority.test.ts @@ -2,6 +2,8 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import type { RpcContext } from '../core' import type { OrchestrationDb } from '../../orchestration/db' import type { OrcaRuntimeService } from '../../orca-runtime' +import { openDecisionGateFromMessage } from '../../orchestration/coordinator-decision-gates' +import { applyEscalationToDispatch } from '../../orchestration/coordinator-escalation-triage' import { createOrchestrationRpcHarness } from './orchestration-rpc-test-harness' describe('orchestration.send Dispatch authority', () => { @@ -115,4 +117,37 @@ describe('orchestration.send Dispatch authority', () => { expect(db.getTask(task.id)?.status).toBe('dispatched') } ) + + it.each(['escalation', 'decision_gate'] as const)( + 'binds queued legacy %s mail to its exact Dispatch before handle reuse', + async (type) => { + setup() + const task = db.createTask({ spec: 'legacy re-dispatch target' }) + const first = db.createDispatchContext(task.id, 'term_legacy') + + const sent = (await send({ + from: 'term_legacy', + type, + subject: 'Queued legacy control', + payload: JSON.stringify({ + taskId: task.id, + ...(type === 'decision_gate' ? { question: 'Proceed?' } : {}) + }) + })) as { message: { id: string; payload: string } } + + expect(JSON.parse(sent.message.payload)).toMatchObject({ dispatchId: first.id }) + db.failDispatch(first.id, 'worker stopped before coordinator read its mail') + const second = db.createDispatchContext(task.id, 'term_legacy') + + if (type === 'escalation') { + applyEscalationToDispatch(db, db.getMessageById(sent.message.id)!, () => {}) + } else { + openDecisionGateFromMessage(db, db.getMessageById(sent.message.id)!, () => {}) + } + + expect(db.getTask(task.id)?.status).toBe('dispatched') + expect(db.getDispatchContextById(second.id)?.status).toBe('dispatched') + expect(db.listGates({ taskId: task.id })).toHaveLength(0) + } + ) }) diff --git a/src/main/runtime/rpc/methods/orchestration.ts b/src/main/runtime/rpc/methods/orchestration.ts index 16a5099c6c8..908673ec7a5 100644 --- a/src/main/runtime/rpc/methods/orchestration.ts +++ b/src/main/runtime/rpc/methods/orchestration.ts @@ -31,6 +31,7 @@ import { OrchestrationError } from '../../orchestration/orchestration-error' import type { OrcaRuntimeService } from '../../orca-runtime' import type { RunRow } from '../../orchestration/types' import { encodeFederatedControlMessage } from '../../orchestration/federation-control-message' +import { bindCoordinatorMutationPayload } from '../../orchestration/dispatch-message-binding' import { ORCHESTRATION_FEDERATION_CONTROL_MAIL_PROTOCOL_VERSION, ORCHESTRATION_FEDERATION_LIFECYCLE_SETTLEMENT_PROTOCOL_VERSION @@ -511,7 +512,11 @@ export const ORCHESTRATION_METHODS: RpcMethod[] = [ type, priority: params.priority ?? 'normal', threadId: params.threadId ?? null, - payload: params.payload ?? null + payload: bindCoordinatorMutationPayload( + type, + params.payload, + remoteAttachment.dispatch_id + ) }), ...(!supportsLifecycleSettlement && outcome ? { settleRemoteOutcome: outcome } : {}) }) @@ -644,15 +649,21 @@ export const ORCHESTRATION_METHODS: RpcMethod[] = [ } // Point-to-point — existing single-recipient behavior revalidateLegacyCoordinator?.() + const dispatch = routing.dispatchId + ? db.getDispatchContextById(routing.dispatchId) + : undefined + const messageType = (params.type ?? 'status') as MessageType const msg = db.insertMessage({ from, to, subject: params.subject, body: params.body, - type: params.type as MessageType, + type: messageType, priority: params.priority as MessagePriority, threadId: params.threadId, - payload: params.payload, + payload: dispatch + ? bindCoordinatorMutationPayload(messageType, params.payload, dispatch.id) + : params.payload, senderPaneKey, runId: routing.run?.id, deliveryContract: legacyWorkerDeliveryContract( @@ -661,9 +672,6 @@ export const ORCHESTRATION_METHODS: RpcMethod[] = [ to ) }) - const dispatch = routing.dispatchId - ? db.getDispatchContextById(routing.dispatchId) - : undefined const dispatchMutationMessage = isDispatchMutationMessageType(msg.type) if (dispatchMutationMessage) { const processIncarnation =