Skip to content

Commit 596ac16

Browse files
juliusmarmingejuliusmarminge
committed
Memoize derived thread reads
- Extract shared thread derivation - Add existence selector for route checks - Cover memoization behavior with tests
1 parent a3dadf3 commit 596ac16

5 files changed

Lines changed: 309 additions & 217 deletions

File tree

‎apps/web/src/routes/_chat.$environmentId.$threadId.tsx‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import {
1717
stripDiffSearchParams,
1818
} from "../diffRouteSearch";
1919
import { useMediaQuery } from "../hooks/useMediaQuery";
20-
import { selectEnvironmentState, selectThreadByRef, useStore } from "../store";
20+
import { selectEnvironmentState, selectThreadExistsByRef, useStore } from "../store";
2121
import { createThreadSelectorByRef } from "../storeSelectors";
2222
import { resolveThreadRouteRef, buildThreadRouteParams } from "../threadRoutes";
2323
import { Sheet, SheetPopup } from "../components/ui/sheet";
@@ -172,7 +172,7 @@ function ChatThreadRouteView() {
172172
(store) => selectEnvironmentState(store, threadRef?.environmentId ?? null).bootstrapComplete,
173173
);
174174
const serverThread = useStore(useMemo(() => createThreadSelectorByRef(threadRef), [threadRef]));
175-
const threadExists = useStore((store) => selectThreadByRef(store, threadRef) !== undefined);
175+
const threadExists = useStore((store) => selectThreadExistsByRef(store, threadRef));
176176
const environmentHasServerThreads = useStore(
177177
(store) => selectEnvironmentState(store, threadRef?.environmentId ?? null).threadIds.length > 0,
178178
);
@@ -208,6 +208,7 @@ function ChatThreadRouteView() {
208208
if (!threadRef) {
209209
return;
210210
}
211+
setHasOpenedDiff(true);
211212
void navigate({
212213
to: "/$environmentId/$threadId",
213214
params: buildThreadRouteParams(threadRef),
@@ -218,12 +219,6 @@ function ChatThreadRouteView() {
218219
});
219220
}, [navigate, threadRef]);
220221

