opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: use cursor-based query for chat stream notifications (#22510)

Browse Source 
## Problem

The pubsub notification handler in `chatd` re-fetched **all** messages
from the DB on every new message notification, then filtered in Go with
`msg.ID > lastMessageID`. This grows linearly with conversation length —
every new message triggers a full table scan of that chat's history.

The `AfterMessageID` field in the pubsub notification payload was
clearly designed for cursor-based fetching, but no matching query
existed.

## Fix

- Add `GetChatMessagesByChatIDAfter` SQL query with `WHERE id >
@after_id`, so the database does the filtering instead of Go.
- Use it in the pubsub notification handler in `chatd.go`, passing
`lastMessageID` as the cursor.
- Implement the dbauthz wrapper (was a `panic("not implemented")` stub
from codegen) with the same read-check-on-parent-chat pattern as
adjacent methods.
- Add dbauthz test coverage for the new method.

**Not changed:** The initial snapshot in `Subscribe()` still loads all
messages — that's correct, since a newly-connecting client needs the
full conversation state. The waste was only in the ongoing notification
path.
This commit is contained in:
Kyle Carberry
2026-03-02 16:31:04 -05:00
committed by GitHub
parent e3c5d734ba
commit 5eebd3829f
14 changed files with 216 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
site/src/api/api.ts
+8 -1
View File
@@ -140,9 +140,16 @@ export const watchWorkspace = (
export const watchChat = (
chatId: string,
afterMessageId?: number,
): OneWayWebSocket<TypesGen.ServerSentEvent> => {
const params = new URLSearchParams();
if (afterMessageId !== undefined && afterMessageId > 0) {
params.set("after_id", afterMessageId.toString());
}
const query = params.toString();
const route = `/api/experimental/chats/${chatId}/stream${query ? `?${query}` : ""}`;
return new OneWayWebSocket({
apiRoute: `/api/experimental/chats/${chatId}/stream`,
apiRoute: route,
});
};
site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx
+26 -26
View File
@@ -202,7 +202,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
act(() => {
@@ -283,7 +283,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
act(() => {
@@ -358,7 +358,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
act(() => {
@@ -460,7 +460,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
const streamBaseline = streamRenderCount;
@@ -526,7 +526,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
act(() => {
@@ -601,7 +601,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
act(() => {
@@ -696,7 +696,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
expect(result.current.queuedMessages.map((message) => message.id)).toEqual([
queuedMessage.id,
@@ -781,7 +781,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
act(() => {
@@ -852,7 +852,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID1);
expect(watchChat).toHaveBeenCalledWith(chatID1, 1);
});
act(() => {
@@ -888,7 +888,7 @@ describe("useChatStore", () => {
});
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID2);
expect(watchChat).toHaveBeenCalledWith(chatID2, 10);
});
// The old WebSocket was closed during effect cleanup.
@@ -935,7 +935,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
act(() => {
@@ -991,7 +991,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
// Build up stream state so we can observe whether it gets cleared.
@@ -1093,7 +1093,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, 1);
});
// Build up stream state first.
@@ -1193,7 +1193,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID1);
expect(watchChat).toHaveBeenCalledWith(chatID1, 1);
});
act(() => {
@@ -1229,7 +1229,7 @@ describe("useChatStore", () => {
});
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID2);
expect(watchChat).toHaveBeenCalledWith(chatID2, 10);
});
expect(result.current.streamState).toBeNull();
@@ -1284,7 +1284,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID1);
expect(watchChat).toHaveBeenCalledWith(chatID1, 1);
});
// Verify queued messages from chat-1 are present.
@@ -1310,7 +1310,7 @@ describe("useChatStore", () => {
// After the switch, queued messages from chat-1 should NOT be
// visible — the store resets them on chatID change.
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID2);
expect(watchChat).toHaveBeenCalledWith(chatID2, undefined);
});
expect(result.current.queuedMessages).toEqual([]);
});
@@ -1352,7 +1352,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
// Emit a batch with message_parts followed by a status change
@@ -1424,7 +1424,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
act(() => {
@@ -1483,7 +1483,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
act(() => {
@@ -1536,7 +1536,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
act(() => {
@@ -1598,7 +1598,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
// Set retry state first.
@@ -1676,7 +1676,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
act(() => {
@@ -1734,7 +1734,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
act(() => {
@@ -1783,7 +1783,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
// Set an error via an error stream event first.
@@ -1847,7 +1847,7 @@ describe("useChatStore", () => {
);
await waitFor(() => {
expect(watchChat).toHaveBeenCalledWith(chatID);
expect(watchChat).toHaveBeenCalledWith(chatID, undefined);
});
// Transition to running — should call clearChatErrorReason.
site/src/pages/AgentsPage/AgentDetail/ChatContext.ts
+14 -1
View File
@@ -435,6 +435,17 @@ export const useChatStore = (
const store = storeRef.current;
// Compute the last REST-fetched message ID so the stream can
// skip messages the client already has. We use a ref so the
// socket effect can read the latest value without including
// chatMessages in its dependency array (which would cause
// unnecessary reconnections).
const lastMessageIdRef = useRef<number | undefined>(undefined);
lastMessageIdRef.current =
chatMessages && chatMessages.length > 0
? chatMessages[chatMessages.length - 1].id
: undefined;
const updateSidebarChat = useCallback(
(updater: (chat: TypesGen.Chat) => TypesGen.Chat) => {
if (!chatID) {
@@ -550,7 +561,9 @@ export const useChatStore = (
return;
}
const socket = watchChat(chatID);
// Pass the last REST-fetched message ID so the stream
// only sends newer messages.
const socket = watchChat(chatID, lastMessageIdRef.current);
const handleMessage = (
payload: OneWayMessageEvent<TypesGen.ServerSentEvent>,
) => {