diff --git a/ui/src/features/agent/wizardWorkflowRuntime.ts b/ui/src/features/agent/wizardWorkflowRuntime.ts index ce709321..1010488c 100644 --- a/ui/src/features/agent/wizardWorkflowRuntime.ts +++ b/ui/src/features/agent/wizardWorkflowRuntime.ts @@ -709,6 +709,10 @@ export class WizardWorkflowRuntime { const workflow = this.find(workflowId) const definition = this.definitions.get(workflow.type) if (!definition) return + // `execute()` can yield. `open()` may replace `this.collection` before the + // step writes its checkpoint, so persist must keep the collection that + // still holds this workflow object. + const owner = this.capturePersistOwner() while (workflow.currentStep < workflow.steps.length) { if (workflow.cancelRequested) return const step = workflow.steps[workflow.currentStep] @@ -717,7 +721,7 @@ export class WizardWorkflowRuntime { workflow.state = failureState(workflow) workflow.recoverableError = `Missing workflow step definition ${step.stepId}.` workflow.updatedAt = Date.now() - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) return } if (step.state === 'completed') { @@ -731,7 +735,7 @@ export class WizardWorkflowRuntime { step.state = 'failed' step.error = workflow.recoverableError workflow.updatedAt = Date.now() - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) return } const now = Date.now() @@ -741,7 +745,7 @@ export class WizardWorkflowRuntime { step.error = '' workflow.state = workflow.resumeRequested ? 'retrying' : 'running' workflow.updatedAt = now - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) let result: WizardWorkflowStepResult try { result = await stepDefinition.execute({ @@ -753,7 +757,7 @@ export class WizardWorkflowRuntime { workflow.state = failureState(workflow) workflow.recoverableError = step.error workflow.updatedAt = Date.now() - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) return } step.output = clone(result.output || {}) @@ -771,7 +775,7 @@ export class WizardWorkflowRuntime { step.completedAt = workflow.updatedAt workflow.currentStep += 1 workflow.state = workflow.currentStep >= workflow.steps.length ? 'completed' : 'running' - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) continue } if (result.state === 'awaiting_input') { @@ -781,7 +785,7 @@ export class WizardWorkflowRuntime { step.error = `Step ${step.stepId} requested input without a reason and declared fields.` workflow.state = failureState(workflow) workflow.recoverableError = step.error - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) return } const previousPending = workflow.pendingInput @@ -812,7 +816,7 @@ export class WizardWorkflowRuntime { } step.state = 'awaiting_input' workflow.state = 'awaiting_input' - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) return } if (!step.taskId) { @@ -820,17 +824,17 @@ export class WizardWorkflowRuntime { step.error = `Step ${step.stepId} returned ${result.state} without a canonical task id.` workflow.state = failureState(workflow) workflow.recoverableError = step.error - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) return } step.state = 'waiting' workflow.state = result.state === 'queued' ? 'queued' : 'waiting' - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) return } workflow.state = 'completed' workflow.updatedAt = Date.now() - await this.persistAndEmit(workflow) + await this.persistAndEmit(workflow, owner) } private find(workflowId: string): WizardWorkflowRecord { @@ -843,21 +847,62 @@ export class WizardWorkflowRuntime { return this.openSequence === openSequence && this.workspace === workspace } - private async persistAndEmit(workflow: WizardWorkflowRecord): Promise { - await this.persist() + private capturePersistOwner(): { + workspace: string + collection: WizardWorkflowCollection + openSequence: number + } { + return { + workspace: this.workspace, + collection: this.collection, + openSequence: this.openSequence, + } + } + + /** + * A CAS retry builds a merged snapshot, but `advanceUnlocked` keeps mutating + * the live workflow object captured in `owner.collection`. Adopt remote-only + * siblings into that same collection instead of pointing persist at the clone + * — otherwise the next persist in this advance overwrites those siblings. + */ + private adoptPersistMerge( + owner: { + workspace: string + collection: WizardWorkflowCollection + openSequence: number + }, + merged: WizardWorkflowCollection, + ): void { + const liveById = new Map(owner.collection.workflows.map(workflow => [workflow.workflowId, workflow])) + owner.collection.workflows = merged.workflows.map(workflow => { + const live = liveById.get(workflow.workflowId) + return live && live.updatedAt >= workflow.updatedAt ? live : workflow + }) + owner.collection.revision = merged.revision + if (this.ownsOpen(owner.openSequence, owner.workspace)) { + this.collection = owner.collection + } + } + + private async persistAndEmit( + workflow: WizardWorkflowRecord, + owner = this.capturePersistOwner(), + ): Promise { + await this.persist(owner) const current = this.collection.workflows.find(item => item.workflowId === workflow.workflowId) if (!current || current.workspace !== this.workspace) return this.emit(current) } - private async persist(): Promise { - // Pin the workspace that owns this snapshot. `open()` rebinds - // `this.workspace` immediately, so a conflict retry must not follow the - // live pointer or it will merge the source checkpoint into the destination + private async persist(owner = this.capturePersistOwner()): Promise { + // Pin the workspace and collection that own this snapshot. `open()` + // rebinds both pointers immediately. A conflict retry — or a persist + // that starts after `execute()` yields — must not follow the live + // pointer or it will write the source checkpoint into the destination // store and then poison the in-memory collection. - const targetWorkspace = this.workspace - const openSequence = this.openSequence - let candidate = clone(this.collection) + const targetWorkspace = owner.workspace + const openSequence = owner.openSequence + let candidate = clone(owner.collection) let lastError: unknown // Workflow updates can race with conversation/library autosaves in another // tab. Re-read and merge a few times before surfacing a real persistence @@ -865,6 +910,7 @@ export class WizardWorkflowRuntime { for (let attempt = 0; attempt < 3; attempt += 1) { try { const saved = await this.persistence.save(targetWorkspace, clone(candidate)) + owner.collection.revision = saved.revision if (this.ownsOpen(openSequence, targetWorkspace)) { this.collection.revision = saved.revision } @@ -873,9 +919,8 @@ export class WizardWorkflowRuntime { lastError = error const remote = await this.persistence.load(targetWorkspace) candidate = mergeCollections(candidate, remote) - if (this.ownsOpen(openSequence, targetWorkspace)) { - this.collection = candidate - } + this.adoptPersistMerge(owner, candidate) + candidate = clone(owner.collection) } } throw lastError instanceof Error ? lastError : new Error('Could not persist Wizard workflow.') diff --git a/ui/tests/wizardWorkflowRuntime.test.mjs b/ui/tests/wizardWorkflowRuntime.test.mjs index bde51214..5186c260 100644 --- a/ui/tests/wizardWorkflowRuntime.test.mjs +++ b/ui/tests/wizardWorkflowRuntime.test.mjs @@ -61,6 +61,11 @@ function workspaceMemoryPersistence() { holdSave?.release() holdSave = null }, + injectWorkflow(workspace, workflow) { + const collection = store(workspace) + collection.workflows = [...collection.workflows, clone(workflow)] + collection.revision += 1 + }, } } @@ -377,3 +382,126 @@ test('an in-flight persist stays on the source workspace after open() retargets' assert.equal(source.workflows[0].workspace, 'workspace-a') assert.deepEqual(persistence.savedWorkspaces.slice(-1), ['workspace-a']) }) + +test('a step that finishes after open() still persists to the source workspace', async () => { + const { WizardWorkflowRuntime } = await import('../src/features/agent/wizardWorkflowRuntime.ts') + const persistence = workspaceMemoryPersistence() + let releaseFinish + let notifyBlocked + const finishBlocked = new Promise(resolve => { notifyBlocked = resolve }) + const finishGate = new Promise(resolve => { releaseFinish = resolve }) + + const runtime = new WizardWorkflowRuntime(persistence) + runtime.register({ + type: 'source_job', + steps: [{ + stepId: 'wait', kind: 'wait for task', + async execute() { return { state: 'waiting', taskId: 'task-source' } }, + }, { + stepId: 'finish', kind: 'publish result', + async execute() { + notifyBlocked() + await finishGate + return { state: 'completed', output: { published: true }, outputRefs: ['final.mp4'] } + }, + }], + }) + + await runtime.open('workspace-a') + await runtime.start({ + workflowId: 'workflow-a', type: 'source_job', workspace: 'workspace-a', + userRequest: 'Keep the finished step on A', + }) + await runtime.open('workspace-b') + await runtime.start({ + workflowId: 'workflow-b', type: 'source_job', workspace: 'workspace-b', + userRequest: 'Destination already has its own checkpoint', + }) + await runtime.open('workspace-a') + assert.equal(runtime.get('workflow-a').state, 'waiting') + + const completion = runtime.handleTaskEvent(taskEvent(71, 'task-source', 'completed', ['song.wav'])) + await finishBlocked + await runtime.open('workspace-b') + releaseFinish() + await completion + + const destination = persistence.snapshot('workspace-b') + assert.deepEqual(destination.workflows.map(item => item.workflowId), ['workflow-b']) + assert.equal(destination.workflows[0].state, 'waiting') + assert.equal(destination.workflows[0].workspace, 'workspace-b') + + const source = persistence.snapshot('workspace-a') + assert.deepEqual(source.workflows.map(item => item.workflowId), ['workflow-a']) + assert.equal(source.workflows[0].state, 'completed') + assert.equal(source.workflows[0].workspace, 'workspace-a') + assert.deepEqual(source.workflows[0].outputRefs, ['song.wav', 'final.mp4']) + assert.equal(runtime.get('workflow-a'), undefined) + assert.equal(runtime.get('workflow-b').workspace, 'workspace-b') +}) + +test('a CAS merge during advance keeps sibling workflows on later persists', async () => { + const { WizardWorkflowRuntime } = await import('../src/features/agent/wizardWorkflowRuntime.ts') + const persistence = workspaceMemoryPersistence() + const runtime = new WizardWorkflowRuntime(persistence) + runtime.register({ + type: 'keeper', + steps: [{ + stepId: 'wait', kind: 'wait for task', + async execute() { return { state: 'waiting', taskId: 'task-keep' } }, + }], + }) + runtime.register({ + type: 'two_step', + steps: [{ + stepId: 'first', kind: 'first step', + async execute() { + persistence.injectWorkflow('demo', { + workflowId: 'workflow-other-tab', + type: 'keeper', + workspace: 'demo', + userRequest: 'Added by another tab', + state: 'completed', + currentStep: 1, + steps: [], + resolvedEntityIds: {}, + inputSnapshot: {}, + taskIds: [], + pipelineIds: [], + outputRefs: [], + confirmationScope: [], + processedEventIds: [], + attempts: 0, + createdAt: 1, + updatedAt: Date.now(), + recoverableError: '', + cancelRequested: false, + resumeRequested: false, + pendingInput: null, + }) + return { state: 'completed', outputRefs: ['first.bin'] } + }, + }, { + stepId: 'second', kind: 'second step', + async execute() { return { state: 'completed', outputRefs: ['second.bin'] } }, + }], + }) + + await runtime.open('demo') + await runtime.start({ + workflowId: 'workflow-keep', type: 'keeper', workspace: 'demo', + userRequest: 'Keep the existing checkpoint', + }) + await runtime.start({ + workflowId: 'workflow-advance', type: 'two_step', workspace: 'demo', + userRequest: 'Finish both steps after a sibling merge', + }) + + const storedIds = persistence.snapshot('demo').workflows.map(item => item.workflowId).sort() + assert.deepEqual(storedIds, ['workflow-advance', 'workflow-keep', 'workflow-other-tab']) + assert.equal(runtime.get('workflow-keep').state, 'waiting') + assert.equal(runtime.get('workflow-other-tab').state, 'completed') + const advanced = runtime.get('workflow-advance') + assert.equal(advanced.state, 'completed') + assert.deepEqual(advanced.outputRefs, ['first.bin', 'second.bin']) +})