jb/harness-tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(pi-workflow): improve workflow usability and recovery

Browse Source
This commit is contained in:
jb
2026-08-26 00:20:24 -04:00
parent c291ca083c
commit ea74dbc6a7
23 changed files with 4422 additions and 913 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extensions/pi-workflow/engine.ts
+315 -77
View File
@@ -35,9 +35,9 @@ import {
const FEATURE_NAME = /^[a-z0-9]+(?:-[a-z0-9]+)*$/;
const MAX_FEATURE_NAME_LENGTH = 80;
const DEFAULT_LEASE_STALE_MS = 30_000;
const DEFAULT_ASSIGNMENT_HEARTBEAT_MS = 10_000;
const DEFAULT_ASSIGNMENT_ABORT_TIMEOUT_MS = 15_000;
const DEFAULT_ASSIGNMENT_EXECUTION_TIMEOUT_MS = 5 * 60 * 1_000;
const ACTIVE_ASSIGNMENT_TOKENS = new Set<string>();
interface ActiveAssignmentExecution {
runId: string;
@@ -59,6 +59,7 @@ const ACTIVE_STARTS = new Map<
string,
{
fingerprint: string;
controller: AbortController;
operation: Promise<WorkflowView>;
published: Promise<WorkflowView>;
}
@@ -78,6 +79,7 @@ export interface StartWorkflowCommand {
idempotencyKey: string;
concurrency?: number;
background?: boolean;
allowDirtySource?: boolean;
script?: FeatureScriptInput;
}
@@ -372,10 +374,10 @@ export interface WorkflowEngineOptions {
now?: () => number;
createId?: () => string;
processId?: number;
leaseStaleMs?: number;
isProcessAlive?: (processId: number) => boolean;
assignmentHeartbeatMs?: number;
assignmentAbortTimeoutMs?: number;
assignmentExecutionTimeoutMs?: number;
afterAssignmentPersist?: (
assignment: WorkflowAssignment,
run: WorkflowRun,
@@ -613,7 +615,7 @@ function validateWorkflowCommand(
assertKeys(
value,
["action", "request", "featureName", "idempotencyKey"],
["background", "concurrency", "script"],
["allowDirtySource", "background", "concurrency", "script"],
);
nonEmptyString(value.request, "request");
nonEmptyString(value.featureName, "featureName");
@@ -632,6 +634,12 @@ function validateWorkflowCommand(
) {
throw new Error("background must be a boolean");
}
if (
value.allowDirtySource !== undefined &&
typeof value.allowDirtySource !== "boolean"
) {
throw new Error("allowDirtySource must be a boolean");
}
if (value.script !== undefined) validateFeatureScript(value.script);
return;
}
@@ -1243,22 +1251,76 @@ export class WorkflowEngine {
private readonly now: () => number;
private readonly createId: () => string;
private readonly processId: number;
private readonly leaseStaleMs: number;
private readonly isProcessAlive: (processId: number) => boolean;
private readonly assignmentHeartbeatMs: number;
private readonly assignmentAbortTimeoutMs: number;
private readonly assignmentExecutionTimeoutMs: number;
constructor(options: WorkflowEngineOptions) {
this.options = options;
this.now = options.now ?? Date.now;
this.createId = options.createId ?? randomUUID;
this.processId = options.processId ?? process.pid;
this.leaseStaleMs = options.leaseStaleMs ?? DEFAULT_LEASE_STALE_MS;
this.isProcessAlive = options.isProcessAlive ?? processAlive;
this.assignmentHeartbeatMs =
options.assignmentHeartbeatMs ?? DEFAULT_ASSIGNMENT_HEARTBEAT_MS;
this.assignmentAbortTimeoutMs =
options.assignmentAbortTimeoutMs ?? DEFAULT_ASSIGNMENT_ABORT_TIMEOUT_MS;
this.assignmentExecutionTimeoutMs =
options.assignmentExecutionTimeoutMs ?? DEFAULT_ASSIGNMENT_EXECUTION_TIMEOUT_MS;
if (this.assignmentExecutionTimeoutMs < 1) {
throw new Error("Assignment execution timeout must be positive");
}
}
async cancelStart(
idempotencyKey: string,
abortIdempotencyKey: string,
reason: Error,
): Promise<void> {
const activeKey = `${this.options.store.root}\0${idempotencyKey}`;
const active = ACTIVE_STARTS.get(activeKey);
if (!active) {
const start = await this.options.store.getStart(idempotencyKey);
if (!start) return;
const terminal = await this.abortCancelledStart(
start.runId,
abortIdempotencyKey,
);
if (terminal) {
await this.options.store.completeStart(idempotencyKey, terminal);
}
return;
}
active.controller.abort(reason);
const published = await Promise.race([
active.published.then(
(result) => result,
() => undefined,
),
active.operation.then(
() => undefined,
() => undefined,
),
]);
let terminal = published
? await this.abortCancelledStart(published.runId, abortIdempotencyKey)
: undefined;
await active.operation.catch(() => {});
if (!terminal) {
const start = await this.options.store.getStart(idempotencyKey);
if (start) {
terminal = await this.abortCancelledStart(
start.runId,
abortIdempotencyKey,
);
}
}
if (terminal) {
await this.options.store.completeStart(idempotencyKey, terminal);
}
}
async execute(command: WorkflowCommand): Promise<WorkflowView> {
@@ -1317,9 +1379,10 @@ export class WorkflowEngine {
}
if (
run.implementationContext?.status === "normal-compact-running" &&
this.leaseIsStale(run)
this.ownershipIsRecoverable(run)
) {
run = await this.transition(run, (next) => {
this.rotateOwnership(next);
const context = next.implementationContext;
if (!context) throw new Error("Implementation Context state is missing");
context.status = "failed";
@@ -1335,22 +1398,21 @@ export class WorkflowEngine {
const running = run.assignments.filter(
(assignment) =>
assignment.status === "running" ||
(assignment.status === "queued" && Boolean(assignment.executionToken)),
((assignment.status === "queued" || assignment.status === "interrupted") &&
Boolean(assignment.executionToken)),
);
const alreadyWaiting = activeGate(run)?.kind === "recovery";
if (running.length === 0 && alreadyWaiting) continue;
if (
running.length === 0 &&
alreadyWaiting &&
!this.ownershipIsRecoverable(run)
) {
continue;
}
if (await this.hasActiveAssignmentLease(run.runId, running)) continue;
if (!this.leaseIsStale(run)) continue;
if (!this.ownershipIsRecoverable(run)) continue;
run = await this.transition(run, (next) => {
next.lease = {
processId: this.processId,
activeParentSessionId: this.options.parentSessionId,
heartbeatAt: this.now(),
worktreeIdentity: next.workspace.root,
};
next.parentSessionIds = [
...new Set([...next.parentSessionIds, this.options.parentSessionId]),
];
this.rotateOwnership(next);
});
if (running.length === 0) {
const claimedRevision = run.revision;
@@ -1374,15 +1436,7 @@ export class WorkflowEngine {
for (const assignment of next.assignments) {
if (interruptedIds.has(assignment.id)) assignment.status = "interrupted";
}
next.lease = {
processId: this.processId,
activeParentSessionId: this.options.parentSessionId,
heartbeatAt: this.now(),
worktreeIdentity: next.workspace.root,
};
next.parentSessionIds = [
...new Set([...next.parentSessionIds, this.options.parentSessionId]),
];
this.rotateOwnership(next);
next.status = "waiting";
if (!activeGate(next)) {
const gate: WorkflowGate = {
@@ -1428,11 +1482,13 @@ export class WorkflowEngine {
resolvePublished = resolve;
rejectPublished = reject;
});
const operation = this.startOwned(command, resolvePublished);
const controller = new AbortController();
const operation = this.startOwned(command, controller.signal, resolvePublished);
void operation.catch(rejectPublished);
void published.catch(() => {});
ACTIVE_STARTS.set(activeKey, {
fingerprint: commandFingerprint,
controller,
operation,
published,
});
@@ -1467,11 +1523,25 @@ export class WorkflowEngine {
}
return active.operation;
}
const operation = this.startAnalysisOwned(command);
let resolvePublished = (_view: WorkflowView) => {};
let rejectPublished = (_error: unknown) => {};
const published = new Promise<WorkflowView>((resolve, reject) => {
resolvePublished = resolve;
rejectPublished = reject;
});
const controller = new AbortController();
const operation = this.startAnalysisOwned(
command,
controller.signal,
resolvePublished,
);
void operation.catch(rejectPublished);
void published.catch(() => {});
ACTIVE_STARTS.set(activeKey, {
fingerprint: commandFingerprint,
controller,
operation,
published: operation,
published,
});
void operation
.finally(() => {
@@ -1483,33 +1553,47 @@ export class WorkflowEngine {
return operation;
}
private async startAnalysisOwned(command: StartAnalysisWorkflowCommand) {
private async startAnalysisOwned(
command: StartAnalysisWorkflowCommand,
cancellationSignal: AbortSignal,
onPublished?: (view: WorkflowView) => void,
) {
this.throwIfStartCancelled(cancellationSignal);
if (!this.options.scriptRuntime) {
throw new Error("Analysis scripts require the restricted script runtime");
}
const commandFingerprint = fingerprint(command);
const existing = await this.options.store.getStart(command.idempotencyKey);
this.throwIfStartCancelled(cancellationSignal);
if (existing) {
this.validateExistingStart(existing, commandFingerprint);
if (existing.status === "completed") {
return existing.result as unknown as WorkflowView;
const result = existing.result as unknown as WorkflowView;
onPublished?.(result);
return result;
}
}
const runId = existing?.runId ?? this.createId();
this.throwIfStartCancelled(cancellationSignal);
const reservation = await this.options.store.reserveStart(
command.idempotencyKey,
commandFingerprint,
runId,
);
this.throwIfStartCancelled(cancellationSignal);
this.validateExistingStart(reservation.record, commandFingerprint);
let run: WorkflowRun;
try {
run = await this.options.store.loadRun(reservation.record.runId);
} catch (error) {
if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
this.throwIfStartCancelled(cancellationSignal);
run = this.initialAnalysisRun(reservation.record.runId, command);
this.throwIfStartCancelled(cancellationSignal);
run = (await this.options.store.createRun(run)).run;
}
onPublished?.(view(run));
this.throwIfStartCancelled(cancellationSignal);
if (run.status !== "completed" && run.status !== "failed") {
run = await this.runWorkflowScript(run);
}
@@ -1520,11 +1604,15 @@ export class WorkflowEngine {
private async startOwned(
command: StartWorkflowCommand,
cancellationSignal: AbortSignal,
onPublished?: (view: WorkflowView) => void,
) {
this.throwIfStartCancelled(cancellationSignal);
this.validateStart(command);
this.throwIfStartCancelled(cancellationSignal);
const commandFingerprint = fingerprint(command);
const existing = await this.options.store.getStart(command.idempotencyKey);
this.throwIfStartCancelled(cancellationSignal);
if (existing) {
this.validateExistingStart(existing, commandFingerprint);
if (existing.status === "completed") {
@@ -1546,6 +1634,7 @@ export class WorkflowEngine {
};
const preflightEligibility = this.preflightEligibility(initialFacts);
selectInitialDiscoveryProfile(this.options.profiles, initialFacts);
this.throwIfStartCancelled(cancellationSignal);
let persistedRun: WorkflowRun | undefined;
if (existing?.status === "pending") {
@@ -1556,6 +1645,7 @@ export class WorkflowEngine {
}
}
let runId = existing?.runId ?? this.createId();
this.throwIfStartCancelled(cancellationSignal);
let workspacePlan = persistedRun
? this.workspacePlanForRun(persistedRun)
: await this.options.workspace.prepare({
@@ -1563,11 +1653,14 @@ export class WorkflowEngine {
configuredRoot: this.options.effectiveConfig.worktrees.root,
featureName: command.featureName,
runId,
allowDirtySource: command.allowDirtySource,
});
this.throwIfStartCancelled(cancellationSignal);
if (workspacePlan.branchExists) {
const owner = await this.options.store.findRunByBranch(
workspacePlan.branch,
);
this.throwIfStartCancelled(cancellationSignal);
const resumable =
existing?.status === "pending" && owner?.runId === runId;
const uncertainPending = existing?.status === "pending" && !owner;
@@ -1583,11 +1676,13 @@ export class WorkflowEngine {
}
}
this.throwIfStartCancelled(cancellationSignal);
const reservation = await this.options.store.reserveStart(
command.idempotencyKey,
commandFingerprint,
runId,
);
this.throwIfStartCancelled(cancellationSignal);
this.validateExistingStart(reservation.record, commandFingerprint);
if (reservation.record.runId !== runId) {
runId = reservation.record.runId;
@@ -1597,6 +1692,7 @@ export class WorkflowEngine {
if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
persistedRun = undefined;
}
this.throwIfStartCancelled(cancellationSignal);
workspacePlan = persistedRun
? this.workspacePlanForRun(persistedRun)
: await this.options.workspace.prepare({
@@ -1604,28 +1700,34 @@ export class WorkflowEngine {
configuredRoot: this.options.effectiveConfig.worktrees.root,
featureName: command.featureName,
runId,
allowDirtySource: command.allowDirtySource,
});
this.throwIfStartCancelled(cancellationSignal);
}
{
this.throwIfStartCancelled(cancellationSignal);
let run: WorkflowRun;
let persistedBeforeThisAttempt = true;
try {
run = persistedRun ?? (await this.options.store.loadRun(runId));
} catch (error) {
if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error;
this.throwIfStartCancelled(cancellationSignal);
run = this.initialRun(
runId,
command,
workspacePlan,
preflightEligibility,
);
this.throwIfStartCancelled(cancellationSignal);
const publication = await this.options.store.createRun(run);
persistedBeforeThisAttempt = !publication.created;
run = publication.run;
}
onPublished?.(view(run));
this.throwIfStartCancelled(cancellationSignal);
if (persistedBeforeThisAttempt) {
const claimed = await this.claimPersistedStartRun(run);
@@ -1634,16 +1736,19 @@ export class WorkflowEngine {
}
if (run.workspace.state === "creating" && persistedBeforeThisAttempt) {
this.throwIfStartCancelled(cancellationSignal);
const reconciled = await this.options.workspace.reconcile(
run.workspace,
run.lease?.worktreeIdentity,
);
this.throwIfStartCancelled(cancellationSignal);
if (reconciled.state !== "creating") {
run = await this.persistWorkspaceOutcome(run, reconciled);
}
}
if (run.workspace.state === "creating") {
this.throwIfStartCancelled(cancellationSignal);
let workspace: WorkflowRun["workspace"];
try {
workspace = await this.options.workspace.create(workspacePlan);
@@ -1654,6 +1759,7 @@ export class WorkflowEngine {
error: error instanceof Error ? error.message : String(error),
};
}
this.throwIfStartCancelled(cancellationSignal);
run = await this.persistWorkspaceOutcome(run, workspace);
}
@@ -1661,12 +1767,14 @@ export class WorkflowEngine {
const discoverStatus = stage(run, "discover").status;
const designStatus = stage(run, "design").status;
if (discoverStatus === "pending" || discoverStatus === "running") {
run = await this.runDiscovery(run);
this.throwIfStartCancelled(cancellationSignal);
run = await this.runDiscovery(run, cancellationSignal);
} else if (
discoverStatus === "completed" &&
(designStatus === "pending" || designStatus === "running")
) {
run = await this.runDesign(run);
this.throwIfStartCancelled(cancellationSignal);
run = await this.runDesign(run, false, cancellationSignal);
}
}
@@ -1691,18 +1799,10 @@ export class WorkflowEngine {
) {
return undefined;
}
if (!this.leaseIsStale(run)) return undefined;
if (!this.ownershipIsRecoverable(run)) return undefined;
try {
return await this.transition(run, (next) => {
next.lease = {
processId: this.processId,
activeParentSessionId: this.options.parentSessionId,
heartbeatAt: this.now(),
worktreeIdentity: next.workspace.root,
};
next.parentSessionIds = [
...new Set([...next.parentSessionIds, this.options.parentSessionId]),
];
this.rotateOwnership(next);
});
} catch (error) {
if (
@@ -1787,6 +1887,9 @@ export class WorkflowEngine {
branch: "",
baseHead: "",
branchExists: false,
sourceChanges: [],
sourceChangesOmitted: 0,
dirtySourceAllowed: false,
},
[],
);
@@ -1880,6 +1983,9 @@ export class WorkflowEngine {
root: workspacePlan.worktreePath,
branch: workspacePlan.branch,
baseHead: workspacePlan.baseHead,
sourceChanges: structuredClone(workspacePlan.sourceChanges),
sourceChangesOmitted: workspacePlan.sourceChangesOmitted,
dirtySourceAllowed: workspacePlan.dirtySourceAllowed,
},
effectiveConfig: {
...structuredClone(this.options.effectiveConfig),
@@ -2340,7 +2446,10 @@ export class WorkflowEngine {
}
}
private async runDiscovery(run: WorkflowRun) {
private async runDiscovery(
run: WorkflowRun,
cancellationSignal?: AbortSignal,
) {
const selected = selectionEvidence(
run.pinnedProfiles,
{ ...run.facts, stage: "discover" },
@@ -2350,6 +2459,8 @@ export class WorkflowEngine {
run,
"discover",
selected,
false,
cancellationSignal,
);
run = execution.run;
if (!execution.output) return run;
@@ -2389,7 +2500,9 @@ export class WorkflowEngine {
});
if (output.questions.length > 0) return run;
run = await this.runWorkflowScript(run, "before-design");
return run.status === "failed" ? run : this.runDesign(run);
return run.status === "failed"
? run
: this.runDesign(run, false, cancellationSignal);
}
private plannerCandidates(run: WorkflowRun) {
@@ -2493,7 +2606,11 @@ export class WorkflowEngine {
return planned;
}
private async runDesign(run: WorkflowRun, forceNewAssignment = false) {
private async runDesign(
run: WorkflowRun,
forceNewAssignment = false,
cancellationSignal?: AbortSignal,
) {
const selected = selectionEvidence(
run.pinnedProfiles,
{ ...run.facts, stage: "design" },
@@ -2504,6 +2621,7 @@ export class WorkflowEngine {
"design",
selected,
forceNewAssignment,
cancellationSignal,
);
run = execution.run;
if (!execution.output) return run;
@@ -2606,6 +2724,7 @@ export class WorkflowEngine {
role: "discover" | "design",
selected: ReturnType<typeof selectionEvidence>,
forceNewAssignment = false,
cancellationSignal?: AbortSignal,
) {
const stageId = stage(run, role).id;
const existing = [...run.assignments]
@@ -2620,6 +2739,7 @@ export class WorkflowEngine {
output: existing.status === "succeeded" ? existing.result : undefined,
};
}
if (cancellationSignal) this.throwIfStartCancelled(cancellationSignal);
return this.runAssignment(run, role, selected);
}
@@ -2658,6 +2778,11 @@ export class WorkflowEngine {
const assignmentId = planned?.id ?? this.createId();
const executionToken = this.createId();
const abortController = new AbortController();
let executionTimeout: ReturnType<typeof setTimeout> | undefined;
let executorExecution: Promise<unknown> | undefined;
let executionTimedOut = false;
let retainExecutionFence = false;
let cleanupStarted = false;
let resolveSettled = () => {};
let resolveResumeAfterAbort = () => {};
const activeExecution: ActiveAssignmentExecution = {
@@ -2740,6 +2865,17 @@ export class WorkflowEngine {
assignment.status = "running";
});
};
const cleanupExecutionFence = async () => {
if (cleanupStarted) return;
cleanupStarted = true;
await stopHeartbeat().catch(() => {});
ACTIVE_ASSIGNMENT_TOKENS.delete(executionToken);
ACTIVE_ASSIGNMENT_EXECUTIONS.delete(executionToken);
await this.options.store
.releaseAssignmentLease(run.runId, assignmentId, executionToken)
.catch(() => {});
activeExecution.resolveSettled();
};
if (!this.options.executor.reservesCapacity) await markRunning();
try {
@@ -2748,7 +2884,7 @@ export class WorkflowEngine {
let transcriptArtifact: WorkflowAssignment["transcriptArtifact"];
try {
const evidence = await this.assignmentEvidence(run, role);
const executionResult = await this.options.executor.run({
executorExecution = this.options.executor.run({
runId: run.runId,
role,
request: planned
@@ -2771,6 +2907,31 @@ export class WorkflowEngine {
? { onReserved: markRunning }
: {}),
});
void executorExecution.catch(() => {});
const timeout = new Promise<never>((_resolve, reject) => {
executionTimeout = setTimeout(() => {
executionTimedOut = true;
const timeoutError = new Error(
`Assignment ${role} timed out after ${this.assignmentExecutionTimeoutMs}ms`,
);
abortController.abort(timeoutError);
reject(timeoutError);
}, this.assignmentExecutionTimeoutMs);
});
const settledExecution = executorExecution.then(
(result) => {
if (executionTimeout) clearTimeout(executionTimeout);
return result;
},
(error: unknown) => {
if (executionTimeout) clearTimeout(executionTimeout);
throw error;
},
);
const executionResult = await Promise.race([
settledExecution,
timeout,
]);
metadata = isAssignmentExecutionResult(executionResult)
? executionResult.metadata
: undefined;
@@ -2799,13 +2960,30 @@ export class WorkflowEngine {
: undefined;
activeExecution.metadata = metadata;
activeExecution.transcriptArtifact = transcriptArtifact;
await stopHeartbeat();
if (executionTimedOut && executorExecution) {
let abortGrace: ReturnType<typeof setTimeout> | undefined;
const settledDuringGrace = await Promise.race([
executorExecution.then(
() => true,
() => true,
),
new Promise<false>((resolve) => {
abortGrace = setTimeout(
() => resolve(false),
this.assignmentAbortTimeoutMs,
);
}),
]);
if (abortGrace) clearTimeout(abortGrace);
if (!settledDuringGrace) retainExecutionFence = true;
}
if (!retainExecutionFence) await stopHeartbeat();
await this.assertAssignmentLease(
run.runId,
assignmentId,
executionToken,
);
if (abortController.signal.aborted) {
if (abortController.signal.aborted && !executionTimedOut) {
activeExecution.resolveSettled();
await activeExecution.resumeAfterAbort;
run = await this.options.store.loadRun(run.runId);
@@ -2816,7 +2994,7 @@ export class WorkflowEngine {
(item) => item.id === assignmentId,
);
if (assignment) {
assignment.status = "failed";
assignment.status = executionTimedOut ? "interrupted" : "failed";
assignment.error =
error instanceof Error ? error.message : String(error);
applyExecutionMetadata(assignment, metadata, transcriptArtifact);
@@ -2849,7 +3027,7 @@ export class WorkflowEngine {
} else {
next.status = "failed";
next.error = {
code: "assignment-failed",
code: executionTimedOut ? "assignment-timeout" : "assignment-failed",
message,
stageId: targetStage.id,
assignmentId,
@@ -2881,14 +3059,16 @@ export class WorkflowEngine {
);
return { run, output };
} finally {
if (executionTimeout) clearTimeout(executionTimeout);
abortController.abort();
await stopHeartbeat().catch(() => {});
ACTIVE_ASSIGNMENT_TOKENS.delete(executionToken);
ACTIVE_ASSIGNMENT_EXECUTIONS.delete(executionToken);
await this.options.store
.releaseAssignmentLease(run.runId, assignmentId, executionToken)
.catch(() => {});
activeExecution.resolveSettled();
if (retainExecutionFence && executorExecution) {
void executorExecution.then(
() => cleanupExecutionFence(),
() => cleanupExecutionFence(),
);
} else {
await cleanupExecutionFence();
}
}
}
@@ -2998,6 +3178,23 @@ export class WorkflowEngine {
}
}
validateDecisionPayload(gate, command.decision);
if (
command.decision.kind === "recover" &&
command.decision.resolution !== "inspect"
) {
const fencedAssignments = run.assignments.filter(
(assignment) =>
Boolean(assignment.executionToken) &&
(assignment.status === "queued" ||
assignment.status === "running" ||
assignment.status === "interrupted"),
);
if (await this.hasActiveAssignmentLease(run.runId, fencedAssignments)) {
throw new Error(
"Recovery is blocked while an interrupted Assignment is still active",
);
}
}
if (
command.decision.kind === "recover" &&
command.decision.resolution === "retry" &&
@@ -4025,7 +4222,6 @@ export class WorkflowEngine {
command: RecoverContextIntentCommand,
) {
if (
command.reason !== "expired" &&
run.lease &&
run.lease.processId !== this.processId &&
this.isProcessAlive(run.lease.processId)
@@ -4101,6 +4297,38 @@ export class WorkflowEngine {
return view(run);
}
private throwIfStartCancelled(signal: AbortSignal) {
if (!signal.aborted) return;
throw signal.reason instanceof Error
? signal.reason
: new Error("Workflow start was cancelled");
}
private async abortCancelledStart(
runId: string,
idempotencyKey: string,
): Promise<WorkflowView | undefined> {
while (true) {
const run = await this.options.store.loadRun(runId).catch(() => undefined);
if (!run) return undefined;
if (isTerminalStatus(run.status)) return view(run);
try {
return await this.execute({
action: "abort",
runId,
expectedRevision: run.revision,
idempotencyKey,
});
} catch (error) {
const latest = await this.options.store.loadRun(runId).catch(() => undefined);
if (!latest) return undefined;
if (isTerminalStatus(latest.status)) return view(latest);
if (latest.revision !== run.revision) continue;
throw error;
}
}
}
private async abort(run: WorkflowRun, command: AbortWorkflowCommand) {
const duplicate = this.commandReceipt(run, command);
if (duplicate) {
@@ -4132,7 +4360,9 @@ export class WorkflowEngine {
(assignment) =>
!localAssignmentIds.has(assignment.id) &&
Boolean(assignment.executionToken) &&
(assignment.status === "queued" || assignment.status === "running"),
(assignment.status === "queued" ||
assignment.status === "running" ||
assignment.status === "interrupted"),
);
for (const assignment of externallyOwned) {
const lease = await this.options.store.loadAssignmentLease(
@@ -4429,6 +4659,9 @@ export class WorkflowEngine {
branch: workspace.branch,
baseHead: workspace.baseHead,
branchExists: false,
sourceChanges: structuredClone(workspace.sourceChanges ?? []),
sourceChangesOmitted: workspace.sourceChangesOmitted ?? 0,
dirtySourceAllowed: workspace.dirtySourceAllowed ?? false,
};
}
@@ -4554,19 +4787,30 @@ export class WorkflowEngine {
command.decision.resolution !== "retry" ||
activeGate(run)?.id !== command.gateId ||
!run.parentSessionIds.includes(this.options.parentSessionId) ||
!run.lease ||
this.isProcessAlive(run.lease.processId)
!this.ownershipIsRecoverable(run)
) {
return run;
}
const claimed = structuredClone(run);
claimed.lease = {
this.rotateOwnership(claimed);
return claimed;
}
private rotateOwnership(run: WorkflowRun) {
run.lease = {
processId: this.processId,
activeParentSessionId: this.options.parentSessionId,
heartbeatAt: this.now(),
worktreeIdentity: claimed.workspace.root,
worktreeIdentity: run.workspace.root,
};
return claimed;
if (run.implementationContext) {
run.implementationContext.activeParentSessionId =
this.options.parentSessionId;
run.implementationContext.runRevision = run.revision + 1;
}
run.parentSessionIds = [
...new Set([...run.parentSessionIds, this.options.parentSessionId]),
];
}
private assertActiveParent(run: WorkflowRun) {
@@ -4642,26 +4886,20 @@ export class WorkflowEngine {
lease.processId === this.processId
? ACTIVE_ASSIGNMENT_TOKENS.has(lease.token)
: this.isProcessAlive(lease.processId);
const heartbeatIsFresh =
this.now() - lease.heartbeatAt <= this.leaseStaleMs;
if (ownerIsRunning && (heartbeatIsFresh || assignment.mutates))
return true;
if (ownerIsRunning) return true;
}
return false;
}
private leaseIsStale(run: WorkflowRun) {
private ownershipIsRecoverable(run: WorkflowRun) {
if (!run.lease) return true;
if (
run.lease?.processId === this.processId &&
run.lease.processId === this.processId &&
run.lease.activeParentSessionId === this.options.parentSessionId
) {
return false;
}
return (
!run.lease ||
this.now() - run.lease.heartbeatAt > this.leaseStaleMs ||
!this.isProcessAlive(run.lease.processId)
);
return !this.isProcessAlive(run.lease.processId);
}
private validateStart(command: StartWorkflowCommand) {