wehub-resource-sync
966 lines
39 KiB
TypeScript
966 lines
39 KiB
TypeScript
import type { UIMessage } from "@ai-sdk/react";
|
|
import { ChatSnapshotV1Schema, SSEStreamSubscription } from "@trigger.dev/core/v3";
|
|
import { useEffect, useMemo, useRef, useState } from "react";
|
|
import { Paragraph } from "~/components/primitives/Paragraph";
|
|
import { Spinner } from "~/components/primitives/Spinner";
|
|
import { AgentMessageView } from "~/components/runs/v3/agent/AgentMessageView";
|
|
import { useAutoScrollToBottom } from "~/hooks/useAutoScrollToBottom";
|
|
import { useEnvironment } from "~/hooks/useEnvironment";
|
|
import { useOrganization } from "~/hooks/useOrganizations";
|
|
import { useProject } from "~/hooks/useProject";
|
|
|
|
export type AgentViewAuth = {
|
|
publicAccessToken: string;
|
|
apiOrigin: string;
|
|
/**
|
|
* Session identifier the AgentView uses to address the backing
|
|
* {@link Session} when subscribing to `.in` / `.out`. Accepts either
|
|
* a `session_*` friendlyId or the transport-supplied externalId
|
|
* (typically the browser's `chatId`) — the dashboard resource route
|
|
* resolves either form via `resolveSessionByIdOrExternalId`.
|
|
*/
|
|
sessionId: string;
|
|
/**
|
|
* User messages extracted from the run's task payload at load time.
|
|
* Empty array for runs started with `trigger: "preload"` — in that
|
|
* case the first user message arrives over the session's `.in`
|
|
* channel and is merged in by the AgentView subscription.
|
|
*/
|
|
initialMessages: UIMessage[];
|
|
/**
|
|
* Presigned GET URL for the session's chat-snapshot S3 blob (written
|
|
* by the agent after each turn-complete; see `ChatSnapshotV1`).
|
|
* Optional — sessions that registered a `hydrateMessages` hook skip
|
|
* snapshot writes and the URL fetch will 404. In that case the
|
|
* dashboard falls back to seq=0 SSE (which, post-trim, shows only the
|
|
* most recent turn). Generated server-side by `SessionPresenter`.
|
|
*/
|
|
snapshotPresignedUrl?: string;
|
|
};
|
|
|
|
/**
|
|
* Max state-update interval while assistant chunks are streaming. Matches
|
|
* the `experimental_throttle: 100` we previously passed to `useChat`.
|
|
* Chunks mutate a staging ref synchronously; a throttled flush copies the
|
|
* ref into React state at most ~10x/sec so tool-call Prism highlighting
|
|
* etc. doesn't re-run on every single text-delta.
|
|
*/
|
|
const STATE_FLUSH_THROTTLE_MS = 100;
|
|
|
|
/**
|
|
* Sentinel timestamp for messages that came from the run's initial task
|
|
* payload — they predate any stream activity, so 0 guarantees they sort
|
|
* first regardless of stream race order.
|
|
*/
|
|
const INITIAL_PAYLOAD_TIMESTAMP = 0;
|
|
|
|
/**
|
|
* Renders a Session's chat conversation as it unfolds.
|
|
*
|
|
* Subscribes to both channels of the {@link Session}:
|
|
* - **`.out`** delivers assistant `UIMessageChunk`s (text deltas, tool
|
|
* calls, reasoning, etc.) produced by the agent's
|
|
* `chatStream.writer(...)` calls — objects, already parsed by the S2
|
|
* SSE reader.
|
|
* - **`.in`** delivers {@link ChatInputChunk}s sent by
|
|
* {@link TriggerChatTransport} (or any other session writer). Each
|
|
* chunk is a tagged union (`{kind: "message", payload}` for user
|
|
* turns, `{kind: "stop"}` for stop signals) — the AgentView only
|
|
* cares about `kind: "message"` and pulls `.payload.messages`.
|
|
*
|
|
* Both streams are read directly via `SSEStreamSubscription` through the
|
|
* dashboard's session-authed resource routes — not through `useChat` or
|
|
* `TriggerChatTransport`. This gives us per-chunk server-side timestamps
|
|
* (S2 sequence numbers) from both streams, which we use to produce a
|
|
* chronologically correct merged message list that works for replays,
|
|
* multi-message turns, cross-run session resumes, and steering messages.
|
|
*
|
|
* Intended to be mounted inside a scrollable container — the component
|
|
* does not own its own scrollbar.
|
|
*/
|
|
export function AgentView({ agentView }: { agentView: AgentViewAuth }) {
|
|
const organization = useOrganization();
|
|
const project = useProject();
|
|
const environment = useEnvironment();
|
|
|
|
const messages = useAgentSessionMessages({
|
|
sessionId: agentView.sessionId,
|
|
apiOrigin: agentView.apiOrigin,
|
|
orgSlug: organization.slug,
|
|
projectSlug: project.slug,
|
|
envSlug: environment.slug,
|
|
initialMessages: agentView.initialMessages,
|
|
snapshotPresignedUrl: agentView.snapshotPresignedUrl,
|
|
});
|
|
|
|
// Sticky-bottom auto-scroll: walks up to find the inspector's scroll
|
|
// container, then scrolls to bottom whenever `messages` changes — but
|
|
// only if the user was at (or near) the bottom at the time. Scrolling
|
|
// away pauses auto-scroll; scrolling back resumes it.
|
|
const rootRef = useAutoScrollToBottom([messages]);
|
|
|
|
return (
|
|
<div ref={rootRef} className="flex min-h-full flex-col py-3">
|
|
{messages.length === 0 ? (
|
|
<div className="flex flex-1 items-center justify-center">
|
|
<div className="flex flex-col items-center gap-3">
|
|
<Spinner className="size-5" color="blue" />
|
|
<Paragraph variant="small" className="text-text-dimmed">
|
|
Loading conversation…
|
|
</Paragraph>
|
|
</div>
|
|
</div>
|
|
) : (
|
|
<AgentMessageView messages={messages} />
|
|
)}
|
|
</div>
|
|
);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// useAgentSessionMessages — reads both realtime streams for a session and
|
|
// maintains a chronologically ordered, merged message list.
|
|
// ---------------------------------------------------------------------------
|
|
|
|
/**
|
|
* Shape of each chunk on the session's `.in` channel. Mirrors the
|
|
* `ChatInputChunk` tagged union produced by {@link TriggerChatTransport}:
|
|
* - `kind: "message"` carries a `ChatTaskWirePayload` in `.payload`
|
|
* (user-submitted messages or regenerate calls); we dedupe by id.
|
|
* - `kind: "stop"` is a stop signal — no messages, nothing to render
|
|
* here, so it's filtered.
|
|
*
|
|
* Wire payloads are slim-wire (one new UIMessage per record, on
|
|
* `payload.message`). The legacy `payload.messages` array shape is kept
|
|
* here as a fallback so any historical records on a long-lived session
|
|
* still render.
|
|
*
|
|
* The server wraps records in `{data, id}` and writes `data` as a JSON
|
|
* string; SSE v2 delivers the parsed string back. {@link parseChunkPayload}
|
|
* re-parses to recover the object.
|
|
*/
|
|
type InputStreamChunk = {
|
|
kind?: "message" | "stop";
|
|
payload?: {
|
|
message?: { id?: string; role?: string; parts?: unknown[] };
|
|
messages?: Array<{ id?: string; role?: string; parts?: unknown[] }>;
|
|
trigger?: string;
|
|
};
|
|
message?: string;
|
|
};
|
|
|
|
/**
|
|
* Minimal typing for the chunks we care about on the chat output stream.
|
|
* Covers the AI SDK `UIMessageChunk` variants that `renderPart` actually
|
|
* knows how to display, plus the Trigger.dev control chunks that we filter.
|
|
*/
|
|
type OutputChunk = { type: string; [key: string]: unknown };
|
|
|
|
/**
|
|
* Per-message orchestration state for the output stream accumulator. Mirrors
|
|
* the active-part tracking that AI SDK's `processUIMessageStream` keeps in
|
|
* its `state` object: a registry of streaming text/reasoning parts so deltas
|
|
* can be matched to the right part by id, plus a way to clear them at step
|
|
* boundaries (`finish-step`) so the next step's `text-start`/`reasoning-start`
|
|
* with the same id starts a fresh part instead of appending to the previous
|
|
* step's part.
|
|
*/
|
|
/**
|
|
* Per-message orchestration state — index-based active-part tracking.
|
|
*
|
|
* Each map points from a part id (text or reasoning) to **the index of the
|
|
* currently-streaming part with that id in `message.parts`**. We need
|
|
* indexes (not just a `Set` of "active ids") because part ids are *only
|
|
* unique within a step*: the SDK happily reuses `text-start id="0"` after
|
|
* a `finish-step` boundary. Without index tracking, a `text-delta` for the
|
|
* reused id would have to find the right part by id alone — and a search
|
|
* would match BOTH the previous step's frozen part and the current step's
|
|
* fresh one, which produces a duplication where the previous text gets
|
|
* the new content appended to it AND a fresh part with the same content
|
|
* also appears.
|
|
*
|
|
* Mirrors AI SDK's `processUIMessageStream`'s `state.activeTextParts` /
|
|
* `state.activeReasoningParts` (which hold direct references in the
|
|
* mutating canonical impl). We use indexes here because we do immutable
|
|
* updates and need indices that survive `parts.map()` rewrites — adding
|
|
* new parts and updating existing ones never reorders, so an index is
|
|
* stable for the lifetime of the part.
|
|
*/
|
|
type MessageOrchestrationState = {
|
|
activeTextPartIndexes: Map<string, number>;
|
|
activeReasoningPartIndexes: Map<string, number>;
|
|
};
|
|
|
|
/**
|
|
* `SSEStreamSubscription`'s v2 batch path delivers `parsedBody.data` as-is
|
|
* — but session channels diverge by direction:
|
|
*
|
|
* - `.in`: {@link TriggerChatTransport.serializeInputChunk} writes the
|
|
* `ChatInputChunk` as a JSON **string**, so `data` is a string that
|
|
* needs a second `JSON.parse` to recover the tagged union.
|
|
* - `.out`: the agent's `chatStream.writer(...)` writes
|
|
* {@link UIMessageChunk} **objects** directly; `data` arrives
|
|
* already-parsed.
|
|
*
|
|
* This helper accepts both shapes defensively: a string is parsed; an
|
|
* object is returned as-is. Returns `null` for unparseable payloads.
|
|
*/
|
|
function parseChunkPayload(raw: unknown): Record<string, unknown> | null {
|
|
if (raw == null) return null;
|
|
if (typeof raw === "string") {
|
|
try {
|
|
const parsed = JSON.parse(raw);
|
|
return parsed && typeof parsed === "object" ? (parsed as Record<string, unknown>) : null;
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
if (typeof raw === "object") return raw as Record<string, unknown>;
|
|
return null;
|
|
}
|
|
|
|
function createOrchestrationState(): MessageOrchestrationState {
|
|
return {
|
|
activeTextPartIndexes: new Map(),
|
|
activeReasoningPartIndexes: new Map(),
|
|
};
|
|
}
|
|
|
|
function useAgentSessionMessages({
|
|
sessionId,
|
|
apiOrigin,
|
|
orgSlug,
|
|
projectSlug,
|
|
envSlug,
|
|
initialMessages,
|
|
snapshotPresignedUrl,
|
|
}: {
|
|
sessionId: string;
|
|
apiOrigin: string;
|
|
orgSlug: string;
|
|
projectSlug: string;
|
|
envSlug: string;
|
|
initialMessages: UIMessage[];
|
|
snapshotPresignedUrl?: string;
|
|
}): UIMessage[] {
|
|
// Seed with the user messages from the run's task payload.
|
|
const seedMessages = useMemo(
|
|
() => initialMessages.filter((m) => m.role === "user"),
|
|
[initialMessages]
|
|
);
|
|
|
|
// The snapshot URL is re-signed by the loader on every navigation
|
|
// (tab switches in the inspector pane re-run the session loader),
|
|
// which would otherwise re-trigger the subscription effect below
|
|
// and replay post-snapshot `.out` chunks on top of the messages we
|
|
// already accumulated — duplicating any assistant content that
|
|
// lives past `snapshot.lastOutEventId` (e.g., a canceled run whose
|
|
// turn never completed). Hold the URL behind a ref and keep it
|
|
// out of the effect's deps so the effect runs exactly once per
|
|
// mount.
|
|
const snapshotUrlRef = useRef(snapshotPresignedUrl);
|
|
useEffect(() => {
|
|
snapshotUrlRef.current = snapshotPresignedUrl;
|
|
}, [snapshotPresignedUrl]);
|
|
|
|
// `pendingRef` is the authoritative, eagerly-updated message state:
|
|
// chunks mutate this synchronously as they arrive. A throttled flush
|
|
// copies it into React state so UI updates are capped at ~10x/sec.
|
|
const pendingRef = useRef<Map<string, UIMessage>>(new Map(seedMessages.map((m) => [m.id, m])));
|
|
const timestampsRef = useRef<Map<string, number>>(
|
|
new Map(seedMessages.map((m) => [m.id, INITIAL_PAYLOAD_TIMESTAMP]))
|
|
);
|
|
// Side-table of orchestration state, keyed by assistant message id. Lives
|
|
// outside the UIMessage so React doesn't see it as a renderable prop.
|
|
const orchestrationRef = useRef<Map<string, MessageOrchestrationState>>(new Map());
|
|
|
|
// Buffered HITL resolutions keyed by toolCallId. `addToolApprovalResponse` /
|
|
// `addToolOutput` send a slim assistant message on the `.in` channel carrying
|
|
// just the resolved tool part; the agent never echoes these on `.out`. We
|
|
// stash them here and overlay onto the matching tool part once it exists, so
|
|
// a denial/approval lands regardless of which stream arrives first.
|
|
const pendingResolutionsRef = useRef<Map<string, Record<string, unknown>>>(new Map());
|
|
|
|
// React state snapshot of pendingRef. Only updated via the throttled
|
|
// `scheduleFlush`. The Map *reference* changes on every flush so React
|
|
// detects the state update and the downstream `useMemo` recomputes.
|
|
const [messagesById, setMessagesById] = useState<Map<string, UIMessage>>(
|
|
() => new Map(pendingRef.current)
|
|
);
|
|
|
|
// Throttled flush scheduler — leading edge within a single throttle
|
|
// window: the first chunk after a quiet period flushes immediately, then
|
|
// subsequent chunks coalesce until the next window opens.
|
|
const lastFlushAtRef = useRef<number>(0);
|
|
const pendingTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
|
|
const scheduleFlush = useRef<() => void>(() => {});
|
|
scheduleFlush.current = () => {
|
|
if (pendingTimerRef.current !== null) return; // already scheduled
|
|
const now = Date.now();
|
|
const sinceLast = now - lastFlushAtRef.current;
|
|
const delay = Math.max(0, STATE_FLUSH_THROTTLE_MS - sinceLast);
|
|
pendingTimerRef.current = setTimeout(() => {
|
|
pendingTimerRef.current = null;
|
|
lastFlushAtRef.current = Date.now();
|
|
setMessagesById(new Map(pendingRef.current));
|
|
}, delay);
|
|
};
|
|
|
|
useEffect(() => {
|
|
const abort = new AbortController();
|
|
|
|
// Overlay a buffered HITL resolution (approval/output delivered on `.in`)
|
|
// onto the matching tool part. Returns true if a part changed. Safe to call
|
|
// repeatedly — after each `.out` tool chunk and whenever a `.in` resolution
|
|
// arrives — so the resolution lands regardless of cross-stream ordering.
|
|
// Never downgrades a part that already reached a terminal output state (so
|
|
// an approved-then-executed tool keeps its `output-available` + output).
|
|
const applyToolResolution = (toolCallId: string): boolean => {
|
|
const res = pendingResolutionsRef.current.get(toolCallId);
|
|
if (!res) return false;
|
|
for (const [mid, msg] of pendingRef.current) {
|
|
const parts = (msg.parts ?? []) as Array<Record<string, unknown>>;
|
|
const idx = parts.findIndex(
|
|
(p) => (p as { toolCallId?: string }).toolCallId === toolCallId
|
|
);
|
|
if (idx < 0) continue;
|
|
const cur = parts[idx]!;
|
|
const terminal =
|
|
cur.state === "output-available" ||
|
|
cur.state === "output-error" ||
|
|
cur.state === "output-denied";
|
|
const nextState = res.state != null && !terminal ? res.state : cur.state;
|
|
const sameApproval = JSON.stringify(cur.approval) === JSON.stringify(res.approval);
|
|
if (
|
|
nextState === cur.state &&
|
|
sameApproval &&
|
|
res.output === undefined &&
|
|
res.errorText === undefined
|
|
) {
|
|
// Already applied. If the part has reached a terminal state, drop
|
|
// the buffered resolution so the Map doesn't grow unbounded on
|
|
// long-lived sessions with many tool calls.
|
|
if (terminal) pendingResolutionsRef.current.delete(toolCallId);
|
|
return false;
|
|
}
|
|
const next = parts.slice();
|
|
next[idx] = {
|
|
...cur,
|
|
...(res.approval != null ? { approval: res.approval } : {}),
|
|
...(res.output !== undefined ? { output: res.output } : {}),
|
|
...(res.errorText !== undefined ? { errorText: res.errorText } : {}),
|
|
state: nextState,
|
|
};
|
|
pendingRef.current.set(mid, { ...msg, parts: next } as UIMessage);
|
|
// Drop the buffered entry once the part has landed at a terminal
|
|
// state — no future `.out` chunk will need this resolution.
|
|
const reachedTerminal =
|
|
nextState === "output-available" ||
|
|
nextState === "output-error" ||
|
|
nextState === "output-denied";
|
|
if (reachedTerminal) pendingResolutionsRef.current.delete(toolCallId);
|
|
return true;
|
|
}
|
|
return false;
|
|
};
|
|
|
|
const encodedSession = encodeURIComponent(sessionId);
|
|
// Always use the page's own origin to avoid CORS preflight failures
|
|
// when the configured `apiOrigin` (e.g. `localhost`) differs from the
|
|
// origin the dashboard was loaded from (e.g. `127.0.0.1`). The dashboard
|
|
// resource route is same-origin by construction.
|
|
const origin = typeof window !== "undefined" ? window.location.origin : apiOrigin;
|
|
const sessionBase =
|
|
`${origin}/resources/orgs/${orgSlug}/projects/${projectSlug}/env/${envSlug}` +
|
|
`/sessions/${encodedSession}/realtime/v1`;
|
|
|
|
const outputUrl = `${sessionBase}/out`;
|
|
const inputUrl = `${sessionBase}/in`;
|
|
|
|
/**
|
|
* Try to seed `pendingRef` from the agent's S3 snapshot blob and return
|
|
* the snapshot's `lastOutEventId` so the `.out` SSE subscription resumes
|
|
* just past the snapshot. Returns undefined for sessions that don't
|
|
* have a snapshot (e.g. `hydrateMessages` customers, or sessions that
|
|
* have never completed a turn).
|
|
*/
|
|
const loadSnapshot = async (): Promise<string | undefined> => {
|
|
const url = snapshotUrlRef.current;
|
|
if (!url) return undefined;
|
|
try {
|
|
const resp = await fetch(url, { signal: abort.signal });
|
|
if (!resp.ok) return undefined;
|
|
const json = (await resp.json()) as unknown;
|
|
const parsed = ChatSnapshotV1Schema.safeParse(json);
|
|
if (!parsed.success) return undefined;
|
|
const snapshot = parsed.data;
|
|
// Preserve the snapshot's array order in the final render by
|
|
// giving each message a unique, monotonically increasing
|
|
// timestamp from `(savedAt - count + index)`. Real chunk
|
|
// timestamps from the SSE path use S2 arrival ms (positive
|
|
// numbers in the present), so anything below `savedAt` sorts
|
|
// before live chunks while preserving snapshot order among
|
|
// themselves.
|
|
const count = snapshot.messages.length;
|
|
snapshot.messages.forEach((raw, i) => {
|
|
const message = raw as UIMessage;
|
|
if (!message?.id) return;
|
|
// The snapshot's seed wins over the task-payload seed for any
|
|
// overlapping ids (the snapshot represents the agent's
|
|
// canonical accumulator, post-turn).
|
|
pendingRef.current.set(message.id, message);
|
|
if (!timestampsRef.current.has(message.id)) {
|
|
timestampsRef.current.set(message.id, snapshot.savedAt - count + i);
|
|
}
|
|
});
|
|
scheduleFlush.current();
|
|
return snapshot.lastOutEventId;
|
|
} catch {
|
|
// 404 / network / parse / abort — fall back to seq=0 SSE
|
|
return undefined;
|
|
}
|
|
};
|
|
|
|
const outputSubOptions = (lastEventId: string | undefined) =>
|
|
({
|
|
signal: abort.signal,
|
|
timeoutInSeconds: 120,
|
|
...(lastEventId !== undefined ? { lastEventId } : {}),
|
|
}) as const;
|
|
|
|
const commonSubOptions = {
|
|
signal: abort.signal,
|
|
timeoutInSeconds: 120,
|
|
} as const;
|
|
|
|
// ---- Output stream: assistant messages ---------------------------------
|
|
//
|
|
// The output stream delivers data records (UIMessageChunks) interleaved
|
|
// with Trigger control records (`turn-complete`, `upgrade-required`) and
|
|
// S2 command records (`trim`). Control + command records ride on
|
|
// `record.headers` with empty bodies; the SSE parser strips S2 command
|
|
// records entirely, and control records arrive with `value.chunk ===
|
|
// undefined`, which `parseChunkPayload` drops below.
|
|
//
|
|
// We fold everything else into an assistant `UIMessage` via our own
|
|
// `applyOutputChunk` accumulator — the AI SDK's `readUIMessageStream`
|
|
// helper is only available in `ai@6`, and the webapp is pinned to
|
|
// `ai@4`, so we re-implement just the chunk types that `renderPart`
|
|
// actually displays.
|
|
//
|
|
// We capture the **server timestamp of each assistant message's first
|
|
// `start` chunk** so later sort-by-timestamp merges with the input
|
|
// stream correctly.
|
|
const runOutput = async () => {
|
|
try {
|
|
// Seed messages from the snapshot first (if available), then
|
|
// resume the SSE from the snapshot's last event id so we don't
|
|
// re-stream chunks already represented in the snapshot. If no
|
|
// snapshot exists (no URL, 404, parse failure), the SSE opens
|
|
// at seq=0 — which, post-trim, contains roughly one turn of
|
|
// records (acceptable fallback for `hydrateMessages` sessions
|
|
// and fresh sessions).
|
|
const snapshotLastEventId = await loadSnapshot();
|
|
if (abort.signal.aborted) return;
|
|
|
|
const sub = new SSEStreamSubscription(outputUrl, outputSubOptions(snapshotLastEventId));
|
|
const raw = await sub.subscribe();
|
|
const reader = raw.getReader();
|
|
|
|
let currentMessageId: string | null = null;
|
|
|
|
try {
|
|
while (!abort.signal.aborted) {
|
|
const { done, value } = await reader.read();
|
|
if (done) return;
|
|
|
|
const chunk = parseChunkPayload(value.chunk) as OutputChunk | null;
|
|
if (!chunk || typeof chunk.type !== "string") continue;
|
|
// Legacy belt-and-suspenders: prior versions of the SDK
|
|
// emitted `trigger:turn-complete` / `trigger:upgrade-required`
|
|
// as data records (`type` field). Current versions use
|
|
// header-form control records, which `parseChunkPayload`
|
|
// drops above. Keep this filter to handle any in-flight
|
|
// sessions whose `.out` was populated by the older SDK.
|
|
if (chunk.type.startsWith("trigger:")) continue;
|
|
|
|
if (chunk.type === "start") {
|
|
const messageId =
|
|
typeof chunk.messageId === "string" && chunk.messageId.length > 0
|
|
? chunk.messageId
|
|
: `asst-${crypto.randomUUID()}`;
|
|
currentMessageId = messageId;
|
|
|
|
if (!timestampsRef.current.has(messageId)) {
|
|
timestampsRef.current.set(messageId, value.timestamp);
|
|
}
|
|
|
|
const existing = pendingRef.current.get(messageId);
|
|
if (existing) {
|
|
// Same message id seen again — merge metadata only, keep
|
|
// existing parts (canonical `processUIMessageStream` does
|
|
// the same on a repeated `start`).
|
|
if (chunk.messageMetadata != null) {
|
|
pendingRef.current.set(messageId, {
|
|
...existing,
|
|
metadata: {
|
|
...((existing as { metadata?: Record<string, unknown> }).metadata ?? {}),
|
|
...(chunk.messageMetadata as Record<string, unknown>),
|
|
},
|
|
} as UIMessage);
|
|
scheduleFlush.current();
|
|
}
|
|
} else {
|
|
const message: UIMessage = {
|
|
id: messageId,
|
|
role: "assistant",
|
|
parts: [],
|
|
...(chunk.messageMetadata != null
|
|
? { metadata: chunk.messageMetadata as UIMessage["metadata"] }
|
|
: {}),
|
|
} as UIMessage;
|
|
pendingRef.current.set(messageId, message);
|
|
orchestrationRef.current.set(messageId, createOrchestrationState());
|
|
scheduleFlush.current();
|
|
}
|
|
continue;
|
|
}
|
|
|
|
if (currentMessageId === null) continue;
|
|
const existing = pendingRef.current.get(currentMessageId);
|
|
if (!existing) continue;
|
|
let orchestration = orchestrationRef.current.get(currentMessageId);
|
|
if (!orchestration) {
|
|
// Defensive: a chunk arrived for a message we never saw a
|
|
// `start` for. Lazily create orchestration state so we can
|
|
// still display the parts.
|
|
orchestration = createOrchestrationState();
|
|
orchestrationRef.current.set(currentMessageId, orchestration);
|
|
}
|
|
|
|
const updated = applyOutputChunk(existing, chunk, orchestration);
|
|
if (updated !== existing) {
|
|
pendingRef.current.set(currentMessageId, updated);
|
|
scheduleFlush.current();
|
|
}
|
|
|
|
// A `.out` chunk just established/updated a tool part — (re)apply any
|
|
// buffered `.in` resolution for it. Covers the `.in`-before-`.out`
|
|
// order and corrects a `.out` chunk that downgraded the state (e.g.
|
|
// a replayed `tool-approval-request` arriving after the denial).
|
|
const outToolCallId = (chunk as { toolCallId?: string }).toolCallId;
|
|
if (typeof outToolCallId === "string" && applyToolResolution(outToolCallId)) {
|
|
scheduleFlush.current();
|
|
}
|
|
}
|
|
} finally {
|
|
try {
|
|
reader.releaseLock();
|
|
} catch {
|
|
// Lock may already be released.
|
|
}
|
|
}
|
|
} catch (err) {
|
|
if (abort.signal.aborted) return;
|
|
// eslint-disable-next-line no-console
|
|
console.debug("[AgentView] output stream subscription failed", err);
|
|
}
|
|
};
|
|
|
|
// ---- Input channel: user messages (`ChatInputChunk`) -------------------
|
|
//
|
|
// The transport appends a `{kind: "message", payload}` ChatInputChunk
|
|
// for every user turn (and `{kind: "stop"}` for stop signals). We pull
|
|
// user messages out of `payload.messages` for `kind: "message"` chunks
|
|
// and ignore the rest.
|
|
const runInput = async () => {
|
|
try {
|
|
const sub = new SSEStreamSubscription(inputUrl, commonSubOptions);
|
|
const raw = await sub.subscribe();
|
|
const reader = raw.getReader();
|
|
try {
|
|
while (!abort.signal.aborted) {
|
|
const { done, value } = await reader.read();
|
|
if (done) return;
|
|
|
|
const chunk = parseChunkPayload(value.chunk) as InputStreamChunk | null;
|
|
if (!chunk || chunk.kind !== "message") continue;
|
|
const payload = chunk.payload;
|
|
if (!payload) continue;
|
|
|
|
// Slim-wire is one UIMessage on `payload.message`; legacy
|
|
// payloads carried an array on `payload.messages`. Accept
|
|
// either so historical records on a long-lived session still
|
|
// render.
|
|
const candidates = Array.isArray(payload.messages)
|
|
? payload.messages
|
|
: payload.message
|
|
? [payload.message]
|
|
: [];
|
|
|
|
let changed = false;
|
|
|
|
// New user turns — merge in (dedupe by id).
|
|
for (const m of candidates) {
|
|
if (
|
|
m == null ||
|
|
(m as { role?: string }).role !== "user" ||
|
|
typeof m.id !== "string"
|
|
) {
|
|
continue;
|
|
}
|
|
if (pendingRef.current.has(m.id)) continue;
|
|
pendingRef.current.set(m.id, m as UIMessage);
|
|
timestampsRef.current.set(m.id, value.timestamp);
|
|
changed = true;
|
|
}
|
|
|
|
// HITL resolutions ride on `.in` as a slim *assistant* message
|
|
// carrying just the resolved tool part (state + approval/output).
|
|
// Buffer each by toolCallId and overlay onto the matching tool part
|
|
// (which usually arrived on `.out` as `tool-approval-request`).
|
|
for (const m of candidates) {
|
|
if (m == null || (m as { role?: string }).role !== "assistant") continue;
|
|
const parts = (m as { parts?: unknown[] }).parts;
|
|
if (!Array.isArray(parts)) continue;
|
|
for (const sp of parts) {
|
|
const part = sp as Record<string, unknown>;
|
|
if (typeof part.type !== "string" || !part.type.startsWith("tool-")) continue;
|
|
const tcId = (part as { toolCallId?: string }).toolCallId;
|
|
if (typeof tcId !== "string") continue;
|
|
pendingResolutionsRef.current.set(tcId, part);
|
|
if (applyToolResolution(tcId)) changed = true;
|
|
}
|
|
}
|
|
|
|
if (changed) scheduleFlush.current();
|
|
}
|
|
} finally {
|
|
try {
|
|
reader.releaseLock();
|
|
} catch {
|
|
// Lock may already be released.
|
|
}
|
|
}
|
|
} catch (err) {
|
|
if (abort.signal.aborted) return;
|
|
// eslint-disable-next-line no-console
|
|
console.debug("[AgentView] input stream subscription failed", err);
|
|
}
|
|
};
|
|
|
|
void runOutput();
|
|
void runInput();
|
|
|
|
return () => {
|
|
abort.abort();
|
|
if (pendingTimerRef.current !== null) {
|
|
clearTimeout(pendingTimerRef.current);
|
|
pendingTimerRef.current = null;
|
|
}
|
|
};
|
|
// `snapshotPresignedUrl` is intentionally NOT in this dep list — see
|
|
// `snapshotUrlRef` above for the reasoning. Including it caused the
|
|
// subscription to tear down + replay on every inspector tab click,
|
|
// which appended duplicate parts to any assistant message whose
|
|
// chunks lived past `snapshot.lastOutEventId`.
|
|
// eslint-disable-next-line react-hooks/exhaustive-deps
|
|
}, [sessionId, apiOrigin, orgSlug, projectSlug, envSlug]);
|
|
|
|
return useMemo(() => {
|
|
const timestamps = timestampsRef.current;
|
|
const arr = Array.from(messagesById.values());
|
|
arr.sort((a, b) => {
|
|
const ta = timestamps.get(a.id) ?? 0;
|
|
const tb = timestamps.get(b.id) ?? 0;
|
|
if (ta !== tb) return ta - tb;
|
|
// Tie-breaker for messages sharing a stream ID bucket (rare): fall
|
|
// back to message id string order so the output is deterministic.
|
|
return a.id < b.id ? -1 : a.id > b.id ? 1 : 0;
|
|
});
|
|
return arr;
|
|
}, [messagesById]);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// applyOutputChunk — minimal UIMessageChunk → UIMessage accumulator.
|
|
// ---------------------------------------------------------------------------
|
|
//
|
|
// A pared-down re-implementation of AI SDK's `processUIMessageStream` (in
|
|
// `ai@6`'s `index.mjs`). The webapp is pinned to `ai@4`, which doesn't ship
|
|
// the v5+ chunk-stream helpers, so we vendor the bits we actually use.
|
|
//
|
|
// Scope vs. canonical:
|
|
// - We render only the chunk shapes that `AgentMessageView`/`renderPart`
|
|
// actually display: text, reasoning, tool-* (input-{start,delta,available}
|
|
// + output-{available,error}), source-url, source-document, file,
|
|
// step-start/finish-step, data-*, plus metadata/finish lifecycle.
|
|
// - Unknown chunk types fall through as no-ops — defensive on purpose for a
|
|
// read-only viewer.
|
|
// - We **do not parse partial JSON for streaming tool inputs.** Canonical
|
|
// uses `parsePartialJson` (which depends on a 300-line `fixJson` state
|
|
// machine to repair incomplete JSON) so users see the input growing
|
|
// character-by-character. We skip it: tool inputs stay `undefined`
|
|
// throughout streaming and snap to the final value when
|
|
// `tool-input-available` lands. Acceptable for a viewer; can be added
|
|
// later by vendoring `fixJson` if the UX warrants it.
|
|
//
|
|
// `orchestration` carries per-message active-part trackers that mirror
|
|
// canonical's `state.activeTextParts` / `state.activeReasoningParts`. They
|
|
// let `text-delta` find the right text part by id and let `finish-step`
|
|
// clear them so a new step can re-use the same id without colliding.
|
|
//
|
|
// Returns the same object reference when nothing changes so the caller can
|
|
// skip unnecessary state flushes + React re-renders.
|
|
|
|
type AnyPart = { [key: string]: unknown; type: string };
|
|
|
|
function applyOutputChunk(
|
|
msg: UIMessage,
|
|
chunk: OutputChunk,
|
|
orchestration: MessageOrchestrationState
|
|
): UIMessage {
|
|
const type = chunk.type;
|
|
|
|
// Text parts ---------------------------------------------------------------
|
|
//
|
|
// Track each streaming text part by its index in `msg.parts`. Part ids
|
|
// are only unique *within a step* — the SDK happily reuses `text-start
|
|
// id="0"` after a `finish-step` boundary — so a delta arriving for a
|
|
// reused id needs to land on the *current* part, not every prior part
|
|
// that ever shared that id. The index map gives us O(1) "which slot is
|
|
// currently streaming this id" without any id-based search.
|
|
if (type === "text-start") {
|
|
const id = chunk.id as string;
|
|
const newIndex = (msg.parts ?? []).length; // index AFTER push
|
|
orchestration.activeTextPartIndexes.set(id, newIndex);
|
|
return withNewPart(msg, {
|
|
type: "text",
|
|
id,
|
|
text: "",
|
|
state: "streaming",
|
|
});
|
|
}
|
|
if (type === "text-delta") {
|
|
const id = chunk.id as string;
|
|
const index = orchestration.activeTextPartIndexes.get(id);
|
|
if (index === undefined) return msg; // delta with no start — drop.
|
|
return updatePartAt(msg, index, (p) => ({
|
|
...p,
|
|
text: ((p as { text?: string }).text ?? "") + String(chunk.delta ?? ""),
|
|
}));
|
|
}
|
|
if (type === "text-end") {
|
|
const id = chunk.id as string;
|
|
const index = orchestration.activeTextPartIndexes.get(id);
|
|
if (index === undefined) return msg;
|
|
orchestration.activeTextPartIndexes.delete(id);
|
|
return updatePartAt(msg, index, (p) => ({ ...p, state: "done" }));
|
|
}
|
|
|
|
// Reasoning parts ----------------------------------------------------------
|
|
if (type === "reasoning-start") {
|
|
const id = chunk.id as string;
|
|
const newIndex = (msg.parts ?? []).length;
|
|
orchestration.activeReasoningPartIndexes.set(id, newIndex);
|
|
return withNewPart(msg, {
|
|
type: "reasoning",
|
|
id,
|
|
text: "",
|
|
state: "streaming",
|
|
});
|
|
}
|
|
if (type === "reasoning-delta") {
|
|
const id = chunk.id as string;
|
|
const index = orchestration.activeReasoningPartIndexes.get(id);
|
|
if (index === undefined) return msg;
|
|
return updatePartAt(msg, index, (p) => ({
|
|
...p,
|
|
text: ((p as { text?: string }).text ?? "") + String(chunk.delta ?? ""),
|
|
}));
|
|
}
|
|
if (type === "reasoning-end") {
|
|
const id = chunk.id as string;
|
|
const index = orchestration.activeReasoningPartIndexes.get(id);
|
|
if (index === undefined) return msg;
|
|
orchestration.activeReasoningPartIndexes.delete(id);
|
|
return updatePartAt(msg, index, (p) => ({ ...p, state: "done" }));
|
|
}
|
|
|
|
// Tool call parts ----------------------------------------------------------
|
|
if (type === "tool-input-start") {
|
|
const toolName = String(chunk.toolName ?? "");
|
|
return withNewPart(msg, {
|
|
type: `tool-${toolName}`,
|
|
toolCallId: chunk.toolCallId,
|
|
toolName,
|
|
state: "input-streaming",
|
|
input: undefined,
|
|
});
|
|
}
|
|
if (type === "tool-input-delta") {
|
|
// We don't parse partial JSON, so streaming tool input deltas are a
|
|
// no-op. The full input snaps in when `tool-input-available` arrives.
|
|
return msg;
|
|
}
|
|
if (type === "tool-input-available") {
|
|
const toolName = String(chunk.toolName ?? "");
|
|
const existingIdx = indexOfPart(
|
|
msg,
|
|
(p) => (p as { toolCallId?: string }).toolCallId === chunk.toolCallId
|
|
);
|
|
if (existingIdx >= 0) {
|
|
return updatePartAt(msg, existingIdx, (p) => ({
|
|
...p,
|
|
state: "input-available",
|
|
input: chunk.input,
|
|
}));
|
|
}
|
|
// Tool input arrived without a preceding tool-input-start (some
|
|
// providers do this for fast tools) — synthesize a new part.
|
|
return withNewPart(msg, {
|
|
type: `tool-${toolName}`,
|
|
toolCallId: chunk.toolCallId,
|
|
toolName,
|
|
state: "input-available",
|
|
input: chunk.input,
|
|
});
|
|
}
|
|
if (type === "tool-output-available") {
|
|
return updatePart(msg, (p) =>
|
|
(p as { toolCallId?: string }).toolCallId === chunk.toolCallId
|
|
? {
|
|
...p,
|
|
state: "output-available",
|
|
output: chunk.output,
|
|
...(chunk.preliminary === true ? { preliminary: true } : {}),
|
|
}
|
|
: null
|
|
);
|
|
}
|
|
if (type === "tool-output-error") {
|
|
return updatePart(msg, (p) =>
|
|
(p as { toolCallId?: string }).toolCallId === chunk.toolCallId
|
|
? { ...p, state: "output-error", errorText: chunk.errorText }
|
|
: null
|
|
);
|
|
}
|
|
|
|
// HITL approval (AI SDK 7) -------------------------------------------------
|
|
//
|
|
// v7 added human-in-the-loop tool approval. A `needsApproval` tool emits a
|
|
// `tool-approval-request` after its input is available; the tool then waits
|
|
// for a `tool-approval-response` (approve/deny) before executing. Mirror AI
|
|
// SDK 7's `processUIMessageStream`: the request marks the matching part
|
|
// `approval-requested` and records `approval.id`; the response (matched by
|
|
// that id) marks it `approval-responded` with the verdict. An approved tool
|
|
// then proceeds to `tool-output-available` as usual.
|
|
if (type === "tool-approval-request") {
|
|
return updatePart(msg, (p) =>
|
|
(p as { toolCallId?: string }).toolCallId === chunk.toolCallId
|
|
? {
|
|
...p,
|
|
state: "approval-requested",
|
|
approval: {
|
|
id: chunk.approvalId,
|
|
...(chunk.isAutomatic === true ? { isAutomatic: true } : {}),
|
|
},
|
|
}
|
|
: null
|
|
);
|
|
}
|
|
if (type === "tool-approval-response") {
|
|
return updatePart(msg, (p) => {
|
|
const approval = (p as { approval?: { id?: string; isAutomatic?: boolean } }).approval;
|
|
if (!approval || approval.id !== chunk.approvalId) return null;
|
|
return {
|
|
...p,
|
|
state: "approval-responded",
|
|
approval: {
|
|
...approval,
|
|
id: chunk.approvalId,
|
|
approved: chunk.approved,
|
|
...(chunk.reason != null ? { reason: chunk.reason } : {}),
|
|
},
|
|
};
|
|
});
|
|
}
|
|
|
|
// Source / file / step / data parts — pass through as a whole -------------
|
|
if (type === "source-url" || type === "source-document" || type === "file") {
|
|
return withNewPart(msg, chunk as unknown as AnyPart);
|
|
}
|
|
if (type === "start-step") {
|
|
return withNewPart(msg, { type: "step-start" });
|
|
}
|
|
if (type === "finish-step") {
|
|
// Step boundary — canonical clears the active part trackers so a new
|
|
// step can re-use the same text/reasoning part IDs cleanly. The
|
|
// message itself doesn't structurally change; the previous step's
|
|
// parts stay frozen at their indexes in `msg.parts`.
|
|
orchestration.activeTextPartIndexes.clear();
|
|
orchestration.activeReasoningPartIndexes.clear();
|
|
return msg;
|
|
}
|
|
if (type.startsWith("data-")) {
|
|
return withNewPart(msg, chunk as unknown as AnyPart);
|
|
}
|
|
|
|
// Metadata / lifecycle -----------------------------------------------------
|
|
if (type === "finish" || type === "message-metadata") {
|
|
if (chunk.messageMetadata == null) return msg;
|
|
return {
|
|
...msg,
|
|
metadata: {
|
|
...((msg as { metadata?: Record<string, unknown> }).metadata ?? {}),
|
|
...(chunk.messageMetadata as Record<string, unknown>),
|
|
},
|
|
} as UIMessage;
|
|
}
|
|
|
|
// Abort / error / unknown — no structural change. (`start` is handled at
|
|
// the orchestration level in the output reader, not here.)
|
|
return msg;
|
|
}
|
|
|
|
// --- Small immutable helpers for UIMessage.parts mutation -------------------
|
|
|
|
function withNewPart(msg: UIMessage, part: AnyPart): UIMessage {
|
|
return {
|
|
...msg,
|
|
parts: [...((msg.parts ?? []) as AnyPart[]), part],
|
|
} as UIMessage;
|
|
}
|
|
|
|
function updatePart(msg: UIMessage, updater: (part: AnyPart) => AnyPart | null): UIMessage {
|
|
const parts = (msg.parts ?? []) as AnyPart[];
|
|
let changed = false;
|
|
const next = parts.map((p) => {
|
|
const updated = updater(p);
|
|
if (updated === null) return p;
|
|
changed = true;
|
|
return updated;
|
|
});
|
|
return changed ? ({ ...msg, parts: next } as UIMessage) : msg;
|
|
}
|
|
|
|
function indexOfPart(msg: UIMessage, predicate: (part: AnyPart) => boolean): number {
|
|
const parts = (msg.parts ?? []) as AnyPart[];
|
|
for (let i = 0; i < parts.length; i++) {
|
|
if (predicate(parts[i]!)) return i;
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
function updatePartAt(
|
|
msg: UIMessage,
|
|
index: number,
|
|
updater: (part: AnyPart) => AnyPart
|
|
): UIMessage {
|
|
const parts = (msg.parts ?? []) as AnyPart[];
|
|
if (index < 0 || index >= parts.length) return msg;
|
|
const next = parts.slice();
|
|
next[index] = updater(parts[index]!);
|
|
return { ...msg, parts: next } as UIMessage;
|
|
}
|