342 lines
11 KiB
TypeScript
342 lines
11 KiB
TypeScript
import { type AssistantMessageFrame, createModels } from "@earendil-works/pi-ai";
|
|
import { afterEach, describe, expect, it, vi } from "vitest";
|
|
import type { WatchHandle } from "../../../src/harness/agent-harness.ts";
|
|
import { DEFAULT_COMPACTION_SETTINGS } from "../../../src/harness/compaction/compaction.ts";
|
|
import { BACKGROUND_CONTEXT, type Context, withAbortSignal } from "../../../src/harness/context.ts";
|
|
import { HookRegistry } from "../../../src/harness/hooks.ts";
|
|
import { Lane } from "../../../src/harness/runtime/lane.ts";
|
|
import { openFrameProgress, openToolProgress } from "../../../src/harness/runtime/progress.ts";
|
|
import { type Config, Drive, type LaneState as RuntimeLaneState } from "../../../src/harness/runtime/types.ts";
|
|
import { MemoryStorage } from "../../../src/harness/session/memory.ts";
|
|
import { StorageBackedSession } from "../../../src/harness/session/session.ts";
|
|
import {
|
|
type AssistantEffectPendingOperation,
|
|
type CheckpointOperation,
|
|
type OperationState,
|
|
operationScopeOf,
|
|
type Session,
|
|
type Write,
|
|
} from "../../../src/harness/session/types.ts";
|
|
import * as storedValues from "../../../src/harness/session/values.ts";
|
|
import { deferred } from "./test-utils.ts";
|
|
|
|
const sessions: Session[] = [];
|
|
const operationId = "operation";
|
|
const configuration = {
|
|
model: { provider: "test", modelId: "model" },
|
|
thinkingLevel: "off" as const,
|
|
activeToolNames: [],
|
|
};
|
|
|
|
class FailingStorage extends MemoryStorage {
|
|
failure: Error | undefined;
|
|
beforeNextCommit: (() => Promise<void>) | undefined;
|
|
|
|
override async commit(writes: Write[], context: Context) {
|
|
const beforeCommit = this.beforeNextCommit;
|
|
this.beforeNextCommit = undefined;
|
|
await beforeCommit?.();
|
|
if (this.failure === undefined) return super.commit(writes, context);
|
|
const failure = this.failure;
|
|
this.failure = undefined;
|
|
throw failure;
|
|
}
|
|
}
|
|
|
|
const runtimeConfig: Config<undefined> = {
|
|
tools: [],
|
|
resources: {},
|
|
streamOptions: {},
|
|
retryPolicy: { enabled: true, maxRetries: 3, baseDelayMs: 1_000 },
|
|
compaction: DEFAULT_COMPACTION_SETTINGS,
|
|
steeringMode: "all",
|
|
followUpMode: "all",
|
|
toolExecution: "parallel",
|
|
toolContext: undefined,
|
|
systemPrompt: undefined,
|
|
toProviderMessages: () => [],
|
|
entryProjectors: {},
|
|
};
|
|
|
|
function unusedWatch<T>(): WatchHandle<T> {
|
|
throw new Error("watch is not used by progress tests");
|
|
}
|
|
|
|
function runScope() {
|
|
return {
|
|
control: { status: "running" as const },
|
|
settings: {
|
|
compaction: DEFAULT_COMPACTION_SETTINGS,
|
|
steeringMode: "all" as const,
|
|
followUpMode: "all" as const,
|
|
toolExecution: "parallel" as const,
|
|
},
|
|
inbox: { steer: [], followUp: [], writes: [] },
|
|
latestAssistantEntryId: null,
|
|
};
|
|
}
|
|
|
|
function assistantEffectPending(responseEntryId: string): AssistantEffectPendingOperation {
|
|
return {
|
|
...runScope(),
|
|
at: "assistant.effect_pending",
|
|
generationContext: {
|
|
stepId: "step",
|
|
triggerEntryId: "trigger",
|
|
configuration,
|
|
streamOptions: {},
|
|
retryPolicy: { maxAttempts: 2, baseDelayMs: 1 },
|
|
overflowRecoveryUsed: false,
|
|
},
|
|
attempt: 1,
|
|
responseEntryId,
|
|
usageId: "usage",
|
|
intendedOutputLimit: 100,
|
|
contextWindow: 1_000,
|
|
};
|
|
}
|
|
|
|
async function createFixture(
|
|
state: OperationState,
|
|
storage = new FailingStorage(),
|
|
): Promise<{ lane: Lane<undefined>; drive: Drive; storage: FailingStorage }> {
|
|
const session = new StorageBackedSession(
|
|
{ id: `progress-${sessions.length}`, createdAt: 1, storageVersion: 1 },
|
|
storage,
|
|
);
|
|
sessions.push(session);
|
|
const meta = {
|
|
operationId,
|
|
lane: "main",
|
|
sourceTipId: null,
|
|
startedAt: 1,
|
|
intent: { kind: "run" as const, promptEntryIds: [] },
|
|
};
|
|
const projection: RuntimeLaneState = {
|
|
tipId: null,
|
|
configuration,
|
|
inbox: [],
|
|
lastOperationId: null,
|
|
operation: { meta, state },
|
|
};
|
|
const lane = new Lane<undefined>(
|
|
"main",
|
|
session,
|
|
createModels(),
|
|
new HookRegistry(() => {}),
|
|
projection,
|
|
(cause) => (cause instanceof Error ? cause : new Error(String(cause))),
|
|
() => Promise.resolve(),
|
|
unusedWatch,
|
|
() => runtimeConfig,
|
|
);
|
|
await session.mutate(
|
|
(mutator) =>
|
|
mutator.commit(
|
|
[
|
|
storedValues.setValue(storedValues.branchTip("main"), null),
|
|
storedValues.setValue(storedValues.laneConfig("main"), configuration),
|
|
storedValues.setValue(storedValues.laneState("main"), {
|
|
currentOperationId: operationId,
|
|
lastOperationId: null,
|
|
inbox: [],
|
|
}),
|
|
storedValues.setValue(storedValues.operationMeta(operationId), meta),
|
|
storedValues.setValue(storedValues.operationState(operationId), state),
|
|
],
|
|
BACKGROUND_CONTEXT,
|
|
),
|
|
BACKGROUND_CONTEXT,
|
|
);
|
|
const drive = new Drive({ operationId }, BACKGROUND_CONTEXT);
|
|
lane.activeDrive = drive;
|
|
return { lane, drive, storage };
|
|
}
|
|
|
|
afterEach(async () => {
|
|
for (const session of sessions.splice(0)) await session.close(BACKGROUND_CONTEXT);
|
|
});
|
|
|
|
describe("runtime active drive ownership", () => {
|
|
it("settles or rejects the shared completion without exposing owner controls through the gate", async () => {
|
|
const settled = new Drive({ operationId }, BACKGROUND_CONTEXT);
|
|
const outcome = { kind: "waiting", operationId, reason: "retry", notBefore: 10 } as const;
|
|
settled.settle(outcome);
|
|
settled.fail(new Error("late failure"));
|
|
await expect(settled.completion).resolves.toBe(outcome);
|
|
|
|
const failed = new Drive({ operationId }, BACKGROUND_CONTEXT);
|
|
const error = new Error("drive failed");
|
|
failed.fail(error);
|
|
await expect(failed.completion).rejects.toBe(error);
|
|
expect("beginAbort" in failed.gate).toBe(false);
|
|
});
|
|
|
|
it("strips invocation cancellation from pass context and owns policy and gate control", () => {
|
|
const controller = new AbortController();
|
|
const drive = new Drive(
|
|
{ operationId, waitForRetry: true, pollDeferred: true },
|
|
withAbortSignal(controller.signal, BACKGROUND_CONTEXT),
|
|
);
|
|
|
|
expect(drive.operationId).toBe(operationId);
|
|
expect(drive.context.abortSignal).toBeUndefined();
|
|
expect(drive.waitForRetry).toBe(true);
|
|
expect(drive.deferredPermits).toBe(1);
|
|
|
|
const closed = new Error("closed");
|
|
drive.closeGate(closed);
|
|
expect(drive.gate.signal.aborted).toBe(true);
|
|
expect(() => drive.gate.admit(() => undefined)).toThrow(closed);
|
|
});
|
|
});
|
|
|
|
describe("runtime progress channels", () => {
|
|
it("enqueues assistant frames in order, seals admission, and exposes the settlement delete", async () => {
|
|
const responseEntryId = "response";
|
|
const { lane, drive } = await createFixture(assistantEffectPending(responseEntryId));
|
|
const progress = openFrameProgress(lane, drive, responseEntryId);
|
|
const frames: AssistantMessageFrame[] = [
|
|
{ type: "text_delta", contentIndex: 0, delta: "a" },
|
|
{ type: "text_delta", contentIndex: 0, delta: "b" },
|
|
];
|
|
|
|
progress.write(frames[0]!);
|
|
progress.write(frames[1]!);
|
|
progress.seal();
|
|
progress.write({ type: "text_delta", contentIndex: 0, delta: "late" });
|
|
await progress.drain();
|
|
|
|
expect(
|
|
(
|
|
await lane.session.readList(
|
|
storedValues.pendingAssistantFrames(operationId, responseEntryId),
|
|
undefined,
|
|
BACKGROUND_CONTEXT,
|
|
)
|
|
).map(({ value }) => value),
|
|
).toEqual(frames);
|
|
});
|
|
|
|
it("declines a queued frame after the authoritative projection leaves its phase without control reads", async () => {
|
|
const responseEntryId = "response";
|
|
const { lane, drive, storage } = await createFixture(assistantEffectPending(responseEntryId));
|
|
const progress = openFrameProgress(lane, drive, responseEntryId);
|
|
const commitStarted = deferred();
|
|
const releaseCommit = deferred();
|
|
storage.beforeNextCommit = async () => {
|
|
commitStarted.resolve();
|
|
await releaseCommit.promise;
|
|
};
|
|
const getValue = vi.spyOn(storage, "getValue");
|
|
const moving = lane.command((state) => {
|
|
const operation = state.operation;
|
|
if (operation === null) throw new Error("missing operation");
|
|
const nextState: CheckpointOperation = {
|
|
...operationScopeOf(operation.state),
|
|
at: "checkpoint",
|
|
continuation: { kind: "may_finish", includeFinalAssistant: true },
|
|
triggerEntryId: responseEntryId,
|
|
};
|
|
return {
|
|
kind: "commit",
|
|
writes: [storedValues.setValue(storedValues.operationState(operationId), nextState)],
|
|
next: { ...state, operation: { meta: operation.meta, state: nextState } },
|
|
materialize: () => undefined,
|
|
};
|
|
}, BACKGROUND_CONTEXT);
|
|
await commitStarted.promise;
|
|
progress.write({ type: "text_delta", contentIndex: 0, delta: "late" });
|
|
releaseCommit.resolve();
|
|
await moving;
|
|
await progress.drain();
|
|
|
|
expect(
|
|
await lane.session.readList(
|
|
storedValues.pendingAssistantFrames(operationId, responseEntryId),
|
|
undefined,
|
|
BACKGROUND_CONTEXT,
|
|
),
|
|
).toEqual([]);
|
|
expect(getValue).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("declines a queued frame after terminal projection publication", async () => {
|
|
const responseEntryId = "response";
|
|
const { lane, drive, storage } = await createFixture(assistantEffectPending(responseEntryId));
|
|
const progress = openFrameProgress(lane, drive, responseEntryId);
|
|
const commitStarted = deferred();
|
|
const releaseCommit = deferred();
|
|
storage.beforeNextCommit = async () => {
|
|
commitStarted.resolve();
|
|
await releaseCommit.promise;
|
|
};
|
|
const getValue = vi.spyOn(storage, "getValue");
|
|
const ending = lane.command(
|
|
(state) => ({
|
|
kind: "commit",
|
|
writes: [
|
|
storedValues.deleteValue(storedValues.operationMeta(operationId)),
|
|
storedValues.deleteValue(storedValues.operationState(operationId)),
|
|
storedValues.setValue(storedValues.laneState("main"), {
|
|
currentOperationId: null,
|
|
lastOperationId: null,
|
|
inbox: state.inbox,
|
|
}),
|
|
],
|
|
next: { ...state, operation: null },
|
|
materialize: () => undefined,
|
|
}),
|
|
BACKGROUND_CONTEXT,
|
|
);
|
|
await commitStarted.promise;
|
|
progress.write({ type: "text_delta", contentIndex: 0, delta: "late" });
|
|
releaseCommit.resolve();
|
|
await ending;
|
|
await progress.drain();
|
|
|
|
expect(
|
|
await lane.session.readList(
|
|
storedValues.pendingAssistantFrames(operationId, responseEntryId),
|
|
undefined,
|
|
BACKGROUND_CONTEXT,
|
|
),
|
|
).toEqual([]);
|
|
expect(getValue).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it("replaces tool checkpoints in invocation order", async () => {
|
|
const invocationId = "result";
|
|
const { lane, drive } = await createFixture({
|
|
...runScope(),
|
|
at: "tools",
|
|
batch: {
|
|
assistantEntryId: "assistant",
|
|
configuration,
|
|
turnId: "turn",
|
|
calls: [{ status: "effect_pending", sourceIndex: 0, resultEntryId: invocationId, replay: "safe" }],
|
|
},
|
|
} satisfies OperationState);
|
|
const progress = openToolProgress(lane, drive, "turn", 0, invocationId);
|
|
progress.write({ content: [{ type: "text", text: "first" }], details: {} });
|
|
progress.write({ content: [{ type: "text", text: "second" }], details: {} });
|
|
await progress.drain();
|
|
expect(
|
|
(await lane.session.getValue(storedValues.pendingToolOutput(operationId, invocationId), BACKGROUND_CONTEXT))
|
|
?.value.content,
|
|
).toEqual([{ type: "text", text: "second" }]);
|
|
});
|
|
|
|
it("retains the rejecting write promise so drain propagates commit failure", async () => {
|
|
const responseEntryId = "response";
|
|
const { lane, drive, storage } = await createFixture(assistantEffectPending(responseEntryId));
|
|
const failure = new Error("frame commit failed");
|
|
storage.failure = failure;
|
|
const progress = openFrameProgress(lane, drive, responseEntryId);
|
|
|
|
progress.write({ type: "text_delta", contentIndex: 0, delta: "lost" });
|
|
|
|
await expect(progress.drain()).rejects.toBe(failure);
|
|
});
|
|
});
|