@@ -33,6 +33,7 @@ import * as Log from "@opencode-ai/core/util/log"
3333import { emptyConsoleState , type ConsoleState } from "@/config/console-state"
3434import path from "path"
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,7 @@ 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 > ( )
117119 let syncedWorkspace = project . workspace . current ( )
118120
119121 function sessionListQuery ( ) : { scope ?: "project" ; path ?: string } {
@@ -228,6 +230,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
228230 break
229231
230232 case "session.deleted" : {
233+ for ( const message of store . message [ event . properties . info . id ] ?? [ ] ) optimisticMessages . delete ( message . id )
231234 const result = Binary . search ( store . session , event . properties . info . id , ( s ) => s . id )
232235 if ( result . found ) {
233236 setStore (
@@ -299,6 +302,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
299302 break
300303 }
301304 case "message.removed" : {
305+ optimisticMessages . delete ( event . properties . messageID )
302306 const messages = store . message [ event . properties . sessionID ]
303307 const result = Binary . search ( messages , event . properties . messageID , ( m ) => m . id )
304308 if ( result . found ) {
@@ -313,6 +317,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
313317 break
314318 }
315319 case "message.part.updated" : {
320+ optimisticMessages . delete ( event . properties . part . messageID )
316321 const parts = store . part [ event . properties . part . messageID ]
317322 if ( ! parts ) {
318323 setStore ( "part" , event . properties . part . messageID , [ event . properties . part ] )
@@ -386,6 +391,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
386391 const workspace = project . workspace . current ( )
387392 if ( workspace !== syncedWorkspace ) {
388393 fullSyncedSessions . clear ( )
394+ optimisticMessages . clear ( )
389395 syncedWorkspace = workspace
390396 }
391397 const projectPromise = project . sync ( )
@@ -528,6 +534,66 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
528534 if ( last . role === "user" ) return "working"
529535 return last . time . completed ? "idle" : "working"
530536 } ,
537+ addOptimisticPrompt ( input : {
538+ sessionID : string
539+ messageID : string
540+ agent : string
541+ model : { providerID : string ; modelID : string }
542+ variant ?: string
543+ parts : OptimisticPromptPart [ ]
544+ } ) {
545+ optimisticMessages . add ( input . messageID )
546+ const messages = store . message [ input . sessionID ]
547+ const match = messages ? Binary . search ( messages , input . messageID , ( m ) => m . id ) : undefined
548+ const info : Message = {
549+ id : input . messageID ,
550+ sessionID : input . sessionID ,
551+ role : "user" ,
552+ time : { created : Date . now ( ) } ,
553+ agent : input . agent ,
554+ model : {
555+ providerID : input . model . providerID ,
556+ modelID : input . model . modelID ,
557+ ...( input . variant ? { variant : input . variant } : { } ) ,
558+ } ,
559+ }
560+ batch ( ( ) => {
561+ if ( ! messages ) {
562+ setStore ( "message" , input . sessionID , [ info ] )
563+ } else if ( ! match ?. found ) {
564+ setStore (
565+ "message" ,
566+ input . sessionID ,
567+ produce ( ( draft ) => {
568+ Binary . insert ( draft , info , ( message ) => message . id )
569+ } ) ,
570+ )
571+ }
572+ setStore ( "part" , input . messageID , reconcile ( optimisticParts ( input ) ) )
573+ } )
574+ } ,
575+ removeOptimisticPrompt ( sessionID : string , messageID : string ) {
576+ if ( ! optimisticMessages . delete ( messageID ) ) return
577+ const messages = store . message [ sessionID ]
578+ const match = messages ? Binary . search ( messages , messageID , ( m ) => m . id ) : undefined
579+ batch ( ( ) => {
580+ if ( match ?. found ) {
581+ setStore (
582+ "message" ,
583+ sessionID ,
584+ produce ( ( draft ) => {
585+ draft . splice ( match . index , 1 )
586+ } ) ,
587+ )
588+ }
589+ setStore (
590+ "part" ,
591+ produce ( ( draft ) => {
592+ delete draft [ messageID ]
593+ } ) ,
594+ )
595+ } )
596+ } ,
531597 async sync ( sessionID : string ) {
532598 if ( fullSyncedSessions . has ( sessionID ) ) return
533599 const [ session , messages , todo , diff ] = await Promise . all ( [
@@ -539,15 +605,22 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
539605 setStore (
540606 produce ( ( draft ) => {
541607 const match = Binary . search ( draft . session , sessionID , ( s ) => s . id )
608+ const merged = mergeFetchedMessages ( {
609+ currentMessages : draft . message [ sessionID ] ?? [ ] ,
610+ currentParts : draft . part ,
611+ fetched : messages . data ! ,
612+ optimisticMessages,
613+ } )
542614 if ( match . found ) draft . session [ match . index ] = session . data !
543615 if ( ! match . found ) draft . session . splice ( match . index , 0 , session . data ! )
544616 draft . todo [ sessionID ] = todo . data ?? [ ]
545- const infos : ( typeof draft . message ) [ string ] = [ ]
546- for ( const message of messages . data ?? [ ] ) {
547- infos . push ( message . info )
548- draft . part [ message . info . id ] = message . parts
617+ draft . message [ sessionID ] = merged . messages
618+ for ( const messageID of merged . resolved ) {
619+ optimisticMessages . delete ( messageID )
620+ }
621+ for ( const [ messageID , parts ] of merged . parts ) {
622+ draft . part [ messageID ] = parts
549623 }
550- draft . message [ sessionID ] = infos
551624 draft . session_diff [ sessionID ] = diff . data ?? [ ]
552625 } ) ,
553626 )
0 commit comments