Thanks to visit codestin.com
Credit goes to github.com

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
59 changes: 59 additions & 0 deletions site/src/pages/AgentsPage/AgentChatPage.stories.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -2398,6 +2398,65 @@ export const DurableUpdateFansOutToOlderPage: Story = {
},
};

/**
* The streamed thinking block is asserted via its disclosure header
* because smoothed body text never reveals in the test iframe
* (requestAnimationFrame is suspended).
*/
export const StreamedPartWhileStatusWaiting: Story = {
parameters: {
queries: buildQueries(
{
id: CHAT_ID,
...baseChatFields,
title: "Stale status stream",
status: "waiting",
},
{
messages: [
{
id: 1,
chat_id: CHAT_ID,
created_at: "2026-02-18T00:05:00.000Z",
role: "user",
content: [{ type: "text", text: "Start the next turn" }],
},
],
queued_messages: [],
has_more: false,
},
{ diffUrl: undefined },
),
webSocket: {
"/chats/": [
{
event: "message",
data: JSON.stringify([
{
type: "message_part",
chat_id: CHAT_ID,
message_part: {
part: {
type: "reasoning",
text: "Streaming while the chat still reads waiting",
},
},
},
] satisfies TypesGen.ChatStreamEvent[]),
},
],
},
},
play: async ({ canvasElement }) => {
const canvas = within(canvasElement);
await canvas.findByText("Start the next turn");
// The disclosure header is the only "Thinking" text in the DOM
// at this point: the generic indicator is suppressed at status
// "waiting".
expect(await canvas.findByText("Thinking")).toBeVisible();
},
};

