wehub-resource-sync
469 lines
17 KiB
TypeScript
469 lines
17 KiB
TypeScript
// Import the test harness FIRST — installs the resource catalog so
|
|
// `chat.agent()` calls register their task functions correctly.
|
|
import { mockChatAgent } from "../src/v3/test/index.js";
|
|
|
|
import type { LanguageModelV3StreamPart } from "@ai-sdk/provider";
|
|
import { simulateReadableStream, streamText } from "ai";
|
|
import { MockLanguageModelV3 } from "ai/test";
|
|
import { describe, expect, it, vi } from "vitest";
|
|
import type { RecoveryBootEvent, RecoveryBootResult } from "../src/v3/ai.js";
|
|
import { __setReplaySessionOutTailImplForTests, chat } from "../src/v3/ai.js";
|
|
|
|
// ── Helpers ────────────────────────────────────────────────────────────
|
|
|
|
function userMessage(text: string, id = "u-" + Math.random().toString(36).slice(2)) {
|
|
return {
|
|
id,
|
|
role: "user" as const,
|
|
parts: [{ type: "text" as const, text }],
|
|
};
|
|
}
|
|
|
|
function assistantMessage(text: string, id = "a-" + Math.random().toString(36).slice(2)) {
|
|
return {
|
|
id,
|
|
role: "assistant" as const,
|
|
parts: [{ type: "text" as const, text }],
|
|
};
|
|
}
|
|
|
|
function partialAssistantWithToolCall(id: string, toolCallId: string, toolName: string) {
|
|
return {
|
|
id,
|
|
role: "assistant" as const,
|
|
parts: [
|
|
{
|
|
type: `tool-${toolName}` as const,
|
|
toolCallId,
|
|
state: "input-available" as const,
|
|
input: { q: "search" },
|
|
},
|
|
],
|
|
} as unknown as ReturnType<typeof assistantMessage>;
|
|
}
|
|
|
|
function textStream(text: string) {
|
|
const chunks: LanguageModelV3StreamPart[] = [
|
|
{ type: "text-start", id: "t1" },
|
|
{ type: "text-delta", id: "t1", delta: text },
|
|
{ type: "text-end", id: "t1" },
|
|
{
|
|
type: "finish",
|
|
finishReason: { unified: "stop", raw: "stop" },
|
|
usage: {
|
|
inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined },
|
|
outputTokens: { total: 10, text: 10, reasoning: undefined },
|
|
},
|
|
},
|
|
];
|
|
return simulateReadableStream({ chunks });
|
|
}
|
|
|
|
// ── Tests ──────────────────────────────────────────────────────────────
|
|
|
|
describe("onRecoveryBoot — chat.agent recovery hook", () => {
|
|
it("does NOT fire on a clean continuation with no recovered state", async () => {
|
|
const onRecoveryBoot = vi.fn();
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => ({ stream: textStream("ok") }),
|
|
});
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.no-state",
|
|
onRecoveryBoot,
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "no-state",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
try {
|
|
// Snapshot is empty, no in-flight users, no partial — guard
|
|
// (partialAssistant !== undefined || inFlightUsers.length > 0) is false.
|
|
await harness.sendMessage(userMessage("fresh message"));
|
|
await new Promise((r) => setTimeout(r, 20));
|
|
expect(onRecoveryBoot).not.toHaveBeenCalled();
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("fires when there's a partial assistant and surfaces it on the ctx", async () => {
|
|
const captured: { event?: RecoveryBootEvent<ReturnType<typeof userMessage>> } = {};
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => ({ stream: textStream("recovered") }),
|
|
});
|
|
const partial = partialAssistantWithToolCall("a-orphan", "tc-1", "search");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.partial-fires-hook",
|
|
onRecoveryBoot: async (event) => {
|
|
captured.event = event as never;
|
|
return {};
|
|
},
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "partial-fires-hook",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionOutPartial(partial as never);
|
|
try {
|
|
await harness.sendMessage(userMessage("next user message"));
|
|
await new Promise((r) => setTimeout(r, 20));
|
|
expect(captured.event).toBeDefined();
|
|
expect(captured.event!.partialAssistant?.id).toBe("a-orphan");
|
|
expect(captured.event!.pendingToolCalls).toHaveLength(1);
|
|
expect(captured.event!.pendingToolCalls[0]!.toolCallId).toBe("tc-1");
|
|
expect(captured.event!.pendingToolCalls[0]!.toolName).toBe("search");
|
|
expect(captured.event!.previousRunId).toBe("run_prior");
|
|
expect(captured.event!.cause).toBe("unknown");
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("pendingToolCalls is extracted from the RAW partial (pre-cleanupAbortedParts)", async () => {
|
|
// Real-world scenario: cancel-mid-tool-call. Session.out has tool-call
|
|
// chunks but the tool never returned. cleanupAbortedParts strips the
|
|
// input-available tool part from the partial used for the chain (you
|
|
// don't want orphan tool calls poisoning the model context), but
|
|
// `pendingToolCalls` should still surface what was happening.
|
|
const cleanedPartial = {
|
|
id: "a-orphan",
|
|
role: "assistant" as const,
|
|
parts: [{ type: "text" as const, text: "Let me look that up" }],
|
|
};
|
|
const rawPartial = {
|
|
id: "a-orphan",
|
|
role: "assistant" as const,
|
|
parts: [
|
|
{ type: "text" as const, text: "Let me look that up" },
|
|
{
|
|
type: "tool-search" as const,
|
|
toolCallId: "tc-pending",
|
|
state: "input-available" as const,
|
|
input: { q: "vietnamese pho" },
|
|
},
|
|
],
|
|
} as unknown as typeof cleanedPartial;
|
|
|
|
const captured: { event?: RecoveryBootEvent } = {};
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => ({ stream: textStream("ok") }),
|
|
});
|
|
const u1 = userMessage("buffered", "u-1");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.pending-tool-from-raw",
|
|
onRecoveryBoot: async (event) => {
|
|
captured.event = event;
|
|
return {};
|
|
},
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "pending-tool-from-raw",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionInTail([u1 as never]);
|
|
// Install AFTER mockChatAgent — its constructor sets its own default
|
|
// override that we want to replace for this test.
|
|
__setReplaySessionOutTailImplForTests(
|
|
async () =>
|
|
({
|
|
settled: [],
|
|
partial: cleanedPartial,
|
|
partialRaw: rawPartial,
|
|
}) as never
|
|
);
|
|
try {
|
|
await new Promise((r) => setTimeout(r, 50));
|
|
expect(captured.event).toBeDefined();
|
|
// Cleaned partial → chain (no input-available tool part)
|
|
expect(captured.event!.partialAssistant?.parts).toHaveLength(1);
|
|
// pendingToolCalls → from raw (input-available tool part visible)
|
|
expect(captured.event!.pendingToolCalls).toHaveLength(1);
|
|
expect(captured.event!.pendingToolCalls[0]!.toolCallId).toBe("tc-pending");
|
|
expect(captured.event!.pendingToolCalls[0]!.toolName).toBe("search");
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("does NOT fire when there are in-flight users but no partial (graceful exit path)", async () => {
|
|
// chat.requestUpgrade(), chat.endRun() before processing, and similar
|
|
// graceful exits leave an unacknowledged user on session.in but no
|
|
// partial assistant on session.out. That's not recovery — the next
|
|
// run just dispatches the message normally.
|
|
const onRecoveryBoot = vi.fn();
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => ({ stream: textStream("ok") }),
|
|
});
|
|
const u1 = userMessage("buffered while dead", "u-buffered");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.inflight-users-no-partial",
|
|
onRecoveryBoot,
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "inflight-users-no-partial",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionInTail([u1 as never]);
|
|
try {
|
|
await new Promise((r) => setTimeout(r, 50));
|
|
expect(onRecoveryBoot).not.toHaveBeenCalled();
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("default behavior re-dispatches each in-flight user as a turn", async () => {
|
|
let turnCount = 0;
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => {
|
|
turnCount++;
|
|
return { stream: textStream(`reply ${turnCount}`) };
|
|
},
|
|
});
|
|
const u1 = userMessage("first buffered", "u-1");
|
|
const u2 = userMessage("second buffered", "u-2");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.default-dispatch",
|
|
// NO onRecoveryBoot — exercise the default path
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "default-dispatch",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionInTail([u1 as never, u2 as never]);
|
|
try {
|
|
await new Promise((r) => setTimeout(r, 100));
|
|
expect(turnCount).toBe(2);
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("smart default: partial + first user spliced into chain, rest dispatched", async () => {
|
|
let observedChain: Array<{ role: string; idHead: string }> = [];
|
|
let turnCount = 0;
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => {
|
|
turnCount++;
|
|
return { stream: textStream("ok") };
|
|
},
|
|
});
|
|
const partial = assistantMessage("partial answer in progress", "a-partial");
|
|
const u1 = userMessage("original question", "u-1");
|
|
const u2 = userMessage("follow-up", "u-2");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.smart-default",
|
|
// NO onRecoveryBoot — exercise the smart default
|
|
onTurnStart: async ({ uiMessages }) => {
|
|
if (turnCount === 0) {
|
|
observedChain = uiMessages.map((m) => ({
|
|
role: m.role,
|
|
idHead: m.id.slice(0, 10),
|
|
}));
|
|
}
|
|
},
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "smart-default",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionOutPartial(partial as never);
|
|
harness.seedSessionInTail([u1 as never, u2 as never]);
|
|
try {
|
|
await new Promise((r) => setTimeout(r, 100));
|
|
// Turn 1 fires with the follow-up user (u2). Its chain should
|
|
// include [u1 (original), a-partial, u2 (follow-up)].
|
|
expect(turnCount).toBe(1);
|
|
expect(observedChain.map((m) => m.role)).toEqual(["user", "assistant", "user"]);
|
|
expect(observedChain[0]!.idHead).toBe("u-1");
|
|
expect(observedChain[1]!.idHead).toBe("a-partial");
|
|
expect(observedChain[2]!.idHead).toBe("u-2");
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("hook's recoveredTurns: [] suppresses re-dispatch of in-flight users", async () => {
|
|
let turnCount = 0;
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => {
|
|
turnCount++;
|
|
return { stream: textStream(`reply ${turnCount}`) };
|
|
},
|
|
});
|
|
const partial = assistantMessage("partial answer", "a-partial");
|
|
const u1 = userMessage("buffered", "u-1");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.suppress-dispatch",
|
|
onRecoveryBoot: async (): Promise<RecoveryBootResult> => ({ recoveredTurns: [] }),
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "suppress-dispatch",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionOutPartial(partial as never);
|
|
harness.seedSessionInTail([u1 as never]);
|
|
try {
|
|
// No turn should fire from the boot-injected queue.
|
|
// Send a fresh user message to confirm the agent is alive.
|
|
await harness.sendMessage(userMessage("real next message"));
|
|
await new Promise((r) => setTimeout(r, 20));
|
|
expect(turnCount).toBe(1); // only the explicit sendMessage turn
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("hook's chain override seeds the accumulator", async () => {
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => ({ stream: textStream("acked") }),
|
|
});
|
|
const custom = assistantMessage("custom-recovered-history", "a-custom");
|
|
const partial = assistantMessage("partial", "a-partial");
|
|
const u1 = userMessage("buffered", "u-1");
|
|
let observedMessageCount = 0;
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.chain-override",
|
|
onRecoveryBoot: async (): Promise<RecoveryBootResult> => ({
|
|
chain: [custom as never],
|
|
recoveredTurns: [u1 as never],
|
|
}),
|
|
onTurnStart: async ({ uiMessages }) => {
|
|
observedMessageCount = uiMessages.length;
|
|
},
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "chain-override",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionOutPartial(partial as never);
|
|
harness.seedSessionInTail([u1 as never]);
|
|
try {
|
|
await new Promise((r) => setTimeout(r, 50));
|
|
// Chain seeded with [custom] before the recovered user message
|
|
// arrives — onTurnStart sees [custom, u1] when the first
|
|
// recovered turn fires.
|
|
expect(observedMessageCount).toBe(2);
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("does NOT fire when hydrateMessages is registered", async () => {
|
|
const onRecoveryBoot = vi.fn();
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => ({ stream: textStream("ok") }),
|
|
});
|
|
const u1 = userMessage("buffered", "u-1");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.hydrate-skips",
|
|
hydrateMessages: async ({ incomingMessages }) => incomingMessages,
|
|
onRecoveryBoot,
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "hydrate-skips",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionInTail([u1 as never]);
|
|
try {
|
|
await harness.sendMessage(userMessage("fresh"));
|
|
await new Promise((r) => setTimeout(r, 20));
|
|
expect(onRecoveryBoot).not.toHaveBeenCalled();
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("beforeBoot runs before the first recovered turn fires", async () => {
|
|
const order: string[] = [];
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => {
|
|
order.push("turn");
|
|
return { stream: textStream("ok") };
|
|
},
|
|
});
|
|
const partial = assistantMessage("partial", "a-partial");
|
|
const u1 = userMessage("buffered original", "u-1");
|
|
const u2 = userMessage("followup", "u-2");
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.before-boot",
|
|
onRecoveryBoot: async (): Promise<RecoveryBootResult> => ({
|
|
beforeBoot: async () => {
|
|
order.push("beforeBoot");
|
|
},
|
|
}),
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "before-boot",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionOutPartial(partial as never);
|
|
// Two users — smart default consumes u1 into the chain, leaves u2 for dispatch
|
|
harness.seedSessionInTail([u1 as never, u2 as never]);
|
|
try {
|
|
await new Promise((r) => setTimeout(r, 50));
|
|
expect(order).toEqual(["beforeBoot", "turn"]);
|
|
} finally {
|
|
await harness.close();
|
|
}
|
|
});
|
|
|
|
it("hook throwing falls back to defaults without sinking the run", async () => {
|
|
let turnCount = 0;
|
|
const model = new MockLanguageModelV3({
|
|
doStream: async () => {
|
|
turnCount++;
|
|
return { stream: textStream("ok") };
|
|
},
|
|
});
|
|
const partial = assistantMessage("partial", "a-partial");
|
|
const u1 = userMessage("buffered original", "u-1");
|
|
const u2 = userMessage("followup", "u-2");
|
|
const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
|
|
const agent = chat.agent({
|
|
id: "recovery-boot.hook-throws",
|
|
onRecoveryBoot: async () => {
|
|
throw new Error("kaboom");
|
|
},
|
|
run: async ({ messages, signal }) => streamText({ model, messages, abortSignal: signal }),
|
|
});
|
|
const harness = mockChatAgent(agent, {
|
|
chatId: "hook-throws",
|
|
continuation: true,
|
|
previousRunId: "run_prior",
|
|
});
|
|
harness.seedSessionOutPartial(partial as never);
|
|
// Two users so smart default leaves u2 to dispatch (u1 spliced into chain)
|
|
harness.seedSessionInTail([u1 as never, u2 as never]);
|
|
try {
|
|
await new Promise((r) => setTimeout(r, 100));
|
|
// Default behavior: the in-flight user is re-dispatched as a turn
|
|
// even though the hook threw.
|
|
expect(turnCount).toBe(1);
|
|
} finally {
|
|
await harness.close();
|
|
warnSpy.mockRestore();
|
|
}
|
|
});
|
|
});
|