salvacybersec/CodeNomad
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ui): restrict selection and xray to post-compaction

Browse Source
This commit is contained in:
Shantur Rathore
2026-03-03 14:09:48 +00:00
parent c766b5ab62
commit 80a02b68b9
3 changed files with 118 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
packages/ui/src/components/message-section.tsx
View File

@@ -103,6 +103,42 @@ export default function MessageSection(props: MessageSectionProps) {
let deleteMenuRef: HTMLDivElement | undefined
let deleteMenuButtonRef: HTMLButtonElement | undefined
// Deletion is only allowed for messages/tool parts that occur AFTER the most
// recent compaction. Compaction effectively resets the stored context; deleting
// earlier items would not reliably reflect what the model sees.
const messageIndexById = createMemo(() => {
const ids = messageIds()
const map = new Map<string, number>()
for (let i = 0; i < ids.length; i++) {
map.set(ids[i], i)
}
return map
})
const lastCompactionIndex = createMemo(() => {
// Depend on a single session revision signal (not every message/part read)
// to keep reactive overhead small.
sessionRevision()
return untrack(() => store().getLastCompactionMessageIndex(props.sessionId))
})
const deletableStartIndex = createMemo(() => {
const idx = lastCompactionIndex()
return idx === -1 ? 0 : idx + 1
})
const deletableMessageIds = createMemo(() => {
const ids = messageIds()
const start = deletableStartIndex()
return new Set(ids.slice(start))
})
const isMessageDeletable = (messageId: string): boolean => {
const idx = messageIndexById().get(messageId)
if (idx === undefined) return false
return idx >= deletableStartIndex()
}
// Build the message group for a segment.
// Tool calls belong to the same assistant turn (between user messages).
// Only assistant badges trigger group selection; user/tool badges are standalone.
@@ -132,12 +168,17 @@ export default function MessageSection(props: MessageSectionProps) {
}
const handleToggleTimelineSelection = (id: string) => {
setLastSelectionAnchorId(id)
const segments = timelineSegments()
const segmentIndex = segments.findIndex((s) => s.id === id)
if (segmentIndex === -1) return
const segment = segments[segmentIndex]
if (!isMessageDeletable(segment.messageId)) {
return
}
setLastSelectionAnchorId(id)
if (selectionMode() === "tools" && segment.type !== "tool") {
return
}
@@ -177,11 +218,16 @@ export default function MessageSection(props: MessageSectionProps) {
}
const handleLongPressTimelineSelection = (segment: TimelineSegment) => {
setLastSelectionAnchorId(segment.id)
const segments = timelineSegments()
const segmentIndex = segments.findIndex((s) => s.id === segment.id)
if (segmentIndex === -1) return
if (!isMessageDeletable(segment.messageId)) {
return
}
setLastSelectionAnchorId(segment.id)
if (selectionMode() === "tools" && segment.type !== "tool") {
return
}
@@ -227,8 +273,8 @@ export default function MessageSection(props: MessageSectionProps) {
const end = Math.max(anchorIndex, targetIndex)
const rangeSegments = selectionMode() === "tools"
? segments.slice(start, end + 1).filter((s) => s.type === "tool")
: segments.slice(start, end + 1)
? segments.slice(start, end + 1).filter((s) => s.type === "tool" && isMessageDeletable(s.messageId))
: segments.slice(start, end + 1).filter((s) => isMessageDeletable(s.messageId))
// Range selection replaces current selection so it can grow or shrink.
setSelectedTimelineIds(new Set(rangeSegments.map((segment) => segment.id)))
}
@@ -242,7 +288,11 @@ export default function MessageSection(props: MessageSectionProps) {
setSelectionMode(mode)
if (mode !== "tools") return
const segments = timelineSegments()
const toolIds = new Set(segments.filter((segment) => segment.type === "tool").map((segment) => segment.id))
const toolIds = new Set(
segments
.filter((segment) => segment.type === "tool" && isMessageDeletable(segment.messageId))
.map((segment) => segment.id),
)
setSelectedTimelineIds((prev) => {
if (prev.size === 0) return prev
const next = new Set([...prev].filter((id) => toolIds.has(id)))
@@ -345,7 +395,8 @@ export default function MessageSection(props: MessageSectionProps) {
const deleteMessageIds = createMemo(() => selectedForDeletion())
const deleteToolParts = createMemo(() => {
const messageIds = deleteMessageIds()
return selectedToolParts().filter((entry) => !messageIds.has(entry.messageId))
const allowed = deletableMessageIds()
return selectedToolParts().filter((entry) => allowed.has(entry.messageId) && !messageIds.has(entry.messageId))
})
const isDeleteMode = createMemo(() => deleteMessageIds().size > 0 || deleteToolParts().length > 0)
const selectedDeleteCount = createMemo(() => deleteMessageIds().size + deleteToolParts().length)
@@ -410,6 +461,7 @@ export default function MessageSection(props: MessageSectionProps) {
const setMessageSelectedForDeletion = (messageId: string, selected: boolean) => {
if (!messageId) return
if (!isMessageDeletable(messageId)) return
setSelectedForDeletion((prev) => {
const next = new Set(prev)
if (selected) {
@@ -440,18 +492,20 @@ export default function MessageSection(props: MessageSectionProps) {
const affectedMessageIds = new Set<string>()
for (const segId of timelineIds) {
const segment = segmentById.get(segId)
if (segment && segment.type !== "tool") affectedMessageIds.add(segment.messageId)
if (segment && segment.type !== "tool" && isMessageDeletable(segment.messageId)) {
affectedMessageIds.add(segment.messageId)
}
}
setSelectedForDeletion(affectedMessageIds)
})
const selectAllForDeletion = () => {
const allMessageIds = messageIds()
const allMessageIds = [...deletableMessageIds()]
setSelectedForDeletion(new Set<string>(allMessageIds))
// Also select all timeline segments — tool visibility is handled by
// isSelectionActive() in isHidden(), no expand/collapse needed.
const segments = timelineSegments()
setSelectedTimelineIds(new Set(segments.map((s) => s.id)))
setSelectedTimelineIds(new Set(segments.filter((s) => isMessageDeletable(s.messageId)).map((s) => s.id)))
}
const deleteSelectedMessages = async () => {
@@ -459,11 +513,13 @@ export default function MessageSection(props: MessageSectionProps) {
const toolParts = deleteToolParts()
if (selected.size === 0 && toolParts.length === 0) return
const allowed = deletableMessageIds()
const idsInSessionOrder = messageIds()
const toDelete: string[] = []
for (let idx = idsInSessionOrder.length - 1; idx >= 0; idx -= 1) {
const id = idsInSessionOrder[idx]
if (selected.has(id)) {
if (allowed.has(id) && selected.has(id)) {
toDelete.push(id)
}
}
@@ -473,6 +529,7 @@ export default function MessageSection(props: MessageSectionProps) {
await deleteMessage(props.instanceId, props.sessionId, messageId)
}
for (const { messageId, partId } of toolParts) {
if (!allowed.has(messageId)) continue
await deleteMessagePart(props.instanceId, props.sessionId, messageId, partId)
}
clearDeleteMode()
@@ -1457,6 +1514,7 @@ export default function MessageSection(props: MessageSectionProps) {
onClearSelection={handleClearTimelineSelection}
selectedIds={selectedTimelineIds}
expandedMessageIds={expandedMessageIds}
deletableMessageIds={deletableMessageIds}
activeSegmentId={activeSegmentId()}
instanceId={props.instanceId}
sessionId={props.sessionId}

25
packages/ui/src/components/message-timeline.tsx
View File

@@ -36,6 +36,9 @@ interface MessageTimelineProps {
onClearSelection?: () => void
selectedIds?: Accessor<Set<string>>
expandedMessageIds?: Accessor<Set<string>>
// Optional: restrict histogram/xray overlay to only show for these message ids.
// Used to hide ribs for messages before the last compaction.
deletableMessageIds?: Accessor<Set<string>>
activeSegmentId?: string | null
instanceId: string
sessionId: string
@@ -319,6 +322,12 @@ const MessageTimeline: Component<MessageTimelineProps> = (props) => {
const showTools = () => props.showToolSegments ?? true
const deleteHover = () => props.deleteHover?.() ?? { kind: "none" as const }
const isHistogramEligible = (segment: TimelineSegment): boolean => {
const allowed = props.deletableMessageIds?.()
if (!allowed) return true
return allowed.has(segment.messageId)
}
const registerButtonRef = (segmentId: string, element: HTMLButtonElement | null) => {
if (element) {
buttonRefs.set(segmentId, element)
@@ -396,6 +405,14 @@ const MessageTimeline: Component<MessageTimelineProps> = (props) => {
// --- Selection & histogram rib state ---
const isSelectionActive = createMemo(() => (props.selectedIds?.().size ?? 0) > 0)
// Segments eligible for xray ribs. We intentionally exclude messages before
// the last compaction (when provided by the parent) to avoid misleading token
// weights for content that's no longer in context.
const xraySegments = createMemo(() => {
if (!isSelectionActive()) return [] as TimelineSegment[]
return props.segments.filter((segment) => isHistogramEligible(segment))
})
// Stable layout offsets per badge (relative to scroll content), recomputed only
// on activation, resize, or expansion — NOT on every scroll frame.
const [badgeOffsets, setBadgeOffsets] = createSignal<Record<string, { layoutTop: number; height: number }>>({})
@@ -495,7 +512,7 @@ const MessageTimeline: Component<MessageTimelineProps> = (props) => {
// O(n) pre-pass: group segments by messageId for O(1) lookups below.
const segmentsByMessageId = new Map<string, TimelineSegment[]>()
for (const s of props.segments) {
for (const s of xraySegments()) {
let list = segmentsByMessageId.get(s.messageId)
if (!list) {
list = []
@@ -540,7 +557,7 @@ const MessageTimeline: Component<MessageTimelineProps> = (props) => {
const aggregateTokensByMessageId = createMemo(() => {
const chars = liveSegmentChars()
const result: Record<string, number> = {}
for (const s of props.segments) {
for (const s of xraySegments()) {
result[s.messageId] = (result[s.messageId] ?? 0) + (chars[s.id] ?? s.totalChars)
}
for (const id of Object.keys(result)) {
@@ -574,7 +591,7 @@ const MessageTimeline: Component<MessageTimelineProps> = (props) => {
const maxTokens = createMemo(() => {
let max = 0
for (const s of props.segments) {
for (const s of xraySegments()) {
const tokens = getSegmentTokens(s)
if (tokens > max) max = tokens
}
@@ -881,7 +898,7 @@ const MessageTimeline: Component<MessageTimelineProps> = (props) => {
<Portal>
<Show when={isSelectionActive()}>
<div class="message-timeline-xray-overlay" style={{ "--max-rib-width": `${maxRibWidth()}px`, "clip-path": `inset(${clipBounds().top}px 0 ${(typeof window !== "undefined" ? window.innerHeight : 0) - clipBounds().bottom}px 0)` }}>
<For each={props.segments}>
<For each={xraySegments()}>
{(segment) => {
// Derive screen position from stable layout offset + scroll state.
// Only arithmetic — no DOM reads per segment per scroll frame.

36
packages/ui/src/stores/message-v2/instance-store.ts
View File

@@ -218,6 +218,9 @@ export interface InstanceMessageStore {
getScrollSnapshot: (sessionId: string, scope: string) => ScrollSnapshot | undefined
getSessionRevision: (sessionId: string) => number
getSessionMessageIds: (sessionId: string) => string[]
// Index of the most recent message in the session that contains a compaction part.
// Returns -1 if there has been no compaction.
getLastCompactionMessageIndex: (sessionId: string) => number
getMessage: (messageId: string) => MessageRecord | undefined
getLatestTodoSnapshot: (sessionId: string) => LatestTodoSnapshot | undefined
clearSession: (sessionId: string) => void
@@ -231,6 +234,24 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt
const messageInfoCache = new Map<string, MessageInfo>()
function getLastCompactionMessageIndex(sessionId: string): number {
if (!sessionId) return -1
const ids = state.sessions[sessionId]?.messageIds ?? []
// Scan from the end: we only care about the most recent compaction.
for (let i = ids.length - 1; i >= 0; i--) {
const messageId = ids[i]
const record = state.messages[messageId]
if (!record || !Array.isArray(record.partIds) || record.partIds.length === 0) continue
for (const partId of record.partIds) {
const part = record.parts[partId]?.data
if ((part as any)?.type === "compaction") {
return i
}
}
}
return -1
}
function isCompletedTodoPart(part: ClientPart | undefined): boolean {
if (!part || (part as any).type !== "tool") {
return false
@@ -1138,8 +1159,8 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt
function clearInstance() {
messageInfoCache.clear()
setState(reconcile(createInitialState(instanceId)))
}
setState(reconcile(createInitialState(instanceId)))
}
return {
@@ -1172,10 +1193,11 @@ export function createInstanceMessageStore(instanceId: string, hooks?: MessageSt
setScrollSnapshot,
getScrollSnapshot,
getSessionRevision: getSessionRevisionValue,
getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [],
getMessage: (messageId: string) => state.messages[messageId],
getLatestTodoSnapshot: (sessionId: string) => state.latestTodos[sessionId],
clearSession,
clearInstance,
getSessionMessageIds: (sessionId: string) => state.sessions[sessionId]?.messageIds ?? [],
getLastCompactionMessageIndex,
getMessage: (messageId: string) => state.messages[messageId],
getLatestTodoSnapshot: (sessionId: string) => state.latestTodos[sessionId],
clearSession,
clearInstance,
}
}