feat(webapp,sdk): dashboard agent plan enforcement, component gallery — and fixes (#4516)
Plan enforcement for the dashboard agent — message quota and watch limits — plus the component gallery, fixes and test hardening from the same stack (#4548, #4549, #4550, #4552, #4556 merged here). ## Plan enforcement ([TRI-12863](https://linear.app/triggerdotdev/issue/TRI-12863)) **Agent message quota.** The Free-plan allowance becomes a real server-side limit with a durable counter. New `agent_message_usage` table keyed `(organization_id, period)` — deliberately not joined to chats, so deleting a chat can't free quota within the period. Both send paths count one user message (wakes never count) and refuse at the cap with `403 message_quota_reached`, which the client renders as an upgrade panel, never a silent drop. The refusal code is a single shared constant on both sides. **Watch limits.** A watch whose window exceeds the plan's `agentWatchMaxHours`, or that would push the org past its `agentWatchers` count, is refused with `watch_limit_reached` (409 on the API, an upgrade hint on the card). Plan limits only tighten the existing code ceilings (`min(plan, 24h)`, per-chat cap of 3 still applies). A plan limit of zero means zero, not unlimited. Questions answerable instantly are answered before any plan refusal — a one-shot consumes no slot and never sees an upgrade nag. **Fails open by design.** Cloud ships the actual per-plan numbers separately (TRI-12863 P0). Until then absent limits resolve to the unlimited sentinel and the upgrade UI is gated on billing presence — self-hosted sees no cap, no upsell, with tests proving the fallback. Both quotas are nudges, not security boundaries: a failing limit read never blocks a send. ## Component gallery An admin-only gallery of every agent card state: five `storybook.agent-*` pages (chat UI, view blocks, report, investigation, watch) with their shared shell and manifest, demo fixtures, two demo-only cards, toast examples, and the screenshot script. No LLM and no data — every state renders from fixtures under `dashboard-agent/demo/`, never reachable from a production path. Designers and reviewers can look at every state, including the report states, without seeding anything. ## And fixes **SDK: watch-mode chat subscriptions survive quiet windows** (TRI-13065, TRI-13070) — watch mode keeps reconnecting across empty long-poll windows and only stops on abort or a settled session; a passive subscriber can no longer stop a turn it doesn't own (`stopOnAbort` is explicit, default off). Review findings fixed alongside: a superseded stream's async teardown no longer removes the live successor's abort controller or multi-tab claim, and stopping a generation hands the chat back to the user's other tabs. **Query boundary pinned end-to-end** ([TRI-11165](https://linear.app/triggerdotdev/issue/TRI-11165)) — a route-level test drives `api.v1.query` with a real signed environment JWT (writes refused before ClickHouse, a read passes); `readonly=1` made non-overridable; a per-turn cap stops the model burning a turn rewriting a query it can't fix (deterministic SQL errors only — busy/transport rejections don't count). **chat.agent durability regression suite** ([TRI-11166](https://linear.app/triggerdotdev/issue/TRI-11166)) — testcontainers-backed coverage of the two audit criticals (cross-tenant isolation, no duplicate mid-stream turn, both control-broken) plus crash-resume, cursor-based refresh, clean rollback of a mid-write turn failure (torn by a real constraint violation), and OOM-restart replay. **Investigation sweep backoff** — stale investigations get an attempt counter and backoff so a poison row can't pin the sweep queue head (migration `0005`: `sweep_attempts`, `last_sweep_attempt_at`). ## Screenshots <img width="1440" height="791" alt="Screenshot 2026-08-06 at 00 36 19" src="https://github.com/user-attachments/assets/6a68cd42-8580-469d-afe7-e28d1eef18e1" /> 🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
Katia Bulatova
@@ -1,6 +1,6 @@
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
|
||||
import type { UIMessage, UIMessageChunk } from "ai";
|
||||
import { TriggerChatTransport, createChatTransport } from "./chat.js";
|
||||
import { TriggerChatTransport, createChatTransport, type ChatTransportEvent } from "./chat.js";
|
||||
|
||||
// ───────────────────────────────────────────────────────────────────────────
|
||||
// Test helpers
|
||||
@@ -132,6 +132,35 @@ function defaultSseResponse(
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* An SSE response whose body stays open until the request signal aborts.
|
||||
* Models a live subscription sitting on a quiet server.
|
||||
*/
|
||||
function openSseResponse(signal?: AbortSignal | null): Response {
|
||||
const body = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
const onAbort = () => {
|
||||
const err = new Error("aborted");
|
||||
err.name = "AbortError";
|
||||
try {
|
||||
controller.error(err);
|
||||
} catch {
|
||||
/* already errored */
|
||||
}
|
||||
};
|
||||
if (signal?.aborted) onAbort();
|
||||
else signal?.addEventListener("abort", onAbort, { once: true });
|
||||
},
|
||||
});
|
||||
return new Response(body, {
|
||||
status: 200,
|
||||
headers: {
|
||||
"content-type": "text/event-stream",
|
||||
"X-Stream-Version": "v2",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function authError(status = 401): Response {
|
||||
return new Response(JSON.stringify({ error: "Unauthorized", name: "TriggerApiError", status }), {
|
||||
status,
|
||||
@@ -1027,6 +1056,37 @@ describe("TriggerChatTransport", () => {
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
it("resumes in watch mode when the session is hydrated with isStreaming=false", async () => {
|
||||
let subscribeCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) {
|
||||
subscribeCount++;
|
||||
const response = defaultSseResponse([{ type: "text-delta", id: "p1", delta: "turn2" }]);
|
||||
const headers = new Headers(response.headers);
|
||||
headers.set("X-Session-Settled", "true");
|
||||
return new Response(response.body, { status: 200, headers });
|
||||
}
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
sessions: {
|
||||
"chat-rc-watch": { publicAccessToken: "p", isStreaming: false },
|
||||
},
|
||||
});
|
||||
|
||||
const stream = await transport.reconnectToStream({ chatId: "chat-rc-watch" });
|
||||
expect(stream).not.toBeNull();
|
||||
const chunks = await drainChunks(stream!);
|
||||
|
||||
expect(subscribeCount).toBe(1);
|
||||
expect(chunks).toEqual([{ type: "text-delta", id: "p1", delta: "turn2" }]);
|
||||
});
|
||||
|
||||
it("opens an SSE subscription with the X-Peek-Settled header set", async () => {
|
||||
let subscribeHeaders: Headers | undefined;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
|
||||
@@ -1176,7 +1236,7 @@ describe("TriggerChatTransport", () => {
|
||||
expect(transport.getSession("chat-slow")?.isStreaming).toBe(false);
|
||||
});
|
||||
|
||||
it("gives up after a bounded number of resubscribes", async () => {
|
||||
it("surfaces an error after the resubscribe budget is exhausted", async () => {
|
||||
// Fake timers so the 100ms..1.6s backoffs don't cost real seconds.
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
@@ -1205,12 +1265,18 @@ describe("TriggerChatTransport", () => {
|
||||
messages: [createUserMessage("hi")],
|
||||
abortSignal: undefined,
|
||||
});
|
||||
// A cut-off turn surfaces an error rather than reading as complete.
|
||||
// Attach the rejection assertion before advancing timers so the
|
||||
// rejection is never unhandled.
|
||||
const drained = drainChunks(stream);
|
||||
const rejects = expect(drained).rejects.toThrow(/reconnect budget exhausted/i);
|
||||
await vi.advanceTimersByTimeAsync(10_000);
|
||||
await drained;
|
||||
await rejects;
|
||||
|
||||
// One initial connect plus the five-attempt resubscribe budget.
|
||||
expect(subscribeCount).toBe(6);
|
||||
// State is cleared before the throw, so a reload won't reopen a
|
||||
// doomed subscription.
|
||||
expect(transport.getSession("chat-empty")?.isStreaming).toBe(false);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
@@ -1218,6 +1284,455 @@ describe("TriggerChatTransport", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("watch mode across long-poll window boundaries", () => {
|
||||
function settled(response: Response): Response {
|
||||
const headers = new Headers(response.headers);
|
||||
headers.set("X-Session-Settled", "true");
|
||||
return new Response(response.body, { status: 200, headers });
|
||||
}
|
||||
|
||||
it("resubscribes after a completed turn and receives a later wake", async () => {
|
||||
let subscribeCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) {
|
||||
subscribeCount++;
|
||||
// Window 1: a turn completes, then the body EOFs with no
|
||||
// settled header — the quiet long-poll boundary.
|
||||
return subscribeCount === 1
|
||||
? defaultSseResponse([
|
||||
{ type: "text-delta", id: "p1", delta: "turn1" },
|
||||
{ type: "trigger:turn-complete" },
|
||||
])
|
||||
: settled(defaultSseResponse([{ type: "text-delta", id: "p2", delta: "wake" }]));
|
||||
}
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
sessions: { "chat-watch-eof": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const stream = await transport.reconnectToStream({ chatId: "chat-watch-eof" });
|
||||
const chunks = await drainChunks(stream!);
|
||||
|
||||
expect(subscribeCount).toBe(2);
|
||||
expect(chunks).toEqual([
|
||||
{ type: "text-delta", id: "p1", delta: "turn1" },
|
||||
{ type: "text-delta", id: "p2", delta: "wake" },
|
||||
]);
|
||||
});
|
||||
|
||||
it("does not peek-settle an idle resubscribe, so the next turn is delivered", async () => {
|
||||
// Watch mode must NOT send X-Peek-Settled between turns: a settled peek
|
||||
// while no turn is in flight closes the standing subscription and the
|
||||
// viewer never sees turn 2. This mock plays the server's peek shortcut —
|
||||
// a peek request with nothing in flight settles — to prove the transport
|
||||
// long-polls instead.
|
||||
const subscribeHeaders: Headers[] = [];
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) {
|
||||
subscribeHeaders.push(new Headers(init?.headers));
|
||||
const n = subscribeHeaders.length;
|
||||
if (n === 1) {
|
||||
// Turn 1 completes, then the body EOFs (no settled header).
|
||||
return defaultSseResponse([
|
||||
{ type: "text-delta", id: "p1", delta: "turn1" },
|
||||
{ type: "trigger:turn-complete" },
|
||||
]);
|
||||
}
|
||||
if (n === 2) {
|
||||
// Idle resubscribe. If it peeked, the server settles and the
|
||||
// subscription would close before turn 2; a long-poll delivers it.
|
||||
if (init && new Headers(init.headers).get("X-Peek-Settled")) {
|
||||
return settled(defaultSseResponse([]));
|
||||
}
|
||||
return defaultSseResponse([
|
||||
{ type: "text-delta", id: "p2", delta: "turn2" },
|
||||
{ type: "trigger:turn-complete" },
|
||||
]);
|
||||
}
|
||||
// Turn 2 done — end the watch cleanly.
|
||||
return settled(defaultSseResponse([]));
|
||||
}
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
sessions: { "chat-watch-turn2": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const stream = await transport.reconnectToStream({ chatId: "chat-watch-turn2" });
|
||||
const chunks = await drainChunks(stream!);
|
||||
|
||||
expect(subscribeHeaders[1]?.get("X-Peek-Settled")).toBeNull();
|
||||
expect(chunks).toEqual([
|
||||
{ type: "text-delta", id: "p1", delta: "turn1" },
|
||||
{ type: "text-delta", id: "p2", delta: "turn2" },
|
||||
]);
|
||||
});
|
||||
|
||||
it("cancelling the reader stops the resubscribe loop", async () => {
|
||||
// A consumer that stops reading without aborting must not leak the
|
||||
// resubscribe loop — the stream's cancel() aborts it.
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
let subscribeCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) {
|
||||
subscribeCount++;
|
||||
// Quiet: EOF, no records, never settled — watch keeps resubscribing.
|
||||
return defaultSseResponse([]);
|
||||
}
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const events: ChatTransportEvent[] = [];
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
onEvent: (e) => events.push(e),
|
||||
sessions: { "chat-watch-cancel": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const stream = await transport.reconnectToStream({ chatId: "chat-watch-cancel" });
|
||||
const reader = stream!.getReader();
|
||||
await vi.advanceTimersByTimeAsync(10_000);
|
||||
expect(subscribeCount).toBeGreaterThan(1);
|
||||
|
||||
const countAtCancel = subscribeCount;
|
||||
await reader.cancel();
|
||||
await vi.advanceTimersByTimeAsync(10_000);
|
||||
expect(subscribeCount).toBe(countAtCancel);
|
||||
// A clean cancel must not surface a spurious stream-error (an
|
||||
// unguarded controller.close() after cancel would throw "Invalid
|
||||
// state" and leak it onto the telemetry channel).
|
||||
expect(events.some((e) => e.type === "stream-error")).toBe(false);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it("stops when the server says the session settled", async () => {
|
||||
let subscribeCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) {
|
||||
subscribeCount++;
|
||||
return settled(
|
||||
defaultSseResponse([
|
||||
{ type: "text-delta", id: "p1", delta: "last" },
|
||||
{ type: "trigger:turn-complete" },
|
||||
])
|
||||
);
|
||||
}
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
sessions: { "chat-watch-settled": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const stream = await transport.reconnectToStream({ chatId: "chat-watch-settled" });
|
||||
const chunks = await drainChunks(stream!);
|
||||
|
||||
expect(subscribeCount).toBe(1);
|
||||
expect(chunks).toHaveLength(1);
|
||||
expect(transport.getSession("chat-watch-settled")?.isStreaming).toBe(false);
|
||||
});
|
||||
|
||||
it("stops promptly when aborted during backoff", async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
let subscribeCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) {
|
||||
subscribeCount++;
|
||||
// Every window is quiet: EOF with no records, never settled.
|
||||
return defaultSseResponse([]);
|
||||
}
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const abortController = new AbortController();
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
sessions: { "chat-watch-abort": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const stream = await transport.reconnectToStream({
|
||||
chatId: "chat-watch-abort",
|
||||
abortSignal: abortController.signal,
|
||||
});
|
||||
const drained = drainChunks(stream!);
|
||||
await vi.advanceTimersByTimeAsync(10_000);
|
||||
// The budget doesn't apply in watch mode, so it is still reconnecting.
|
||||
expect(subscribeCount).toBeGreaterThan(6);
|
||||
|
||||
const countAtAbort = subscribeCount;
|
||||
abortController.abort();
|
||||
await drained;
|
||||
await vi.advanceTimersByTimeAsync(10_000);
|
||||
expect(subscribeCount).toBe(countAtAbort);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("reconnectToStream stop-on-abort ownership (TRI-13070)", () => {
|
||||
// A quiet stream: EOF, no records, never settled — the subscription
|
||||
// stays alive (watch mode) so an abort mid-flight exercises the stop path.
|
||||
function quietWatchTransport(): {
|
||||
transport: TriggerChatTransport;
|
||||
appends: () => number;
|
||||
} {
|
||||
let appendCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionStreamAppendUrl(urlStr)) {
|
||||
appendCount++;
|
||||
return defaultAppendResponse();
|
||||
}
|
||||
if (isSessionOutSubscribeUrl(urlStr)) return defaultSseResponse([]);
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
sessions: { "chat-own": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
return { transport, appends: () => appendCount };
|
||||
}
|
||||
|
||||
it("passive subscriber aborting writes no stop chunk to .in", async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
const { transport, appends } = quietWatchTransport();
|
||||
const abort = new AbortController();
|
||||
const stream = await transport.reconnectToStream({
|
||||
chatId: "chat-own",
|
||||
abortSignal: abort.signal,
|
||||
});
|
||||
const drained = drainChunks(stream!);
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
abort.abort();
|
||||
await drained;
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
expect(appends()).toBe(0);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it("owning subscriber with stopOnAbort:true sends a stop chunk on abort", async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
const { transport, appends } = quietWatchTransport();
|
||||
const abort = new AbortController();
|
||||
const stream = await transport.reconnectToStream({
|
||||
chatId: "chat-own",
|
||||
abortSignal: abort.signal,
|
||||
stopOnAbort: true,
|
||||
});
|
||||
const drained = drainChunks(stream!);
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
abort.abort();
|
||||
await drained;
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
expect(appends()).toBe(1);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it("abortSignal presence alone (stopOnAbort unset) sends no stop", async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
const { transport, appends } = quietWatchTransport();
|
||||
const abort = new AbortController();
|
||||
const stream = await transport.reconnectToStream({
|
||||
chatId: "chat-own",
|
||||
abortSignal: abort.signal,
|
||||
});
|
||||
const drained = drainChunks(stream!);
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
abort.abort();
|
||||
await drained;
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
expect(appends()).toBe(0);
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("superseded stream teardown", () => {
|
||||
it("keeps the successor's controller registered when the aborted stream tears down", async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
let appendCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionStreamAppendUrl(urlStr)) {
|
||||
appendCount++;
|
||||
return defaultAppendResponse();
|
||||
}
|
||||
// Quiet stream: EOF, no records, never settled — watch keeps it open.
|
||||
if (isSessionOutSubscribeUrl(urlStr)) return defaultSseResponse([]);
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
watch: true,
|
||||
sessions: { "chat-race": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const send = () =>
|
||||
transport.sendMessages({
|
||||
trigger: "submit-message" as const,
|
||||
chatId: "chat-race",
|
||||
messageId: undefined,
|
||||
messages: [createUserMessage("hi")],
|
||||
abortSignal: undefined,
|
||||
});
|
||||
|
||||
const first = drainChunks(await send());
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
|
||||
// Supersede: the new stream registers its controller synchronously,
|
||||
// the aborted one tears down a microtask later.
|
||||
const second = await send();
|
||||
let secondClosed = false;
|
||||
const secondDrain = drainChunks(second).then(() => {
|
||||
secondClosed = true;
|
||||
});
|
||||
await first;
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
|
||||
// stopGeneration posts the stop chunk either way — only the
|
||||
// closing assertion proves it found the successor to abort.
|
||||
appendCount = 0;
|
||||
expect(await transport.stopGeneration("chat-race")).toBe(true);
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
expect(appendCount).toBe(1);
|
||||
expect(secondClosed).toBe(true);
|
||||
|
||||
transport.dispose();
|
||||
await secondDrain;
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it("keeps the tab claim the successor took (multi-tab)", async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse();
|
||||
// Open SSE that only ends when the subscription is aborted, so
|
||||
// the superseded stream tears down while the successor is live.
|
||||
if (isSessionOutSubscribeUrl(urlStr)) return openSseResponse(init?.signal);
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
multiTab: true,
|
||||
sessions: { "chat-race-tab": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const send = () =>
|
||||
transport.sendMessages({
|
||||
trigger: "submit-message" as const,
|
||||
chatId: "chat-race-tab",
|
||||
messageId: undefined,
|
||||
messages: [createUserMessage("hi")],
|
||||
abortSignal: undefined,
|
||||
});
|
||||
|
||||
const first = drainChunks(await send());
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
const secondDrain = drainChunks(await send());
|
||||
await first;
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
|
||||
// The superseded stream must not release the claim its successor
|
||||
// holds — otherwise this tab flips to read-only mid-turn.
|
||||
expect(transport.hasClaim("chat-race-tab")).toBe(true);
|
||||
|
||||
transport.dispose();
|
||||
await secondDrain;
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
|
||||
it("releases the tab claim when the user stops generation (multi-tab)", async () => {
|
||||
vi.useFakeTimers();
|
||||
try {
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) return openSseResponse(init?.signal);
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
const transport = new TriggerChatTransport({
|
||||
task: "my-chat-task",
|
||||
accessToken: () => "pat",
|
||||
multiTab: true,
|
||||
sessions: { "chat-stop-tab": { publicAccessToken: "p", isStreaming: true } },
|
||||
});
|
||||
|
||||
const drain = drainChunks(
|
||||
await transport.sendMessages({
|
||||
trigger: "submit-message" as const,
|
||||
chatId: "chat-stop-tab",
|
||||
messageId: undefined,
|
||||
messages: [createUserMessage("hi")],
|
||||
abortSignal: undefined,
|
||||
})
|
||||
);
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
expect(transport.hasClaim("chat-stop-tab")).toBe(true);
|
||||
|
||||
expect(await transport.stopGeneration("chat-stop-tab")).toBe(true);
|
||||
await vi.advanceTimersByTimeAsync(1_000);
|
||||
|
||||
// The turn ends here with no successor stream, so the claim must be
|
||||
// freed or other tabs stay read-only until this one closes.
|
||||
expect(transport.hasClaim("chat-stop-tab")).toBe(false);
|
||||
|
||||
transport.dispose();
|
||||
await drain;
|
||||
} finally {
|
||||
vi.useRealTimers();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("multi-tab coordination", () => {
|
||||
it("isReadOnly defaults to false when multiTab is disabled", () => {
|
||||
const transport = new TriggerChatTransport({
|
||||
@@ -1488,9 +2003,18 @@ describe("TriggerChatTransport", () => {
|
||||
{ type: "text-delta", id: "p2", delta: "Again" },
|
||||
{ type: "trigger:turn-complete" },
|
||||
];
|
||||
let subscribeCount = 0;
|
||||
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
|
||||
const urlStr = typeof url === "string" ? url : url.toString();
|
||||
if (isSessionOutSubscribeUrl(urlStr)) return defaultSseResponse(turn1);
|
||||
if (isSessionOutSubscribeUrl(urlStr)) {
|
||||
subscribeCount++;
|
||||
if (subscribeCount === 1) return defaultSseResponse(turn1);
|
||||
// Watch mode reconnects past the body EOF; settle so the drain ends.
|
||||
const response = defaultSseResponse([]);
|
||||
const headers = new Headers(response.headers);
|
||||
headers.set("X-Session-Settled", "true");
|
||||
return new Response(response.body, { status: 200, headers });
|
||||
}
|
||||
throw new Error(`Unexpected URL: ${urlStr}`);
|
||||
});
|
||||
|
||||
|
||||
@@ -873,7 +873,11 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
state.isStreaming = true;
|
||||
this.notifySessionChange(chatId, state);
|
||||
|
||||
return this.subscribeToSessionStream(state, abortSignal, chatId, { sinceInSeq: inSeq });
|
||||
// Owning turn: aborting this live send stops the turn the user drives.
|
||||
return this.subscribeToSessionStream(state, abortSignal, chatId, {
|
||||
sinceInSeq: inSeq,
|
||||
sendStopOnAbort: true,
|
||||
});
|
||||
};
|
||||
|
||||
/**
|
||||
@@ -1146,12 +1150,21 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
options: {
|
||||
chatId: string;
|
||||
abortSignal?: AbortSignal | undefined;
|
||||
/**
|
||||
* Whether aborting this subscription sends `{kind:"stop"}` on `.in`.
|
||||
* A subscription ending is not session ownership — a passive/watch
|
||||
* reader unmounting must never stop a turn it doesn't drive. Only
|
||||
* pass `true` from a caller that owns the live turn. @default false
|
||||
*/
|
||||
stopOnAbort?: boolean;
|
||||
} & ChatRequestOptions
|
||||
): Promise<ReadableStream<UIMessageChunk> | null> => {
|
||||
const state = this.sessions.get(options.chatId);
|
||||
if (!state) return null;
|
||||
|
||||
if (state.isStreaming === false) return null;
|
||||
// Watch is a standing subscription: a settled session is exactly the
|
||||
// state it waits in, so a completed last turn must not block the resume.
|
||||
if (state.isStreaming === false && !this.watchMode) return null;
|
||||
if (this.activeStreams.has(options.chatId)) return null;
|
||||
|
||||
const abortController = new AbortController();
|
||||
@@ -1163,12 +1176,14 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
|
||||
return this.subscribeToSessionStream(state, abortSignal, options.chatId, {
|
||||
resumed: true,
|
||||
sendStopOnAbort: !!options.abortSignal,
|
||||
sendStopOnAbort: options.stopOnAbort ?? false,
|
||||
// Reconnect-on-reload opts into the server's settled-peek shortcut
|
||||
// so the SSE doesn't hang for 60s when no turn is in flight. Active
|
||||
// send-a-message paths must keep wait=60 to avoid racing the
|
||||
// freshly-triggered turn's first chunk.
|
||||
peekSettled: true,
|
||||
// freshly-triggered turn's first chunk. Watch mode must NOT peek: a
|
||||
// settled peek between turns sets sessionSettled and closes the
|
||||
// standing subscription, so the viewer never sees the next turn.
|
||||
peekSettled: !this.watchMode,
|
||||
});
|
||||
};
|
||||
|
||||
@@ -1205,6 +1220,11 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
activeStream.abort();
|
||||
this.activeStreams.delete(chatId);
|
||||
}
|
||||
// Release here, not in the stream teardown: that only releases while it
|
||||
// still owns the map entry, and we just deleted it. Unlike a supersede,
|
||||
// no successor stream follows a stop, so the claim would never be freed
|
||||
// and other tabs would stay read-only until this one closes.
|
||||
this.coordinator?.release(chatId);
|
||||
|
||||
// The turn won't reach its turn-complete on this client (we just
|
||||
// aborted the reader), so clear the streaming flag here and persist —
|
||||
@@ -1266,7 +1286,11 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
state.isStreaming = true;
|
||||
this.notifySessionChange(chatId, state);
|
||||
|
||||
return this.subscribeToSessionStream(state, undefined, chatId, { sinceInSeq: inSeq });
|
||||
// Owning action: aborting this send stops the turn the user drives.
|
||||
return this.subscribeToSessionStream(state, undefined, chatId, {
|
||||
sinceInSeq: inSeq,
|
||||
sendStopOnAbort: true,
|
||||
});
|
||||
};
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
@@ -1780,11 +1804,13 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
let eofResubscribes = 0;
|
||||
|
||||
const resumeAfterEof = async () => {
|
||||
// Watch mode is a standing subscription: it outlives turn-complete
|
||||
// (which clears `isStreaming`) and idle windows EOF by design, so the
|
||||
// give-up budget doesn't apply. Only abort or a settled session ends it.
|
||||
while (
|
||||
state.isStreaming &&
|
||||
(this.watchMode || (state.isStreaming && eofResubscribes < MAX_EOF_RESUBSCRIBES)) &&
|
||||
!currentSubscription?.sessionSettled &&
|
||||
!combinedSignal.aborted &&
|
||||
eofResubscribes < MAX_EOF_RESUBSCRIBES
|
||||
!combinedSignal.aborted
|
||||
) {
|
||||
eofResubscribes++;
|
||||
// Sleep, but wake immediately on abort — otherwise a stop lands
|
||||
@@ -1796,7 +1822,10 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
combinedSignal.removeEventListener("abort", done);
|
||||
resolve();
|
||||
};
|
||||
timer = setTimeout(done, Math.min(100 * 2 ** (eofResubscribes - 1), 5_000));
|
||||
// Jitter the backoff so many clients reconnecting after the same
|
||||
// dropped window don't resubscribe in lockstep.
|
||||
const backoff = Math.min(100 * 2 ** (eofResubscribes - 1), 5_000);
|
||||
timer = setTimeout(done, backoff * (0.5 + Math.random() * 0.5));
|
||||
combinedSignal.addEventListener("abort", done);
|
||||
});
|
||||
if (combinedSignal.aborted) break;
|
||||
@@ -1804,6 +1833,25 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
if (opened) return opened;
|
||||
}
|
||||
|
||||
// A settled session or an abort ends the turn cleanly. Exhausting the
|
||||
// resubscribe budget while the turn is still streaming means it was cut
|
||||
// off — surface an error so the UI doesn't read a truncated reply as
|
||||
// complete. The caller's catch emits stream-error and errors the stream.
|
||||
if (
|
||||
state.isStreaming &&
|
||||
!currentSubscription?.sessionSettled &&
|
||||
!combinedSignal.aborted
|
||||
) {
|
||||
// Clear + persist before throwing so the surfaced error leaves
|
||||
// consistent state — otherwise a reload sees isStreaming: true
|
||||
// and reopens a doomed subscription.
|
||||
state.isStreaming = false;
|
||||
this.notifySessionChange(chatId, state);
|
||||
throw new Error(
|
||||
"Chat stream ended before the turn completed (reconnect budget exhausted)."
|
||||
);
|
||||
}
|
||||
|
||||
// Settled close, or the turn is gone — tell the UI instead of
|
||||
// leaving it spinning on a stream nobody will finish.
|
||||
if (state.isStreaming) {
|
||||
@@ -1823,7 +1871,11 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
|
||||
const opened = (await openWithAuthRetry()) ?? (await resumeAfterEof());
|
||||
if (opened === null) {
|
||||
controller.close();
|
||||
try {
|
||||
controller.close();
|
||||
} catch {
|
||||
/* already closed by a consumer cancel */
|
||||
}
|
||||
return;
|
||||
}
|
||||
reader = opened.reader;
|
||||
@@ -1853,7 +1905,11 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
if (next.done) {
|
||||
const resumed = await resumeAfterEof();
|
||||
if (resumed === null) {
|
||||
controller.close();
|
||||
try {
|
||||
controller.close();
|
||||
} catch {
|
||||
/* already closed by a consumer cancel */
|
||||
}
|
||||
return;
|
||||
}
|
||||
reader = resumed.reader;
|
||||
@@ -1866,7 +1922,11 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
if (combinedSignal.aborted) {
|
||||
internalAbort.abort();
|
||||
await reader.cancel();
|
||||
controller.close();
|
||||
try {
|
||||
controller.close();
|
||||
} catch {
|
||||
/* already closed by a consumer cancel */
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -2001,10 +2061,20 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
|
||||
controller.error(error);
|
||||
} finally {
|
||||
teardownWakeListeners();
|
||||
this.activeStreams.delete(chatId);
|
||||
this.coordinator?.release(chatId);
|
||||
// Only clear the entry (and drop the tab claim) if it is still
|
||||
// ours — a superseding send registers its controller before this
|
||||
// teardown runs, and owns the claim from then on.
|
||||
if (this.activeStreams.get(chatId) === internalAbort) {
|
||||
this.activeStreams.delete(chatId);
|
||||
this.coordinator?.release(chatId);
|
||||
}
|
||||
}
|
||||
},
|
||||
// A consumer that stops reading without aborting (drops the reader)
|
||||
// would otherwise leave the resubscribe loop running forever.
|
||||
cancel() {
|
||||
internalAbort.abort();
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -211,11 +211,20 @@ describe("transport stream events", () => {
|
||||
``,
|
||||
].join("\n");
|
||||
|
||||
let subscribes = 0;
|
||||
const { transport, events } = makeTransport({
|
||||
watch: true,
|
||||
sessions: { c1: { publicAccessToken: "tok_test", isStreaming: true } },
|
||||
fetch: async (_url, _init, ctx) =>
|
||||
ctx.endpoint === "in" ? jsonOk() : sseResponse(TWO_TURNS),
|
||||
fetch: async (_url, _init, ctx) => {
|
||||
if (ctx.endpoint === "in") return jsonOk();
|
||||
if (subscribes++ > 0) {
|
||||
// Watch mode reconnects past the body EOF; settle so the read ends.
|
||||
const settled = sseResponse("");
|
||||
settled.headers.set("X-Session-Settled", "true");
|
||||
return settled;
|
||||
}
|
||||
return sseResponse(TWO_TURNS);
|
||||
},
|
||||
});
|
||||
|
||||
const stream = await transport.reconnectToStream({ chatId: "c1" });
|
||||
|
||||
Reference in New Issue
Block a user