diff --git a/apps/mesh/src/api/routes/decopilot/project-chunks.test.ts b/apps/mesh/src/api/routes/decopilot/project-chunks.test.ts index b718f4699a..899f7b4e8c 100644 --- a/apps/mesh/src/api/routes/decopilot/project-chunks.test.ts +++ b/apps/mesh/src/api/routes/decopilot/project-chunks.test.ts @@ -440,4 +440,64 @@ describe("projectChunks", () => { }), ).toBe(true); }); + + test("ignores data-user-message chunks when projecting assistant parts", async () => { + const emitted: Array<{ id: string; parts?: unknown[] }> = []; + + await projectChunks({ + chunks: (async function* () { + yield { + type: "data-user-message", + data: { + id: "u-1", + role: "user", + parts: [{ type: "text", text: "hi" }], + }, + } as unknown as UIMessageChunk; + yield { type: "start", messageId: "m-1" } as UIMessageChunk; + yield { type: "text-start", id: "txt" } as UIMessageChunk; + yield { + type: "text-delta", + id: "txt", + delta: "hello", + } as UIMessageChunk; + yield { type: "text-end", id: "txt" } as UIMessageChunk; + yield { type: "finish", finishReason: "stop" } as UIMessageChunk; + })(), + persistence: { + emitStepParts: async (message) => { + emitted.push({ id: message.id, parts: message.parts }); + }, + emitFinal: async (message) => { + emitted.push({ id: message.id, parts: message.parts }); + }, + emitError: async () => {}, + }, + }); + + // The user-message chunk must not leak into any projected assistant part, + // and the assistant text still folds normally. + expect( + emitted + .flatMap((message) => message.parts ?? []) + .some( + (part) => + typeof part === "object" && + part !== null && + (part as { type?: unknown }).type === "data-user-message", + ), + ).toBe(false); + expect(emitted.some((m) => m.id === "u-1")).toBe(false); + expect( + emitted + .flatMap((message) => message.parts ?? []) + .some( + (part) => + typeof part === "object" && + part !== null && + (part as { type?: unknown; text?: unknown }).type === "text" && + (part as { text?: unknown }).text === "hello", + ), + ).toBe(true); + }); }); diff --git a/apps/mesh/src/api/routes/decopilot/project-chunks.ts b/apps/mesh/src/api/routes/decopilot/project-chunks.ts index e06308b553..910bc65841 100644 --- a/apps/mesh/src/api/routes/decopilot/project-chunks.ts +++ b/apps/mesh/src/api/routes/decopilot/project-chunks.ts @@ -5,6 +5,15 @@ import { type HarnessUsage, } from "./consume-harness-stream"; import { isRunStatusControlChunk } from "./run-status-stage"; +import { isUserMessageControlChunk } from "./user-message-stream"; + +/** Transient control chunks that ride the run stream for the live UI but must + * never be folded into the durable assistant projection. `fenceFilter` + * already drops the fence-less ones, but filtering by type here keeps the + * projection correct even if one is ever published with a run fence. */ +function isTransientControlChunk(chunk: unknown): boolean { + return isRunStatusControlChunk(chunk) || isUserMessageControlChunk(chunk); +} /** * Title persistence for the projector. When supplied, the projector becomes the @@ -192,7 +201,7 @@ export async function projectChunks( ? (async function* () { try { for await (const chunk of options.chunks!) { - if (isRunStatusControlChunk(chunk)) continue; + if (isTransientControlChunk(chunk)) continue; yield chunk; } } catch (e) { @@ -207,7 +216,7 @@ export async function projectChunks( }).pipeThrough( new TransformStream({ transform(chunk, controller) { - if (!isRunStatusControlChunk(chunk)) controller.enqueue(chunk); + if (!isTransientControlChunk(chunk)) controller.enqueue(chunk); }, }), ) diff --git a/apps/mesh/src/api/routes/decopilot/routes.ts b/apps/mesh/src/api/routes/decopilot/routes.ts index 93e7d208f2..68b18c9832 100644 --- a/apps/mesh/src/api/routes/decopilot/routes.ts +++ b/apps/mesh/src/api/routes/decopilot/routes.ts @@ -45,6 +45,7 @@ import { publishRunStatusStage, shouldPublishClusterRunStatus, } from "./run-status-stage"; +import { publishUserMessage } from "./user-message-stream"; import { wrapWithSseKeepalive } from "./sse-keepalive"; import { resolveDispatchTarget } from "../../../links/resolve-dispatch-target"; import { @@ -746,6 +747,10 @@ export function createDecopilotRoutes(deps: DecopilotDeps) { }) ) { await emitter.emitRequestMessage(persistedRequestMessage); + // Mirror the prompt onto the run stream so OTHER viewers of a shared + // thread render it live, not only after a DB refetch. Best-effort and + // published before the run's assistant chunks so it sorts first. + await publishUserMessage(streamBuffer, taskId, persistedRequestMessage); } const serializableRequest = buildDurableDispatchInput(input, { diff --git a/apps/mesh/src/api/routes/decopilot/user-message-stream.test.ts b/apps/mesh/src/api/routes/decopilot/user-message-stream.test.ts new file mode 100644 index 0000000000..e0fba7b19a --- /dev/null +++ b/apps/mesh/src/api/routes/decopilot/user-message-stream.test.ts @@ -0,0 +1,74 @@ +import { describe, expect, test } from "bun:test"; +import type { UIMessage } from "ai"; +import { + buildUserMessageChunk, + isUserMessageControlChunk, + publishUserMessage, + USER_MESSAGE_CHUNK_TYPE, +} from "./user-message-stream"; + +const MESSAGE = { + id: "u-1", + role: "user", + parts: [{ type: "text", text: "hello" }], +} as unknown as UIMessage; + +describe("buildUserMessageChunk", () => { + test("wraps the message in a data-user-message chunk", () => { + expect(buildUserMessageChunk(MESSAGE)).toEqual({ + type: USER_MESSAGE_CHUNK_TYPE, + data: MESSAGE, + }); + }); +}); + +describe("isUserMessageControlChunk", () => { + test("matches the control chunk and nothing else", () => { + expect(isUserMessageControlChunk(buildUserMessageChunk(MESSAGE))).toBe( + true, + ); + expect(isUserMessageControlChunk({ type: "data-run-status" })).toBe(false); + expect(isUserMessageControlChunk({ type: "text-delta" })).toBe(false); + expect(isUserMessageControlChunk(null)).toBe(false); + expect(isUserMessageControlChunk("data-user-message")).toBe(false); + }); +}); + +describe("publishUserMessage", () => { + test("publishes the chunk via streamBuffer.publishRawChunk", async () => { + const calls: Array<{ taskId: string; chunk: unknown }> = []; + await publishUserMessage( + { + publishRawChunk: async (taskId, chunk) => { + calls.push({ taskId, chunk }); + return true; + }, + }, + "task-1", + MESSAGE, + ); + expect(calls).toEqual([ + { taskId: "task-1", chunk: buildUserMessageChunk(MESSAGE) }, + ]); + }); + + test("is a no-op without a stream buffer", async () => { + await expect( + publishUserMessage(undefined, "task-1", MESSAGE), + ).resolves.toBeUndefined(); + }); + + test("swallows publish errors (best-effort)", async () => { + await expect( + publishUserMessage( + { + publishRawChunk: async () => { + throw new Error("nats down"); + }, + }, + "task-1", + MESSAGE, + ), + ).resolves.toBeUndefined(); + }); +}); diff --git a/apps/mesh/src/api/routes/decopilot/user-message-stream.ts b/apps/mesh/src/api/routes/decopilot/user-message-stream.ts new file mode 100644 index 0000000000..98926ec2e1 --- /dev/null +++ b/apps/mesh/src/api/routes/decopilot/user-message-stream.ts @@ -0,0 +1,48 @@ +import type { UIMessage, UIMessageChunk } from "ai"; +import type { StreamBuffer } from "./stream-buffer"; + +/** + * Live mirror of a just-posted USER prompt onto the run stream. + * + * The run's JetStream subject carries only assistant/harness output, so a + * SECOND viewer of a shared thread never saw the other user's prompt until a DB + * refetch. We publish the materialized request message as a transient control + * chunk at POST time (before the run's assistant chunks) so every tailing + * viewer renders it live. It is deliberately fence-less and NOT persisted from + * the stream — the durable copy is the `emitRequestMessage` DB write; the + * projector filters this chunk out (see `project-chunks.ts`). + * + * Large prompts are handled for free: `publishRawChunk` → `serializeChunk` + * fragments anything over `MAX_PUBLISH_BYTES` and drops over `MAX_CHUNKED_BYTES`. + */ +export const USER_MESSAGE_CHUNK_TYPE = "data-user-message"; + +export type UserMessageChunk = Extract< + UIMessageChunk, + { type: `data-${string}` } +> & { + type: typeof USER_MESSAGE_CHUNK_TYPE; + data: UIMessage; +}; + +export function buildUserMessageChunk(message: UIMessage): UserMessageChunk { + return { type: USER_MESSAGE_CHUNK_TYPE, data: message } as UserMessageChunk; +} + +export function isUserMessageControlChunk(chunk: unknown): boolean { + if (!chunk || typeof chunk !== "object") return false; + return (chunk as { type?: unknown }).type === USER_MESSAGE_CHUNK_TYPE; +} + +export async function publishUserMessage( + streamBuffer: Pick | undefined, + taskId: string, + message: UIMessage, +): Promise { + if (!streamBuffer) return; + try { + await streamBuffer.publishRawChunk(taskId, buildUserMessageChunk(message)); + } catch { + // Best-effort live mirror. Never fail the POST because a viewer hint failed. + } +} diff --git a/apps/mesh/src/web/components/chat/store/thread-connection.ts b/apps/mesh/src/web/components/chat/store/thread-connection.ts index 7508dba0cf..edde440571 100644 --- a/apps/mesh/src/web/components/chat/store/thread-connection.ts +++ b/apps/mesh/src/web/components/chat/store/thread-connection.ts @@ -952,6 +952,15 @@ export class ThreadConnection { this.chunkBuffer.push(chunk); return; } + // A user prompt mirrored onto the run stream (see server + // `user-message-stream.ts`). Intercept BEFORE the assistant fold — the AI + // SDK reassembler hardcodes `role:"assistant"` — and upsert it by id so the + // author's optimistic copy dedupes while other viewers get it live. + if ((chunk as { type?: unknown }).type === "data-user-message") { + const message = (chunk as { data?: UIMessage }).data; + if (message) this.applyLocalMessage(message); + return; + } const isRunStatusControl = isRunStatusControlChunk(chunk); const runStatusStage = parseRunStatusStageChunk(chunk); if (this.waitingForNewRun) {