221-
useEffect(() => {
222-
if (diffOpen) {
223-
setHasOpenedDiff(true);
224-
}
225-
}, [diffOpen]);
226-
227222
useEffect(() => {
228223
if (!threadRef || !bootstrapComplete) {
229224
return;

‎apps/web/src/store.test.ts‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ import {
1818
applyOrchestrationEvents,
1919
selectEnvironmentState,
2020
selectProjectsAcrossEnvironments,
21+
selectThreadByRef,
22+
selectThreadExistsByRef,
2123
setThreadBranch,
2224
selectThreadsAcrossEnvironments,
2325
syncServerReadModel,
@@ -245,6 +247,128 @@ function makeEvent<T extends OrchestrationEvent["type"]>(
245247
} as Extract<OrchestrationEvent, { type: T }>;
246248
}
247249

250+
describe("thread selection memoization", () => {
251+
it("returns stable thread references for repeated reads of the same state", () => {
252+
const thread = makeThread({
253+
messages: [
254+
{
255+
id: MessageId.make("message-1"),
256+
role: "user",
257+
text: "hello",
258+
createdAt: "2026-02-13T00:01:00.000Z",
259+
streaming: false,
260+
},
261+
],
262+
activities: [
263+
{
264+
id: EventId.make("activity-1"),
265+
tone: "info",
266+
kind: "step",
267+
summary: "working",
268+
payload: {},
269+
turnId: TurnId.make("turn-1"),
270+
createdAt: "2026-02-13T00:01:30.000Z",
271+
},
272+
],
273+
proposedPlans: [
274+
{
275+
id: "plan-1",
276+
turnId: null,
277+
planMarkdown: "plan",
278+
implementedAt: null,
279+
implementationThreadId: null,
280+
createdAt: "2026-02-13T00:02:00.000Z",
281+
updatedAt: "2026-02-13T00:02:00.000Z",
282+
},
283+
],
284+
turnDiffSummaries: [
285+
{
286+
turnId: TurnId.make("turn-1"),
287+
completedAt: "2026-02-13T00:03:00.000Z",
288+
files: [],
289+
},
290+
],
291+
});
292+
const state = makeState(thread);
293+
const ref = scopeThreadRef(thread.environmentId, thread.id);
294+
295+
const first = selectThreadByRef(state, ref);
296+
const second = selectThreadByRef(state, ref);
297+
298+
expect(first).toBeDefined();
299+
expect(second).toBe(first);
300+
expect(second?.messages).toBe(first?.messages);
301+
expect(second?.activities).toBe(first?.activities);
302+
expect(second?.proposedPlans).toBe(first?.proposedPlans);
303+
expect(second?.turnDiffSummaries).toBe(first?.turnDiffSummaries);
304+
});
305+
306+
it("reuses the derived thread when the app state wrapper changes but thread data does not", () => {
307+
const thread = makeThread({
308+
messages: [
309+
{
310+
id: MessageId.make("message-1"),
311+
role: "assistant",
312+
text: "done",
313+
createdAt: "2026-02-13T00:01:00.000Z",
314+
streaming: false,
315+
},
316+
],
317+
});
318+
const state = makeState(thread);
319+
const ref = scopeThreadRef(thread.environmentId, thread.id);
320+
const wrappedState: AppState = {
321+
...state,
322+
environmentStateById: { ...state.environmentStateById },
323+
};
324+
325+
const first = selectThreadByRef(state, ref);
326+
const second = selectThreadByRef(wrappedState, ref);
327+
328+
expect(second).toBe(first);
329+
});
330+
331+
it("updates the derived thread when the underlying thread data changes", () => {
332+
const thread = makeThread();
333+
const ref = scopeThreadRef(thread.environmentId, thread.id);
334+
const firstState = makeState(thread);
335+
const secondState = makeState({
336+
...thread,
337+
messages: [
338+
{
339+
id: MessageId.make("message-2"),
340+
role: "user",
341+
text: "new",
342+
createdAt: "2026-02-13T00:04:00.000Z",
343+
streaming: false,
344+
},
345+
],
346+
});
347+
348+
const first = selectThreadByRef(firstState, ref);
349+
const second = selectThreadByRef(secondState, ref);
350+
351+
expect(second).not.toBe(first);
352+
expect(second?.messages).toHaveLength(1);
353+
expect(second?.messages[0]?.text).toBe("new");
354+
});
355+
356+
it("checks thread existence without materializing the full thread", () => {
357+
const thread = makeThread();
358+
const state = makeState(thread);
359+
const ref = scopeThreadRef(thread.environmentId, thread.id);
360+
361+
expect(selectThreadExistsByRef(state, ref)).toBe(true);
362+
expect(
363+
selectThreadExistsByRef(
364+
state,
365+
scopeThreadRef(thread.environmentId, ThreadId.make("missing")),
366+
),
367+
).toBe(false);
368+
expect(selectThreadExistsByRef(state, null)).toBe(false);
369+
});
370+
});
371+
248372
function makeReadModelThread(overrides: Partial<OrchestrationReadModel["threads"][number]>) {
249373
return {
250374
id: ThreadId.make("thread-1"),

‎apps/web/src/store.ts‎

Lines changed: 14 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import {
3838
} from "./types";
3939
import { resolveEnvironmentHttpUrl } from "./environments/runtime";
4040
import { sanitizeThreadErrorMessage } from "./rpc/transportError";
41+
import { getThreadFromEnvironmentState } from "./threadDerivation";
4142

4243
export interface EnvironmentState {
4344
projectIds: ProjectId[];
@@ -94,19 +95,6 @@ const MAX_THREAD_CHECKPOINTS = 500;
9495
const MAX_THREAD_PROPOSED_PLANS = 200;
9596
const MAX_THREAD_ACTIVITIES = 500;
9697
const EMPTY_THREAD_IDS: ThreadId[] = [];
97-
const EMPTY_MESSAGE_IDS: MessageId[] = [];
98-
const EMPTY_ACTIVITY_IDS: string[] = [];
99-
const EMPTY_PROPOSED_PLAN_IDS: string[] = [];
100-
const EMPTY_TURN_IDS: TurnId[] = [];
101-
const EMPTY_MESSAGES: ChatMessage[] = [];
102-
const EMPTY_ACTIVITIES: OrchestrationThreadActivity[] = [];
103-
const EMPTY_PROPOSED_PLANS: ProposedPlan[] = [];
104-
const EMPTY_TURN_DIFF_SUMMARIES: TurnDiffSummary[] = [];
105-
const EMPTY_MESSAGE_MAP: Record<MessageId, ChatMessage> = {};
106-
const EMPTY_ACTIVITY_MAP: Record<string, OrchestrationThreadActivity> = {};
107-
const EMPTY_PROPOSED_PLAN_MAP: Record<string, ProposedPlan> = {};
108-
const EMPTY_TURN_DIFF_MAP: Record<TurnId, TurnDiffSummary> = {};
109-
const EMPTY_THREAD_TURN_STATE: ThreadTurnState = Object.freeze({ latestTurn: null });
11098

11199
function arraysEqual<T>(left: readonly T[], right: readonly T[]): boolean {
112100
return left.length === right.length && left.every((value, index) => value === right[index]);
@@ -403,78 +391,6 @@ function buildTurnDiffSlice(thread: Thread): {
403391
};
404392
}
405393

406-
function selectThreadMessages(state: EnvironmentState, threadId: ThreadId): ChatMessage[] {
407-
const ids = state.messageIdsByThreadId[threadId] ?? EMPTY_MESSAGE_IDS;
408-
const byId = state.messageByThreadId[threadId] ?? EMPTY_MESSAGE_MAP;
409-
if (ids.length === 0) {
410-
return EMPTY_MESSAGES;
411-
}
412-
return ids.flatMap((id) => {
413-
const message = byId[id];
414-
return message ? [message] : [];
415-
});
416-
}
417-
418-
function selectThreadActivities(
419-
state: EnvironmentState,
420-
threadId: ThreadId,
421-
): OrchestrationThreadActivity[] {
422-
const ids = state.activityIdsByThreadId[threadId] ?? EMPTY_ACTIVITY_IDS;
423-
const byId = state.activityByThreadId[threadId] ?? EMPTY_ACTIVITY_MAP;
424-
if (ids.length === 0) {
425-
return EMPTY_ACTIVITIES;
426-
}
427-
return ids.flatMap((id) => {
428-
const activity = byId[id];
429-
return activity ? [activity] : [];
430-
});
431-
}
432-
433-
function selectThreadProposedPlans(state: EnvironmentState, threadId: ThreadId): ProposedPlan[] {
434-
const ids = state.proposedPlanIdsByThreadId[threadId] ?? EMPTY_PROPOSED_PLAN_IDS;
435-
const byId = state.proposedPlanByThreadId[threadId] ?? EMPTY_PROPOSED_PLAN_MAP;
436-
if (ids.length === 0) {
437-
return EMPTY_PROPOSED_PLANS;
438-
}
439-
return ids.flatMap((id) => {
440-
const plan = byId[id];
441-
return plan ? [plan] : [];
442-
});
443-
}
444-
445-
function selectThreadTurnDiffSummaries(
446-
state: EnvironmentState,
447-
threadId: ThreadId,
448-
): TurnDiffSummary[] {
449-
const ids = state.turnDiffIdsByThreadId[threadId] ?? EMPTY_TURN_IDS;
450-
const byId = state.turnDiffSummaryByThreadId[threadId] ?? EMPTY_TURN_DIFF_MAP;
451-
if (ids.length === 0) {
452-
return EMPTY_TURN_DIFF_SUMMARIES;
453-
}
454-
return ids.flatMap((id) => {
455-
const summary = byId[id];
456-
return summary ? [summary] : [];
457-
});
458-
}
459-
460-
function getThread(state: EnvironmentState, threadId: ThreadId): Thread | undefined {
461-
const shell = state.threadShellById[threadId];
462-
if (!shell) {
463-
return undefined;
464-
}
465-
const turnState = state.threadTurnStateById[threadId] ?? EMPTY_THREAD_TURN_STATE;
466-
return {
467-
...shell,
468-
session: state.threadSessionById[threadId] ?? null,
469-
latestTurn: turnState.latestTurn,
470-
pendingSourceProposedPlan: turnState.pendingSourceProposedPlan,
471-
messages: selectThreadMessages(state, threadId),
472-
activities: selectThreadActivities(state, threadId),
473-
proposedPlans: selectThreadProposedPlans(state, threadId),
474-
turnDiffSummaries: selectThreadTurnDiffSummaries(state, threadId),
475-
};
476-
}
477-
478394
function getProjects(state: EnvironmentState): Project[] {
479395
return state.projectIds.flatMap((projectId) => {
480396
const project = state.projectById[projectId];
@@ -484,7 +400,7 @@ function getProjects(state: EnvironmentState): Project[] {
484400

485401
function getThreads(state: EnvironmentState): Thread[] {
486402
return state.threadIds.flatMap((threadId) => {
487-
const thread = getThread(state, threadId);
403+
const thread = getThreadFromEnvironmentState(state, threadId);
488404
return thread ? [thread] : [];
489405
});
490406
}
@@ -896,7 +812,7 @@ function updateThreadState(
896812
threadId: ThreadId,
897813
updater: (thread: Thread) => Thread,
898814
): EnvironmentState {
899-
const currentThread = getThread(state, threadId);
815+
const currentThread = getThreadFromEnvironmentState(state, threadId);
900816
if (!currentThread) {
901817
return state;
902818
}
@@ -1163,7 +1079,7 @@ function applyEnvironmentOrchestrationEvent(
11631079
}
11641080

11651081
case "thread.created": {
1166-
const previousThread = getThread(state, event.payload.threadId);
1082+
const previousThread = getThreadFromEnvironmentState(state, event.payload.threadId);
11671083
const nextThread = mapThread(
11681084
{
11691085
id: event.payload.threadId,
@@ -1669,10 +1585,19 @@ export function selectThreadByRef(
16691585
ref: ScopedThreadRef | null | undefined,
16701586
): Thread | undefined {
16711587
return ref
1672-
? getThread(selectEnvironmentState(state, ref.environmentId), ref.threadId)
1588+
? getThreadFromEnvironmentState(selectEnvironmentState(state, ref.environmentId), ref.threadId)
16731589
: undefined;
16741590
}
16751591

1592+
export function selectThreadExistsByRef(
1593+
state: AppState,
1594+
ref: ScopedThreadRef | null | undefined,
1595+
): boolean {
1596+
return ref
1597+
? selectEnvironmentState(state, ref.environmentId).threadShellById[ref.threadId] !== undefined
1598+
: false;
1599+
}
1600+
16761601
export function selectSidebarThreadSummaryByRef(
16771602
state: AppState,
16781603
ref: ScopedThreadRef | null | undefined,

0 commit comments

Comments
 (0)