fix(tui): preserve live parts during session hydration (#30300)

This commit is contained in:
Kit Langton
2026-06-01 21:40:29 -04:00
committed by GitHub
parent 8dc2ffd48f
commit 1cae8f89c5
2 changed files with 361 additions and 22 deletions

View File

@@ -113,6 +113,14 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
const kv = useKV()
const fullSyncedSessions = new Set<string>()
const syncingSessions = new Map<string, Promise<void>>()
const hydratingSessions = new Map<string, { messages: Set<string>; parts: Set<string> }>()
const touchMessage = (sessionID: string, messageID: string) => {
hydratingSessions.get(sessionID)?.messages.add(messageID)
}
const touchPart = (sessionID: string, partID: string) => {
hydratingSessions.get(sessionID)?.parts.add(partID)
}
function sessionListQuery(): { scope?: "project"; path?: string } {
if (!kv.get("session_directory_filter_enabled", true)) return { scope: "project" }
@@ -251,6 +259,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
}
case "message.updated": {
touchMessage(event.properties.info.sessionID, event.properties.info.id)
const messages = store.message[event.properties.info.sessionID]
if (!messages) {
setStore("message", event.properties.info.sessionID, [event.properties.info])
@@ -290,6 +299,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
break
}
case "message.removed": {
touchMessage(event.properties.sessionID, event.properties.messageID)
const messages = store.message[event.properties.sessionID]
const result = Binary.search(messages, event.properties.messageID, (m) => m.id)
if (result.found) {
@@ -304,6 +314,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
break
}
case "message.part.updated": {
touchPart(event.properties.part.sessionID, event.properties.part.id)
const parts = store.part[event.properties.part.messageID]
if (!parts) {
setStore("part", event.properties.part.messageID, [event.properties.part])
@@ -329,6 +340,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
if (!parts) break
const result = Binary.search(parts, event.properties.partID, (p) => p.id)
if (!result.found) break
touchPart(event.properties.sessionID, event.properties.partID)
setStore(
"part",
event.properties.messageID,
@@ -343,6 +355,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
}
case "message.part.removed": {
touchPart(event.properties.sessionID, event.properties.partID)
const parts = store.part[event.properties.messageID]
const result = Binary.search(parts, event.properties.partID, (p) => p.id)
if (result.found) {
@@ -520,28 +533,76 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
},
async sync(sessionID: string) {
if (fullSyncedSessions.has(sessionID)) return
const [session, messages, todo, diff] = await Promise.all([
sdk.client.session.get({ sessionID }, { throwOnError: true }),
sdk.client.session.messages({ sessionID, limit: 100 }),
sdk.client.session.todo({ sessionID }),
sdk.client.session.diff({ sessionID }),
])
setStore(
produce((draft) => {
const match = Binary.search(draft.session, sessionID, (s) => s.id)
if (match.found) draft.session[match.index] = session.data!
if (!match.found) draft.session.splice(match.index, 0, session.data!)
draft.todo[sessionID] = todo.data ?? []
const infos: (typeof draft.message)[string] = []
for (const message of messages.data ?? []) {
infos.push(message.info)
draft.part[message.info.id] = message.parts
}
draft.message[sessionID] = infos
draft.session_diff[sessionID] = diff.data ?? []
}),
)
fullSyncedSessions.add(sessionID)
const syncing = syncingSessions.get(sessionID)
if (syncing) return syncing
const tracker = { messages: new Set<string>(), parts: new Set<string>() }
hydratingSessions.set(sessionID, tracker)
const task = (async () => {
const [session, messages, todo, diff] = await Promise.all([
sdk.client.session.get({ sessionID }, { throwOnError: true }),
sdk.client.session.messages({ sessionID, limit: 100 }),
sdk.client.session.todo({ sessionID }),
sdk.client.session.diff({ sessionID }),
])
setStore(
produce((draft) => {
const match = Binary.search(draft.session, sessionID, (s) => s.id)
if (match.found) draft.session[match.index] = session.data!
if (!match.found) draft.session.splice(match.index, 0, session.data!)
draft.todo[sessionID] = todo.data ?? []
const currentMessages = draft.message[sessionID] ?? []
const infos = (messages.data ?? []).flatMap((message) => {
if (!tracker.messages.has(message.info.id)) return [message.info]
const current = currentMessages.find((item) => item.id === message.info.id)
return current ? [current] : []
})
infos.push(
...currentMessages.filter(
(message) => tracker.messages.has(message.id) && !infos.some((item) => item.id === message.id),
),
)
const removed = infos.slice(0, -100)
const visible = infos.slice(-100)
const visibleIDs = new Set(visible.map((message) => message.id))
for (const message of messages.data ?? []) {
if (!visibleIDs.has(message.info.id)) {
delete draft.part[message.info.id]
continue
}
const currentParts = draft.part[message.info.id] ?? []
const parts = message.parts.flatMap((part) => {
const current = currentParts.find((item) => item.id === part.id)
if (tracker.parts.has(part.id)) return current ? [current] : []
if (
current &&
(part.type === "text" || part.type === "reasoning") &&
(current.type === "text" || current.type === "reasoning") &&
part.text.length === 0 &&
current.text.length > 0
) {
return [current]
}
return [part]
})
parts.push(
...currentParts.filter(
(part) => tracker.parts.has(part.id) && !parts.some((item) => item.id === part.id),
),
)
draft.part[message.info.id] = parts
}
for (const message of removed) delete draft.part[message.id]
draft.message[sessionID] = visible
draft.session_diff[sessionID] = diff.data ?? []
}),
)
fullSyncedSessions.add(sessionID)
})().finally(() => {
syncingSessions.delete(sessionID)
hydratingSessions.delete(sessionID)
})
syncingSessions.set(sessionID, task)
return task
},
},
bootstrap,