diff --git a/src/main/runtime/rpc/methods/orchestration-federated-attach-receipt.test.ts b/src/main/runtime/rpc/methods/orchestration-federated-attach-receipt.test.ts new file mode 100644 index 00000000000..400232fdffb --- /dev/null +++ b/src/main/runtime/rpc/methods/orchestration-federated-attach-receipt.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from 'vitest' +import { parseRemoteFederatedWorkerStartReceipt } from './orchestration-federated-attach-receipt' + +describe('remote federated worker start receipt', () => { + it.each([ + { name: 'missing runtime epoch', receipt: {} }, + { name: 'non-string runtime epoch', receipt: { runtimeEpoch: 42 } }, + { name: 'missing worktree id', receipt: { runtimeEpoch: 'epoch_remote' } }, + { + name: 'missing terminal handle', + receipt: { + runtimeEpoch: 'epoch_remote', + worktreeId: 'worktree_remote', + terminalHandle: undefined + } + } + ])('rejects a ready receipt with $name', ({ receipt }) => { + expect(() => + parseRemoteFederatedWorkerStartReceipt({ + dispatchId: 'ctx_remote', + state: 'ready', + terminalHandle: 'term_remote', + ...receipt + }) + ).toThrow('invalid ready receipt') + }) + + it('accepts a non-ready receipt without ready-only resources', () => { + expect( + parseRemoteFederatedWorkerStartReceipt({ + dispatchId: 'ctx_remote', + state: 'outcome_unknown' + }) + ).toMatchObject({ dispatchId: 'ctx_remote', state: 'outcome_unknown' }) + }) +}) diff --git a/src/main/runtime/rpc/methods/orchestration-federated-attach-receipt.ts b/src/main/runtime/rpc/methods/orchestration-federated-attach-receipt.ts index 1813a8da7cd..fc62a31788b 100644 --- a/src/main/runtime/rpc/methods/orchestration-federated-attach-receipt.ts +++ b/src/main/runtime/rpc/methods/orchestration-federated-attach-receipt.ts @@ -3,7 +3,7 @@ import type { OrchestrationWorkerLaunchReceipt } from './orchestration-worker-la export type RemoteFederatedWorkerStartReceipt = { dispatchId: string state: string - runtimeEpoch: string + runtimeEpoch?: string worktreeId?: string terminalHandle?: string setup?: { state: string } @@ -13,3 +13,41 @@ export type RemoteFederatedWorkerStartReceipt = { failedStage?: string lastError?: string } + +export type ReadyRemoteFederatedWorkerStartReceipt = RemoteFederatedWorkerStartReceipt & { + state: 'ready' + runtimeEpoch: string + worktreeId: string + terminalHandle: string +} + +export function parseRemoteFederatedWorkerStartReceipt( + value: unknown +): RemoteFederatedWorkerStartReceipt { + if (!isRecord(value) || !isNonEmptyString(value.dispatchId) || !isNonEmptyString(value.state)) { + throw new Error('The worker server returned an invalid attachment receipt.') + } + if ( + value.state === 'ready' && + (!isNonEmptyString(value.runtimeEpoch) || + !isNonEmptyString(value.worktreeId) || + !isNonEmptyString(value.terminalHandle)) + ) { + throw new Error('The worker server returned an invalid ready receipt.') + } + return value as RemoteFederatedWorkerStartReceipt +} + +export function isReadyRemoteFederatedWorkerStartReceipt( + receipt: RemoteFederatedWorkerStartReceipt +): receipt is ReadyRemoteFederatedWorkerStartReceipt { + return receipt.state === 'ready' +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +function isNonEmptyString(value: unknown): value is string { + return typeof value === 'string' && value.length > 0 +} diff --git a/src/main/runtime/rpc/methods/orchestration-federated-worker-start-receipt.test.ts b/src/main/runtime/rpc/methods/orchestration-federated-worker-start-receipt.test.ts new file mode 100644 index 00000000000..270c5f6a322 --- /dev/null +++ b/src/main/runtime/rpc/methods/orchestration-federated-worker-start-receipt.test.ts @@ -0,0 +1,84 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { + ORCHESTRATION_CONTRACT_RUNTIME_CAPABILITY, + ORCHESTRATION_FEDERATION_RUNTIME_CAPABILITY +} from '../../../../shared/protocol-version' +import { OrcaRuntimeService } from '../../orca-runtime' +import { OrchestrationDb } from '../../orchestration/db' +import { startFederatedWorker } from './orchestration-federated-worker-start' + +describe('federated worker start receipt validation', () => { + const databases: OrchestrationDb[] = [] + + afterEach(() => { + for (const database of databases.splice(0)) { + database.close() + } + }) + + it('marks a malformed ready receipt outcome unknown without persisting resources', async () => { + const db = new OrchestrationDb(':memory:') + const runtime = new OrcaRuntimeService() + runtime.setOrchestrationDb(db) + databases.push(db) + const run = db.createRun({ + objective: 'federated worker', + coordinatorHandle: 'term_coord', + coordinatorPaneKey: 'tab_coord:leaf_coord' + }) + const task = db.createTask({ spec: 'remote work', runId: run.id }) + vi.spyOn(runtime, 'resolveOrchestrationWorkerServer').mockReturnValue({ + environmentId: 'environment_remote', + name: 'remote', + peerFingerprint: 'remote_peer' + }) + vi.spyOn(runtime, 'callOrchestrationWorkerServer').mockImplementation( + async (_environmentId, method, params) => { + if (method === 'status.get') { + return { + capabilities: [ + ORCHESTRATION_CONTRACT_RUNTIME_CAPABILITY, + ORCHESTRATION_FEDERATION_RUNTIME_CAPABILITY + ] + } + } + return { + dispatchId: (params as { dispatchId: string }).dispatchId, + state: 'ready', + worktreeId: 'worktree_remote', + terminalHandle: 'term_remote' + } + } + ) + + const result = (await startFederatedWorker({ + params: { + task: task.id, + from: 'term_coord', + on: 'remote', + worktree: 'id:worktree_remote', + terminal: 'term_remote' + }, + runtime, + db, + runId: run.id, + task, + orchestrationMutation: { + callerFingerprint: 'caller', + requestId: 'remote_start', + method: 'orchestration.workerStart', + payloadHash: 'payload' + } + })) as { dispatchId: string; state: string; lastError?: string } + + expect(result).toMatchObject({ + state: 'outcome_unknown', + lastError: 'The worker server returned an invalid ready receipt.' + }) + expect(db.getFederatedDispatch(result.dispatchId)).toMatchObject({ + remote_runtime_epoch: null, + remote_worktree_id: null, + remote_terminal_handle: null + }) + }) +}) diff --git a/src/main/runtime/rpc/methods/orchestration-federated-worker-start.ts b/src/main/runtime/rpc/methods/orchestration-federated-worker-start.ts index a7d8bfc8b94..ebf5db1da5e 100644 --- a/src/main/runtime/rpc/methods/orchestration-federated-worker-start.ts +++ b/src/main/runtime/rpc/methods/orchestration-federated-worker-start.ts @@ -23,7 +23,10 @@ import { import { validateFederatedWorkerStartPlacement } from './orchestration-worker-start-validation' import { resolveFederatedWorkerStartBudgets } from './orchestration-worker-start-budgets' import { resolveDispatchCreator } from './orchestration-dispatch-creator' -import type { RemoteFederatedWorkerStartReceipt } from './orchestration-federated-attach-receipt' +import { + isReadyRemoteFederatedWorkerStartReceipt, + parseRemoteFederatedWorkerStartReceipt +} from './orchestration-federated-attach-receipt' export async function startFederatedWorker(args: { params: WorkerStartInput @@ -133,40 +136,42 @@ export async function startFederatedWorker(args: { }) db.recordWorkerStage({ dispatchId: started.dispatch.id, stage: 'remote_attach_requested' }) try { - const remote = (await runtime.callOrchestrationWorkerServer( - server.environmentId, - 'orchestration.federationAttachStart', - { - dispatchId: started.dispatch.id, - taskId: task.id, - taskSpec: task.spec, - // Carry the home dispatch depth across the federation boundary so a - // remote worker cannot be mistaken for a root when it dispatches again. - depth: started.dispatch.depth, - protocolVersion: federationProtocolVersion, - worktree, - name: params.name, - repo: params.repo, - baseBranch: params.baseBranch, - displayName: params.displayName, - comment: params.comment, - setup: createsWorktree ? (params.setup ?? 'run') : undefined, - setupSource: createsWorktree - ? params.setup - ? 'explicit_request' - : 'orchestration_default' - : undefined, - terminal: params.terminal, - agent: params.agent, - model: params.model, - effort: params.effort, - timeoutMs: params.timeoutMs, - devMode: params.devMode - }, - budgets.attachDeadlineMs, - { orchestrationRequestId: orchestrationMutation.requestId }, - { contractVerified: true } - )) as RemoteFederatedWorkerStartReceipt + const remote = parseRemoteFederatedWorkerStartReceipt( + await runtime.callOrchestrationWorkerServer( + server.environmentId, + 'orchestration.federationAttachStart', + { + dispatchId: started.dispatch.id, + taskId: task.id, + taskSpec: task.spec, + // Carry the home dispatch depth across the federation boundary so a + // remote worker cannot be mistaken for a root when it dispatches again. + depth: started.dispatch.depth, + protocolVersion: federationProtocolVersion, + worktree, + name: params.name, + repo: params.repo, + baseBranch: params.baseBranch, + displayName: params.displayName, + comment: params.comment, + setup: createsWorktree ? (params.setup ?? 'run') : undefined, + setupSource: createsWorktree + ? params.setup + ? 'explicit_request' + : 'orchestration_default' + : undefined, + terminal: params.terminal, + agent: params.agent, + model: params.model, + effort: params.effort, + timeoutMs: params.timeoutMs, + devMode: params.devMode + }, + budgets.attachDeadlineMs, + { orchestrationRequestId: orchestrationMutation.requestId }, + { contractVerified: true } + ) + ) if (remote.dispatchId !== started.dispatch.id) { throw new OrchestrationError( 'resource_server_mismatch', @@ -178,7 +183,7 @@ export async function startFederatedWorker(args: { requestedLaunch, remote.state === 'ready' ) - if (remote.state === 'ready' && remote.worktreeId && remote.terminalHandle) { + if (isReadyRemoteFederatedWorkerStartReceipt(remote)) { db.updateFederatedDispatchResources({ dispatchId: started.dispatch.id, remoteRuntimeEpoch: remote.runtimeEpoch,