Squashed 'third_party/pi/' content from commit d981de122
git-subtree-dir: third_party/pi git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477
This commit is contained in:
341
packages/agent/test/harness/runtime/progress.test.ts
Normal file
341
packages/agent/test/harness/runtime/progress.test.ts
Normal file
@@ -0,0 +1,341 @@
|
||||
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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user