Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 60 additions & 0 deletions apps/mesh/src/api/routes/decopilot/project-chunks.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});
13 changes: 11 additions & 2 deletions apps/mesh/src/api/routes/decopilot/project-chunks.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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) {
Expand All @@ -207,7 +216,7 @@ export async function projectChunks(
}).pipeThrough(
new TransformStream<UIMessageChunk, UIMessageChunk>({
transform(chunk, controller) {
if (!isRunStatusControlChunk(chunk)) controller.enqueue(chunk);
if (!isTransientControlChunk(chunk)) controller.enqueue(chunk);
},
}),
)
Expand Down
5 changes: 5 additions & 0 deletions apps/mesh/src/api/routes/decopilot/routes.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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, {
Expand Down
74 changes: 74 additions & 0 deletions apps/mesh/src/api/routes/decopilot/user-message-stream.test.ts
Original file line number Diff line number Diff line change
@@ -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();
});
});
48 changes: 48 additions & 0 deletions apps/mesh/src/api/routes/decopilot/user-message-stream.ts
Original file line number Diff line number Diff line change
@@ -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<StreamBuffer, "publishRawChunk"> | undefined,
taskId: string,
message: UIMessage,
): Promise<void> {
if (!streamBuffer) return;
try {
await streamBuffer.publishRawChunk(taskId, buildUserMessageChunk(message));
} catch {
// Best-effort live mirror. Never fail the POST because a viewer hint failed.
}
}
9 changes: 9 additions & 0 deletions apps/mesh/src/web/components/chat/store/thread-connection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Loading