fix(tui): preserve live parts during session hydration (#30300)
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user