Files
harness-tools/extensions/pi-workflow/context-transition.test.ts
T

850 lines
28 KiB
TypeScript

import assert from "node:assert/strict";
import { createHash } from "node:crypto";
import { readFileSync } from "node:fs";
import { mkdtemp, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import * as path from "node:path";
import test from "node:test";
import type {
ExtensionAPI,
ExtensionCommandContext,
ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import type { ImplementationContextState, WorkflowRun } from "./domain.ts";
import {
registerWorkflowContextTransitions,
type ContextTransitionRuntime,
} from "./context-transition.ts";
import type { WorkflowCommand, WorkflowEngine, WorkflowView } from "./engine.ts";
async function withTempDir(run: (directory: string) => Promise<void>) {
const directory = await mkdtemp(path.join(tmpdir(), "pi-workflow-context-"));
try {
await run(directory);
} finally {
await rm(directory, { recursive: true, force: true });
}
}
type CommandHandler = (
args: string,
ctx: ExtensionCommandContext,
) => Promise<void>;
type ShutdownHandler = (
event: { reason: string; targetSessionFile?: string },
ctx: ExtensionContext,
) => Promise<void>;
type StartHandler = (
event: { reason: string; previousSessionFile?: string },
ctx: ExtensionContext,
) => Promise<void>;
function fakeApi() {
let command: CommandHandler | undefined;
let shutdown: ShutdownHandler | undefined;
let start: StartHandler | undefined;
const commandNames: string[] = [];
const userMessages: string[] = [];
const api = {
registerCommand(name: string, options: { handler: CommandHandler }) {
commandNames.push(name);
command = options.handler;
},
sendUserMessage(message: string) {
userMessages.push(message);
},
on(event: string, handler: ShutdownHandler | StartHandler) {
if (event === "session_shutdown") shutdown = handler as ShutdownHandler;
if (event === "session_start") start = handler as StartHandler;
},
} as unknown as ExtensionAPI;
return {
api,
commandNames,
userMessages,
command: () => {
assert.ok(command);
return command;
},
shutdown: () => {
assert.ok(shutdown);
return shutdown;
},
start: () => {
assert.ok(start);
return start;
},
};
}
function workflowView(
planPath: string,
revision: number,
context: ImplementationContextState,
): WorkflowView {
return {
runId: "run-1",
revision,
status: context.status === "ready" ? "running" : "waiting",
background: false,
workspace: {
state: "ready",
sourceRoot: "/repo",
root: "/repo/.worktrees/run-1",
branch: "feature/example",
baseHead: "abc123",
},
openGate:
context.status === "ready"
? undefined
: {
id: context.gateId,
kind: "implementation-context",
status: "open",
revision: 10,
openedAt: 1,
},
planArtifact: {
path: planPath,
revision: 1,
digest: createHash("sha256").update(readFileSync(planPath)).digest("hex"),
},
implementationContext: structuredClone(context),
stages: [],
assignments: [],
mutatorProfileNames: [],
reviewCycles: 0,
verificationCycles: 0,
uiImpactEvidence: [],
affectedUiSurfaces: [],
};
}
function commandContext(options: {
sessionId: string;
sessionFile: string;
compact?: ExtensionCommandContext["compact"];
newSession?: ExtensionCommandContext["newSession"];
mode?: ExtensionCommandContext["mode"];
hasUI?: boolean;
input?: (title: string, placeholder?: string) => Promise<string | undefined>;
select?: (title: string, options: string[]) => Promise<string | undefined>;
}) {
const notifications: string[] = [];
const context = {
cwd: "/repo",
sessionManager: {
getSessionId: () => options.sessionId,
getSessionFile: () => options.sessionFile,
},
mode: options.mode ?? "print",
hasUI: options.hasUI ?? false,
ui: {
notify(message: string) {
notifications.push(message);
},
input: options.input ?? (async () => undefined),
select: options.select ?? (async () => undefined),
},
async waitForIdle() {},
compact: options.compact ?? (() => {}),
newSession:
options.newSession ??
(async () => ({ cancelled: false })) as ExtensionCommandContext["newSession"],
} as unknown as ExtensionCommandContext;
return { context, notifications };
}
test("/workflow routes non-continuation arguments to the dashboard", async () => {
const registered = fakeApi();
const opened: string[] = [];
registerWorkflowContextTransitions(
registered.api,
{
store: {
async listRuns() {
return [];
},
},
engineForContext: () => {
throw new Error("dashboard should not create an engine");
},
} as unknown as ContextTransitionRuntime,
async (args) => {
opened.push(args);
return {};
},
);
await registered.command()("run-1", commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
}).context);
assert.deepEqual(opened, ["run-1"]);
});
test("/workflow start routes the exact request through the assistant", async () => {
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: { listRuns: async () => [] },
engineForContext: () => {
throw new Error("start should route through the assistant");
},
} as unknown as ContextTransitionRuntime);
await registered.command()(
"start Build searchable workflow cards",
commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
}).context,
);
assert.equal(registered.userMessages.length, 1);
assert.match(registered.userMessages[0]!, /invoke pi_workflow/);
assert.ok(
registered.userMessages[0]!.endsWith("\n\nBuild searchable workflow cards"),
);
});
test("the misspelled /worlflow alias is not registered", () => {
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: { listRuns: async () => [] },
engineForContext: () => {
throw new Error("command registration should not create an engine");
},
} as unknown as ContextTransitionRuntime);
assert.deepEqual(registered.commandNames, ["workflow"]);
});
test("/workflow start prompts in the TUI and requires a message otherwise", async () => {
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: { listRuns: async () => [] },
engineForContext: () => {
throw new Error("start should route through the assistant");
},
} as unknown as ContextTransitionRuntime);
let prompt = "";
await registered.command()(
"start",
commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
mode: "tui",
hasUI: true,
input: async (title) => {
prompt = title;
return "Add workflow lifecycle commands";
},
}).context,
);
assert.equal(prompt, "What do you want to start working on?");
assert.ok(
registered.userMessages[0]!.endsWith("\n\nAdd workflow lifecycle commands"),
);
const fallback = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
});
await registered.command()("start", fallback.context);
assert.ok(fallback.notifications.some((message) => message.includes("required")));
});
test("/workflow continue lists only valid runs and resume is an alias", async () => {
const validRun = (runId: string, activeParentSessionId = "session-1") =>
({
runId,
revision: 10,
status: "waiting",
request: `Request ${runId}`,
parentSessionIds: [activeParentSessionId],
lease: { activeParentSessionId },
planArtifact: { path: "/plan.md", revision: 1, digest: "digest" },
implementationContext: {
gateId: `gate-${runId}`,
runRevision: 10,
planRevision: 1,
preparationId: `prep-${runId}`,
activeParentSessionId,
strategy: "fresh-session",
status: "handoff-pending",
method: "fresh-session",
},
gates: [
{
id: `gate-${runId}`,
kind: "implementation-context",
status: "open",
revision: 10,
openedAt: 1,
},
],
assignments: [],
}) as unknown as WorkflowRun;
const selectedTitles: string[] = [];
const selectedOptions: string[][] = [];
const inspected: string[] = [];
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: {
listRuns: async () => [
validRun("run-1"),
validRun("run-2"),
validRun("historical", "session-2"),
],
},
engineForContext: () =>
({
async execute(command: WorkflowCommand) {
if (command.action === "inspect") {
inspected.push(command.runId);
return {
runId: command.runId,
revision: 1,
status: "running",
background: false,
workspace: { state: "ready", sourceRoot: "/repo" },
stages: [],
assignments: [],
mutatorProfileNames: [],
reviewCycles: 0,
verificationCycles: 0,
uiImpactEvidence: [],
affectedUiSurfaces: [],
} as WorkflowView;
}
throw new Error(`Unexpected command ${command.action}`);
},
}) as WorkflowEngine,
} as unknown as ContextTransitionRuntime);
const interactive = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
mode: "tui",
hasUI: true,
select: async (title, options) => {
selectedTitles.push(title);
selectedOptions.push(options);
return undefined;
},
});
await registered.command()("continue", interactive.context);
await registered.command()("resume", interactive.context);
assert.deepEqual(selectedTitles, [
"Continue Workflow Run",
"Continue Workflow Run",
]);
assert.equal(selectedOptions[0]?.length, 2);
assert.equal(selectedOptions[1]?.length, 2);
assert.ok(
selectedOptions.every((options) =>
options.every((label) => !label.includes("historical")),
),
);
await registered.command()("resume run-2", interactive.context);
assert.deepEqual(inspected, ["run-2"]);
});
test("/workflow continue without a run ID gives actionable non-UI guidance", async () => {
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: { listRuns: async () => [] },
engineForContext: () => {
throw new Error("non-UI fallback should not create an engine");
},
} as unknown as ContextTransitionRuntime);
const command = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
});
await registered.command()("continue", command.context);
assert.ok(
command.notifications.some((message) =>
message.includes("/workflow continue <run-id>"),
),
);
});
test("/workflow stop selects among active runs owned by the current session", async () => {
const run = (
runId: string,
status: WorkflowRun["status"],
activeParentSessionId: string,
) =>
({
runId,
revision: 4,
status,
request: `Request ${runId}`,
parentSessionIds: [activeParentSessionId],
lease: { activeParentSessionId },
}) as unknown as WorkflowRun;
const commands: WorkflowCommand[] = [];
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: {
listRuns: async () => [
run("run-1", "running", "session-1"),
run("run-2", "waiting", "session-1"),
run("completed", "completed", "session-1"),
run("historical", "running", "session-2"),
],
},
engineForContext: () =>
({
async execute(command: WorkflowCommand) {
commands.push(command);
if (command.action !== "inspect" && command.action !== "abort") {
throw new Error(`Unexpected command ${command.action}`);
}
return {
runId: command.runId,
revision: 4,
status: command.action === "abort" ? "aborted" : "waiting",
} as WorkflowView;
},
}) as WorkflowEngine,
} as unknown as ContextTransitionRuntime);
let choices: string[] = [];
const command = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
mode: "tui",
hasUI: true,
select: async (title, options) => {
assert.equal(title, "Stop Workflow Run");
choices = options;
return options[1];
},
});
await registered.command()("stop", command.context);
assert.equal(choices.length, 2);
assert.ok(choices.every((label) => !label.includes("completed")));
assert.deepEqual(
commands.map((entry) =>
`${entry.action}:${"runId" in entry ? entry.runId : ""}`,
),
["inspect:run-2", "abort:run-2"],
);
const abort = commands[1];
assert.match(
abort?.action === "abort" ? abort.idempotencyKey : "",
/^workflow-stop-run-2-4-/,
);
});
test("/workflow stop immediately stops the only active run", async () => {
const commands: WorkflowCommand[] = [];
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: {
listRuns: async () => [
{
runId: "run-1",
revision: 2,
status: "running",
request: "Only active run",
lease: { activeParentSessionId: "session-1" },
} as unknown as WorkflowRun,
],
},
engineForContext: () =>
({
async execute(command: WorkflowCommand) {
commands.push(command);
if (command.action !== "inspect" && command.action !== "abort") {
throw new Error(`Unexpected command ${command.action}`);
}
return {
runId: command.runId,
revision: 2,
status: command.action === "abort" ? "aborted" : "running",
} as WorkflowView;
},
}) as WorkflowEngine,
} as unknown as ContextTransitionRuntime);
const command = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
});
await registered.command()("stop", command.context);
assert.deepEqual(
commands.map((entry) => entry.action),
["inspect", "abort"],
);
assert.ok(command.notifications.some((message) => message.includes("stopped")));
});
test("/workflow stop requires a run ID for multiple active runs without UI", async () => {
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: {
listRuns: async () =>
["run-1", "run-2"].map(
(runId) =>
({
runId,
status: "running",
request: runId,
lease: { activeParentSessionId: "session-1" },
}) as unknown as WorkflowRun,
),
},
engineForContext: () => {
throw new Error("selection fallback should not create an engine");
},
} as unknown as ContextTransitionRuntime);
const command = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
});
await registered.command()("stop", command.context);
assert.ok(
command.notifications.some((message) =>
message.includes("/workflow stop <run-id>"),
),
);
});
test("/workflow continue invokes and fences normal compaction callbacks", async () => {
await withTempDir(async (directory) => {
const planPath = path.join(directory, "plan.md");
await writeFile(planPath, "# Approved plan\n\nImplement safely.", "utf8");
const context: ImplementationContextState = {
gateId: "gate-1",
runRevision: 10,
planRevision: 1,
preparationId: "prep-1",
activeParentSessionId: "session-1",
sourceParentSessionId: "session-1",
strategy: "smart-compact",
status: "normal-compact-pending",
method: "smart_compact",
outcome: "unavailable",
};
let view = workflowView(planPath, 10, context);
const commands: WorkflowCommand[] = [];
const engine = {
async execute(command: WorkflowCommand) {
commands.push(structuredClone(command));
if (command.action === "inspect") return structuredClone(view);
if (command.action === "continue-context") {
const next = structuredClone(context);
next.status = "normal-compact-running";
next.method = "normal-compact";
next.runRevision = 11;
view = workflowView(planPath, 11, next);
return structuredClone(view);
}
if (command.action === "complete-normal-compaction") {
if (command.expectedRevision !== view.revision) {
throw new Error("stale compaction callback");
}
const next = structuredClone(view.implementationContext!);
next.status = command.outcome === "failed" ? "failed" : "ready";
next.outcome = command.outcome;
view = workflowView(planPath, 12, next);
return structuredClone(view);
}
throw new Error(`Unexpected command ${command.action}`);
},
} as WorkflowEngine;
const runtime = {
store: { listRuns: async () => [] },
engineForContext: () => engine,
} as unknown as ContextTransitionRuntime;
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, runtime);
let compactOptions: Parameters<ExtensionCommandContext["compact"]>[0];
const { context: ctx, notifications } = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
compact(options) {
compactOptions = options;
},
});
await registered.command()("continue run-1", ctx);
assert.match(compactOptions?.customInstructions ?? "", /Approved plan/);
compactOptions?.onComplete?.({
summary: "summary",
firstKeptEntryId: "entry-1",
tokensBefore: 100,
});
await new Promise((resolve) => setTimeout(resolve, 0));
assert.deepEqual(
commands.map((command) => command.action),
["inspect", "continue-context", "complete-normal-compaction"],
);
const completed = commands.at(-1);
assert.equal(
completed?.action === "complete-normal-compaction"
? completed.outcome
: undefined,
"succeeded",
);
compactOptions?.onComplete?.({
summary: "duplicate",
firstKeptEntryId: "entry-1",
tokensBefore: 100,
});
await new Promise((resolve) => setTimeout(resolve, 0));
assert.ok(notifications.some((message) => message.includes("stale compaction callback")));
});
});
test("cancelled fresh-session replacement clears the source-bound intent", async () => {
await withTempDir(async (directory) => {
const planPath = path.join(directory, "plan.md");
await writeFile(planPath, "# Plan", "utf8");
const context: ImplementationContextState = {
gateId: "gate-1",
runRevision: 10,
planRevision: 1,
preparationId: "prep-1",
activeParentSessionId: "session-1",
sourceParentSessionId: "session-1",
strategy: "fresh-session",
status: "handoff-pending",
method: "fresh-session",
};
let view = workflowView(planPath, 10, context);
const commands: WorkflowCommand[] = [];
const engine = {
async execute(command: WorkflowCommand) {
commands.push(structuredClone(command));
if (command.action === "inspect") return structuredClone(view);
if (command.action === "continue-context") {
const next = structuredClone(view.implementationContext!);
next.sourceSessionFile = command.sourceSessionFile;
next.runRevision = 11;
view = workflowView(planPath, 11, next);
return structuredClone(view);
}
if (command.action === "recover-context-intent") {
const next = structuredClone(view.implementationContext!);
next.sourceSessionFile = undefined;
next.intentCreatedAt = undefined;
next.intentExpiresAt = undefined;
next.runRevision = 12;
view = workflowView(planPath, 12, next);
return structuredClone(view);
}
throw new Error(`Unexpected command ${command.action}`);
},
} as WorkflowEngine;
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, {
store: { listRuns: async () => [] },
engineForContext: () => engine,
} as unknown as ContextTransitionRuntime);
const command = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
newSession: (async () => ({ cancelled: true })) as ExtensionCommandContext["newSession"],
});
await registered.command()("continue run-1", command.context);
assert.deepEqual(
commands.map((entry) => entry.action),
["inspect", "continue-context", "recover-context-intent"],
);
assert.equal(view.implementationContext?.sourceSessionFile, undefined);
});
});
test("a later startup recovers an unconsumed claimed handoff", async () => {
await withTempDir(async (directory) => {
const planPath = path.join(directory, "plan.md");
await writeFile(planPath, "# Plan", "utf8");
const context: ImplementationContextState = {
gateId: "gate-1",
runRevision: 30,
planRevision: 1,
preparationId: "prep-1",
activeParentSessionId: "session-1",
sourceParentSessionId: "session-1",
sourceSessionFile: "/sessions/source.jsonl",
targetSessionFile: "/sessions/target.jsonl",
strategy: "fresh-session",
status: "claimed-target-awaiting-withSession",
method: "fresh-session",
intentExpiresAt: Date.now() + 60_000,
};
let view = workflowView(planPath, 30, context);
const commands: WorkflowCommand[] = [];
const engine = {
async execute(command: WorkflowCommand) {
commands.push(structuredClone(command));
if (command.action === "recover-context-intent") {
const next = structuredClone(view.implementationContext!);
next.status = "handoff-pending";
next.targetSessionFile = undefined;
next.runRevision = 31;
view = workflowView(planPath, 31, next);
return structuredClone(view);
}
throw new Error(`Unexpected command ${command.action}`);
},
} as WorkflowEngine;
const runtime = {
store: {
async listRuns() {
return [
{
runId: view.runId,
revision: view.revision,
status: view.status,
implementationContext: structuredClone(view.implementationContext),
} as WorkflowRun,
];
},
},
engineForContext: () => engine,
} as unknown as ContextTransitionRuntime;
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, runtime);
const target = commandContext({
sessionId: "session-2",
sessionFile: "/sessions/target.jsonl",
}).context;
for (const scenario of [
{ eventReason: "startup", expectedReason: "startup", expired: false },
{ eventReason: "reload", expectedReason: "reload", expired: false },
{ eventReason: "resume", expectedReason: "resume", expired: false },
{ eventReason: "fork", expectedReason: "fork", expired: false },
{ eventReason: "startup", expectedReason: "expired", expired: true },
] as const) {
const pending = structuredClone(context);
pending.intentExpiresAt = scenario.expired
? Date.now() - 1
: Date.now() + 60_000;
view = workflowView(planPath, 30, pending);
commands.length = 0;
await registered.start()({ reason: scenario.eventReason }, target);
const recovered = commands[0];
assert.equal(recovered?.action, "recover-context-intent");
assert.equal(
recovered?.action === "recover-context-intent"
? recovered.reason
: undefined,
scenario.expectedReason,
);
assert.equal(view.implementationContext?.status, "handoff-pending");
}
});
});
test("fresh continuation claims shutdown target and consumes only in withSession", async () => {
await withTempDir(async (directory) => {
const planPath = path.join(directory, "plan.md");
const planPrefix = "# Handoff plan\n";
const boundaryPlan = `${planPrefix}${"a".repeat(
32 * 1024 - 1 - Buffer.byteLength(planPrefix, "utf8"),
)}😀tail`;
await writeFile(planPath, boundaryPlan, "utf8");
const context: ImplementationContextState = {
gateId: "gate-1",
runRevision: 20,
planRevision: 1,
preparationId: "prep-1",
activeParentSessionId: "session-1",
sourceParentSessionId: "session-1",
strategy: "fresh-session",
status: "handoff-pending",
method: "fresh-session",
};
let view = workflowView(planPath, 20, context);
const commands: WorkflowCommand[] = [];
const sent: unknown[] = [];
const engine = {
async execute(command: WorkflowCommand) {
commands.push(structuredClone(command));
if (command.action === "inspect") return structuredClone(view);
if (command.action === "continue-context") {
const next = structuredClone(view.implementationContext!);
next.sourceSessionFile = command.sourceSessionFile;
next.runRevision = 21;
view = workflowView(planPath, 21, next);
return structuredClone(view);
}
if (command.action === "claim-context-handoff") {
const next = structuredClone(view.implementationContext!);
next.status = "claimed-target-awaiting-withSession";
next.targetSessionFile = command.targetSessionFile;
next.runRevision = 22;
view = workflowView(planPath, 22, next);
return structuredClone(view);
}
if (command.action === "consume-context-handoff") {
const next = structuredClone(view.implementationContext!);
next.status = "ready";
next.activeParentSessionId = "session-2";
next.runRevision = 23;
view = workflowView(planPath, 23, next);
return structuredClone(view);
}
throw new Error(`Unexpected command ${command.action}`);
},
} as WorkflowEngine;
const runtime = {
store: {
async listRuns() {
return [
{
runId: view.runId,
revision: view.revision,
status: view.status,
implementationContext: structuredClone(view.implementationContext),
} as WorkflowRun,
];
},
},
engineForContext: () => engine,
} as unknown as ContextTransitionRuntime;
const registered = fakeApi();
registerWorkflowContextTransitions(registered.api, runtime);
const replacement = commandContext({
sessionId: "session-2",
sessionFile: "/sessions/target.jsonl",
}).context as unknown as ExtensionCommandContext & {
sendMessage(message: unknown): Promise<void>;
sendUserMessage(message: string): Promise<void>;
};
replacement.sendMessage = async (message) => {
sent.push(message);
};
replacement.sendUserMessage = async () => {};
const old = commandContext({
sessionId: "session-1",
sessionFile: "/sessions/source.jsonl",
newSession: (async (options) => {
await registered.shutdown()(
{ reason: "new", targetSessionFile: "/sessions/target.jsonl" },
old.context,
);
await registered.start()(
{ reason: "new", previousSessionFile: "/sessions/source.jsonl" },
replacement,
);
await options?.withSession?.(replacement);
return { cancelled: false };
}) as ExtensionCommandContext["newSession"],
});
await registered.command()("continue run-1", old.context);
assert.equal(sent.length, 1);
const handoff = JSON.stringify(sent[0]);
assert.match(handoff, /Handoff plan/);
assert.doesNotMatch(handoff, /�|😀/);
assert.deepEqual(
commands.map((command) => command.action),
[
"inspect",
"continue-context",
"claim-context-handoff",
"inspect",
"consume-context-handoff",
],
);
});
});