/**
* Live agent turn with streaming reasoning and a back-to-back flurry of
* in-progress file tool calls. The persisted history establishes context
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1409,7 +1409,7 @@ describe("useChatStore", () => {
});
});

it("ignores message_part updates while chat is waiting", async () => {
it("applies message_part updates while a REST-hydrated chat status reads waiting", async () => {
immediateAnimationFrame();

const chatID = "chat-1";
Expand All @@ -1427,7 +1427,10 @@ describe("useChatStore", () => {
const { store } = useChatStore({
chatID,
chatMessages: [existingMessage],
chatRecord: buildChat(chatID),
// REST hydrates the store with a waiting status, but the
// stream has not delivered any status event, so the
// server's view may already be ahead.
chatRecord: { ...buildChat(chatID), status: "waiting" },
chatMessagesData: {
messages: [existingMessage],
queued_messages: [],
Expand Down Expand Up @@ -1456,57 +1459,171 @@ describe("useChatStore", () => {
role: "assistant",
part: {
type: "text",
text: "first",
text: "live output",
},
},
});
});

await waitFor(() => {
expect(result.current.streamState?.blocks).toEqual([
{ type: "response", text: "first" },
{ type: "response", text: "live output" },
]);
});
});

it("drops message_part updates after the stream reports waiting", async () => {
immediateAnimationFrame();

const chatID = "chat-1";
const existingMessage = buildMessage(chatID, 1, "user", "hello");
const mockSocket = createMockSocket();
mockWatchChatReturn(mockSocket);

const queryClient = createTestQueryClient();
const wrapper = createWrapper(queryClient);
const setChatErrorReason = vi.fn();
const clearChatErrorReason = vi.fn();

const { result } = renderHook(
() => {
const { store } = useChatStore({
chatID,
chatMessages: [existingMessage],
chatRecord: { ...buildChat(chatID), status: "waiting" },
chatMessagesData: {
messages: [existingMessage],
queued_messages: [],
has_more: false,
},
chatQueuedMessages: [],
setChatErrorReason,
clearChatErrorReason,
});
return {
streamState: useChatSelector(store, selectStreamState),
chatStatus: useChatSelector(store, selectChatStatus),
store,
};
},
{ wrapper },
);

await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});

act(() => {
mockSocket.emitData({
type: "status",
chat_id: chatID,
status: { status: "waiting" },
status: { status: "running" },
});
});

act(() => {
mockSocket.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: { type: "text", text: "first" },
},
});
});

await waitFor(() => {
// Stream state is preserved after status=waiting (the
// durable message event handles cleanup via
// needsStreamReset). Only new message_parts should be
// blocked by the shouldApplyMessagePart gate.
expect(result.current.streamState).not.toBeNull();
expect(result.current.streamState?.blocks).toEqual([
{ type: "response", text: "first" },
]);
});

// The stream reports waiting: the turn is over server-side.
// A part arriving now comes from the closed episode draining
// (the server keeps closed episodes subscribed for replay for
// up to 15s) and must not repopulate the stream.
act(() => {
mockSocket.emitData({
type: "status",
chat_id: chatID,
status: { status: "waiting" },
});
});

await waitFor(() => {
expect(result.current.chatStatus).toBe("waiting");
});

act(() => {
mockSocket.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: {
type: "text",
text: "late",
},
part: { type: "text", text: "late drain" },
},
});
});

// Wait past the coalesced flush window so the drop is
// observable rather than "not yet flushed".
await act(async () => {
await new Promise((resolve) => setTimeout(resolve, 10));
});

expect(result.current.streamState?.blocks).toEqual([
{ type: "response", text: "first" },
]);

// An optimistic send status must not reopen the window; drain
// parts from the closed episode are still dropped.
act(() => {
result.current.store.setChatStatus("running");
});

act(() => {
mockSocket.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: { type: "text", text: "drain after send" },
},
});
});

await act(async () => {
await new Promise((resolve) => setTimeout(resolve, 10));
});

expect(result.current.streamState?.blocks).toEqual([
{ type: "response", text: "first" },
]);

// The stream reporting running reopens the window: the next
// turn's parts must flow again.
act(() => {
mockSocket.emitData({
type: "status",
chat_id: chatID,
status: { status: "running" },
});
});

act(() => {
mockSocket.emitData({
type: "message_part",
chat_id: chatID,
message_part: {
role: "assistant",
part: { type: "text", text: "next turn" },
},
});
});

await waitFor(() => {
// The late message_part should not be applied because
// shouldApplyMessagePart gates on waiting.
// Stream state still shows the original "first".
expect(result.current.streamState?.blocks).toEqual([
{ type: "response", text: "first" },
{ type: "response", text: "firstnext turn" },
]);
});
});
Expand Down Expand Up @@ -4384,7 +4501,7 @@ describe("thinking indicator event ordering", () => {
});
});

it("discards buffered parts when status transitions to pending", async () => {
it("discards buffered parts when status transitions to waiting", async () => {
vi.useFakeTimers({ shouldAdvanceTime: true });
immediateAnimationFrame();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -143,8 +143,8 @@ export const useChatStore = (
// source for chatStatus and the REST-fetched chatRecord.status
// must not overwrite it. Without this guard, a React Query
// refetch (e.g. on window focus) can regress chatStatus to a
// stale value like "waiting", causing shouldApplyMessagePart()
// to drop all incoming parts.
// stale value like "waiting", hiding live status from the user
// until the next server event arrives.
const wsStatusReceivedRef = useRef(false);
const [pendingStatusResync, setPendingStatusResync] = useState(false);
const pendingStatusResyncUpdatedAtRef = useRef<number | null>(null);
Expand Down Expand Up @@ -423,9 +423,13 @@ export const useChatStore = (
let historyResetPending = false;
const historyReplacementBuf: TypesGen.ChatMessage[] = [];

const shouldApplyMessagePart = (): boolean => {
return store.getSnapshot().chatStatus !== "waiting";
};
// Set when the stream reports "waiting", cleared by any other
// stream status. While set, parts are dropped: they are late
// leftovers from the finished turn. REST and optimistic
// statuses never set it, because they can lag a live turn.
let streamReportedWaiting = false;

const shouldKeepMessagePart = (): boolean => !streamReportedWaiting;

const schedulePartsFlush = () => {
if (partsFlushTimer !== null || partsBuf.length === 0) {
Expand All @@ -436,11 +440,11 @@ export const useChatStore = (
if (disposed || activeChatIDRef.current !== chatID) {
return;
}
const parts = partsBuf.splice(0);
if (parts.length === 0 || !shouldApplyMessagePart()) {
if (!shouldKeepMessagePart()) {
partsBuf.length = 0;
return;
}
store.applyMessageParts(parts);
store.applyMessageParts(partsBuf.splice(0));
}, 0);
};

Expand All @@ -455,11 +459,11 @@ export const useChatStore = (
clearTimeout(partsFlushTimer);
partsFlushTimer = null;
}
const parts = partsBuf.splice(0);
if (activeChatIDRef.current !== chatID || !shouldApplyMessagePart()) {
if (activeChatIDRef.current !== chatID || !shouldKeepMessagePart()) {
partsBuf.length = 0;
return;
}
store.applyMessageParts(parts);
store.applyMessageParts(partsBuf.splice(0));
};

// Discard buffered parts without applying them. Used when
Expand Down Expand Up @@ -521,7 +525,7 @@ export const useChatStore = (
continue;
}
commitHistoryReplacement();
if (!shouldApplyMessagePart()) {
if (!shouldKeepMessagePart()) {
continue;
}
const part = streamEvent.message_part?.part;
Expand Down Expand Up @@ -622,6 +626,7 @@ export const useChatStore = (
continue;
}

streamReportedWaiting = nextStatus === "waiting";
wsStatusReceivedRef.current = true;
store.clearRetryState();
store.applyServerChatStatus(nextStatus);
Expand All @@ -643,6 +648,7 @@ export const useChatStore = (
kind: "generic",
message: "Chat processing failed.",
};
streamReportedWaiting = false;
wsStatusReceivedRef.current = true;
store.applyServerChatStatus("error");
store.setStreamError(reason);
Expand Down Expand Up @@ -698,9 +704,8 @@ export const useChatStore = (
clearTimeout(partsFlushTimer);
partsFlushTimer = null;
}
const nextParts = partsBuf.splice(0);
if (shouldApplyMessagePart()) {
store.applyMessageParts(nextParts);
if (shouldKeepMessagePart()) {
store.applyMessageParts(partsBuf.splice(0));
}
}
}
Expand Down
Loading