From 5701a9da9c47d8970a85fcba3dc8b82bf6ee177e Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 16 Jun 2026 15:29:50 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20fix:=20Codex=20review=20on=20con?= =?UTF-8?q?text=20projection=20(G1=20guard,=20IDOR,=20recount,=20summary)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Guard `currentActive` against a stale window: a model/window switch on the current branch left the live snapshot outranking the projection (G1 didn't fire). Now defers to the projection unless streaming or the window matches. - Scope branch lookups to the authenticated user (`getMessages` filter + injected `userId`) — was loading any conversation by id (IDOR). - Recount messages with no stored `tokenCount` via the tokenizer instead of charging 0, so snapshot-less/imported histories don't under-report. - Fall back (null) for already-summarized branches rather than projecting from the full raw parent chain (the next call would send summary + tail); the client's summary-baseline-aware estimate handles them until a follow-up replays the summary boundary. --- .../ContextProjectionController.js | 2 +- client/src/hooks/Chat/useTokenUsage.ts | 12 +++++-- packages/api/src/endpoints/projection.ts | 36 ++++++++++++++----- 3 files changed, 38 insertions(+), 12 deletions(-) diff --git a/api/server/controllers/ContextProjectionController.js b/api/server/controllers/ContextProjectionController.js index df5c79d3bf..88bbbed366 100644 --- a/api/server/controllers/ContextProjectionController.js +++ b/api/server/controllers/ContextProjectionController.js @@ -18,7 +18,7 @@ async function contextProjectionController(req, res) { return; } const projection = await resolveContextProjection( - { getMessages: db.getMessages, getAgent: db.getAgent }, + { userId: req.user?.id, getMessages: db.getMessages, getAgent: db.getAgent }, params, ); res.json(projection ?? null); diff --git a/client/src/hooks/Chat/useTokenUsage.ts b/client/src/hooks/Chat/useTokenUsage.ts index 3e0749e2b3..f586c3db22 100644 --- a/client/src/hooks/Chat/useTokenUsage.ts +++ b/client/src/hooks/Chat/useTokenUsage.ts @@ -237,7 +237,13 @@ export default function useTokenUsage({ * one branch's breakdown onto its siblings. */ const currentActive = snapshot != null && - (isSubmitting || (snapshot.anchorMessageId != null && branchTotals.containsAnchor)); + (isSubmitting || + (snapshot.anchorMessageId != null && + branchTotals.containsAnchor && + /** G1: once streaming ends, a model/window switch leaves the live + * snapshot's baked window stale — defer to the projection instead of + * showing the old window/prune boundary on the current branch. */ + (resolvedMax == null || snapshot.breakdown.maxContextTokens === resolvedMax))); /** Precedence: live/active snapshot → fresh persisted branch snapshot → * server projection (covers G1 stale-window + G2 snapshot-less branches) → @@ -258,8 +264,7 @@ export default function useTokenUsage({ if (effective != null) { const breakdown = effective.breakdown; const maxTokens = effective.contextBudget ?? breakdown.maxContextTokens; - const instructionTokens = - effective.effectiveInstructionTokens ?? breakdown.instructionTokens; + const instructionTokens = effective.effectiveInstructionTokens ?? breakdown.instructionTokens; const baseUsed = effective.remainingContextTokens != null ? maxTokens - effective.remainingContextTokens @@ -324,5 +329,6 @@ export default function useTokenUsage({ branchSnapshot, branchSnapshotFresh, projection, + resolvedMax, ]); } diff --git a/packages/api/src/endpoints/projection.ts b/packages/api/src/endpoints/projection.ts index a126545eb6..2362e219e0 100644 --- a/packages/api/src/endpoints/projection.ts +++ b/packages/api/src/endpoints/projection.ts @@ -7,6 +7,7 @@ interface ProjectionMessage { messageId: string; parentMessageId?: string | null; tokenCount?: number; + summaryTokenCount?: number; isCreatedByUser?: boolean; text?: string; } @@ -19,8 +20,10 @@ interface ProjectionAgent { } export interface ContextProjectionDeps { + /** Authenticated requester — branch lookups are scoped to this user. */ + userId?: string; getMessages: ( - filter: { conversationId: string }, + filter: { conversationId: string; user?: string }, select?: string, ) => Promise; getAgent: (filter: { id: string }) => Promise; @@ -87,14 +90,23 @@ export async function resolveContextProjection( params: TContextProjectionRequest, ): Promise { const stored = await deps.getMessages( - { conversationId: params.conversationId }, - 'messageId parentMessageId tokenCount isCreatedByUser text', + { conversationId: params.conversationId, user: deps.userId }, + 'messageId parentMessageId tokenCount summaryTokenCount isCreatedByUser text', ); const branch = resolveBranch(stored, params.messageId); if (branch.length === 0) { return null; } + /** A summarized/compacted branch's next call sends the saved summary + the + * post-summary tail, NOT this raw parent chain — projecting from the full + * history would prune/count the wrong context and omit the summary. Until the + * follow-up replays the summary boundary, fall back (null) so the client's + * summary-baseline-aware estimate handles these branches. */ + if (branch.some((message) => (message.summaryTokenCount ?? 0) > 0)) { + return null; + } + let instructions: string | undefined; let providerValue: string | undefined = params.endpoint; let model = params.model; @@ -112,18 +124,26 @@ export async function resolveContextProjection( return null; } + const encoding = (model ?? '').toLowerCase().includes('claude') ? 'claude' : 'o200k_base'; + const tokenCounter = await createTokenCounter(encoding); + const messages: BaseMessage[] = []; const indexTokenCountMap: Record = {}; for (let i = 0; i < branch.length; i++) { const message = branch[i]; const text = message.text ?? ''; - messages.push(message.isCreatedByUser === true ? new HumanMessage(text) : new AIMessage(text)); - indexTokenCountMap[String(i)] = message.tokenCount ?? 0; + const lcMessage = + message.isCreatedByUser === true ? new HumanMessage(text) : new AIMessage(text); + messages.push(lcMessage); + /** Recount messages with no stored count (imported / pre-feature) rather + * than charging 0 — a real 0 and "unknown" must not collapse, or the + * snapshot-less histories this endpoint targets would under-report. */ + indexTokenCountMap[String(i)] = + message.tokenCount != null && message.tokenCount > 0 + ? message.tokenCount + : tokenCounter(lcMessage); } - const encoding = (model ?? '').toLowerCase().includes('claude') ? 'claude' : 'o200k_base'; - const tokenCounter = await createTokenCounter(encoding); - return projectAgentContextUsage({ agent: { agentId: params.agentId ?? 'projection',