@@ -33,6 +33,7 @@ import { emptyConsoleState, type ConsoleState } from "@/config/console-state"
3333import path from "path"
3434import { useKV } from "./kv"
3535import { aggregateFailures } from "./aggregate-failures"
36+ import { mergeFetchedMessages , optimisticParts , type OptimisticPromptPart } from "./sync-optimistic"
3637
3738export const { use : useSync , provider : SyncProvider } = createSimpleContext ( {
3839 name : "Sync" ,
@@ -114,6 +115,8 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
114115 const [ autoaccept ] = kv . signal < "none" | "edit" > ( "permission_auto_accept" , "edit" )
115116
116117 const fullSyncedSessions = new Set < string > ( )
118+ const optimisticMessages = new Set < string > ( )
119+ let syncedWorkspace = project . workspace . current ( )
117120
118121 function sessionListQuery ( ) : { scope ?: "project" ; path ?: string } {
119122 if ( ! kv . get ( "session_directory_filter_enabled" , true ) ) return { scope : "project" }
@@ -227,6 +230,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
227230 break
228231
229232 case "session.deleted" : {
233+ for ( const message of store . message [ event . properties . info . id ] ?? [ ] ) optimisticMessages . delete ( message . id )
230234 const result = Binary . search ( store . session , event . properties . info . id , ( s ) => s . id )
231235 if ( result . found ) {
232236 setStore (
@@ -298,6 +302,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
298302 break
299303 }
300304 case "message.removed" : {
305+ optimisticMessages . delete ( event . properties . messageID )
301306 const messages = store . message [ event . properties . sessionID ]
302307 const result = Binary . search ( messages , event . properties . messageID , ( m ) => m . id )
303308 if ( result . found ) {
@@ -312,6 +317,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
312317 break
313318 }
314319 case "message.part.updated" : {
320+ optimisticMessages . delete ( event . properties . part . messageID )
315321 const parts = store . part [ event . properties . part . messageID ]
316322 if ( ! parts ) {
317323 setStore ( "part" , event . properties . part . messageID , [ event . properties . part ] )
@@ -386,6 +392,11 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
386392 async function bootstrap ( input : { fatal ?: boolean } = { } ) {
387393 const fatal = input . fatal ?? true
388394 const workspace = project . workspace . current ( )
395+ if ( workspace !== syncedWorkspace ) {
396+ fullSyncedSessions . clear ( )
397+ optimisticMessages . clear ( )
398+ syncedWorkspace = workspace
399+ }
389400 const projectPromise = project . sync ( )
390401 const sessionListPromise = projectPromise . then ( ( ) => listSessions ( ) )
391402
@@ -526,6 +537,66 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
526537 if ( last . role === "user" ) return "working"
527538 return last . time . completed ? "idle" : "working"
528539 } ,
540+ addOptimisticPrompt ( input : {
541+ sessionID : string
542+ messageID : string
543+ agent : string
544+ model : { providerID : string ; modelID : string }
545+ variant ?: string
546+ parts : OptimisticPromptPart [ ]
547+ } ) {
548+ optimisticMessages . add ( input . messageID )
549+ const messages = store . message [ input . sessionID ]
550+ const match = messages ? Binary . search ( messages , input . messageID , ( m ) => m . id ) : undefined
551+ const info : Message = {
552+ id : input . messageID ,
553+ sessionID : input . sessionID ,
554+ role : "user" ,
555+ time : { created : Date . now ( ) } ,
556+ agent : input . agent ,
557+ model : {
558+ providerID : input . model . providerID ,
559+ modelID : input . model . modelID ,
560+ ...( input . variant ? { variant : input . variant } : { } ) ,
561+ } ,
562+ }
563+ batch ( ( ) => {
564+ if ( ! messages ) {
565+ setStore ( "message" , input . sessionID , [ info ] )
566+ } else if ( ! match ?. found ) {
567+ setStore (
568+ "message" ,
569+ input . sessionID ,
570+ produce ( ( draft ) => {
571+ Binary . insert ( draft , info , ( message ) => message . id )
572+ } ) ,
573+ )
574+ }
575+ setStore ( "part" , input . messageID , reconcile ( optimisticParts ( input ) ) )
576+ } )
577+ } ,
578+ removeOptimisticPrompt ( sessionID : string , messageID : string ) {
579+ if ( ! optimisticMessages . delete ( messageID ) ) return
580+ const messages = store . message [ sessionID ]
581+ const match = messages ? Binary . search ( messages , messageID , ( m ) => m . id ) : undefined
582+ batch ( ( ) => {
583+ if ( match ?. found ) {
584+ setStore (
585+ "message" ,
586+ sessionID ,
587+ produce ( ( draft ) => {
588+ draft . splice ( match . index , 1 )
589+ } ) ,
590+ )
591+ }
592+ setStore (
593+ "part" ,
594+ produce ( ( draft ) => {
595+ delete draft [ messageID ]
596+ } ) ,
597+ )
598+ } )
599+ } ,
529600 async sync ( sessionID : string ) {
530601 if ( fullSyncedSessions . has ( sessionID ) ) return
531602 const [ session , messages , todo , diff ] = await Promise . all ( [
@@ -537,15 +608,22 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
537608 setStore (
538609 produce ( ( draft ) => {
539610 const match = Binary . search ( draft . session , sessionID , ( s ) => s . id )
611+ const merged = mergeFetchedMessages ( {
612+ currentMessages : draft . message [ sessionID ] ?? [ ] ,
613+ currentParts : draft . part ,
614+ fetched : messages . data ! ,
615+ optimisticMessages,
616+ } )
540617 if ( match . found ) draft . session [ match . index ] = session . data !
541618 if ( ! match . found ) draft . session . splice ( match . index , 0 , session . data ! )
542619 draft . todo [ sessionID ] = todo . data ?? [ ]
543- const infos : ( typeof draft . message ) [ string ] = [ ]
544- for ( const message of messages . data ?? [ ] ) {
545- infos . push ( message . info )
546- draft . part [ message . info . id ] = message . parts
620+ draft . message [ sessionID ] = merged . messages
621+ for ( const messageID of merged . resolved ) {
622+ optimisticMessages . delete ( messageID )
623+ }
624+ for ( const [ messageID , parts ] of merged . parts ) {
625+ draft . part [ messageID ] = parts
547626 }
548- draft . message [ sessionID ] = infos
549627 draft . session_diff [ sessionID ] = diff . data ?? [ ]
550628 } ) ,
551629 )
0 commit comments