@@ -32,6 +32,7 @@ import * as Log from "@opencode-ai/core/util/log"
3232import { emptyConsoleState , type ConsoleState } from "@/config/console-state"
3333import path from "path"
3434import { useKV } from "./kv"
35+ import { mergeFetchedMessages , optimisticParts , type OptimisticPromptPart } from "./sync-optimistic"
3536
3637export const { use : useSync , provider : SyncProvider } = createSimpleContext ( {
3738 name : "Sync" ,
@@ -113,6 +114,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
113114 const [ autoaccept ] = kv . signal < "none" | "edit" > ( "permission_auto_accept" , "edit" )
114115
115116 const fullSyncedSessions = new Set < string > ( )
117+ const optimisticMessages = new Set < string > ( )
116118 let syncedWorkspace = project . workspace . current ( )
117119
118120 function sessionListQuery ( ) : { scope ?: "project" ; path ?: string } {
@@ -227,6 +229,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
227229 break
228230
229231 case "session.deleted" : {
232+ for ( const message of store . message [ event . properties . info . id ] ?? [ ] ) optimisticMessages . delete ( message . id )
230233 const result = Binary . search ( store . session , event . properties . info . id , ( s ) => s . id )
231234 if ( result . found ) {
232235 setStore (
@@ -298,6 +301,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
298301 break
299302 }
300303 case "message.removed" : {
304+ optimisticMessages . delete ( event . properties . messageID )
301305 const messages = store . message [ event . properties . sessionID ]
302306 const result = Binary . search ( messages , event . properties . messageID , ( m ) => m . id )
303307 if ( result . found ) {
@@ -312,6 +316,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
312316 break
313317 }
314318 case "message.part.updated" : {
319+ optimisticMessages . delete ( event . properties . part . messageID )
315320 const parts = store . part [ event . properties . part . messageID ]
316321 if ( ! parts ) {
317322 setStore ( "part" , event . properties . part . messageID , [ event . properties . part ] )
@@ -385,6 +390,7 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
385390 const workspace = project . workspace . current ( )
386391 if ( workspace !== syncedWorkspace ) {
387392 fullSyncedSessions . clear ( )
393+ optimisticMessages . clear ( )
388394 syncedWorkspace = workspace
389395 }
390396 const projectPromise = project . sync ( )
@@ -520,6 +526,66 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
520526 if ( last . role === "user" ) return "working"
521527 return last . time . completed ? "idle" : "working"
522528 } ,
529+ addOptimisticPrompt ( input : {
530+ sessionID : string
531+ messageID : string
532+ agent : string
533+ model : { providerID : string ; modelID : string }
534+ variant ?: string
535+ parts : OptimisticPromptPart [ ]
536+ } ) {
537+ optimisticMessages . add ( input . messageID )
538+ const messages = store . message [ input . sessionID ]
539+ const match = messages ? Binary . search ( messages , input . messageID , ( m ) => m . id ) : undefined
540+ const info : Message = {
541+ id : input . messageID ,
542+ sessionID : input . sessionID ,
543+ role : "user" ,
544+ time : { created : Date . now ( ) } ,
545+ agent : input . agent ,
546+ model : {
547+ providerID : input . model . providerID ,
548+ modelID : input . model . modelID ,
549+ ...( input . variant ? { variant : input . variant } : { } ) ,
550+ } ,
551+ }
552+ batch ( ( ) => {
553+ if ( ! messages ) {
554+ setStore ( "message" , input . sessionID , [ info ] )
555+ } else if ( ! match ?. found ) {
556+ setStore (
557+ "message" ,
558+ input . sessionID ,
559+ produce ( ( draft ) => {
560+ Binary . insert ( draft , info , ( message ) => message . id )
561+ } ) ,
562+ )
563+ }
564+ setStore ( "part" , input . messageID , reconcile ( optimisticParts ( input ) ) )
565+ } )
566+ } ,
567+ removeOptimisticPrompt ( sessionID : string , messageID : string ) {
568+ if ( ! optimisticMessages . delete ( messageID ) ) return
569+ const messages = store . message [ sessionID ]
570+ const match = messages ? Binary . search ( messages , messageID , ( m ) => m . id ) : undefined
571+ batch ( ( ) => {
572+ if ( match ?. found ) {
573+ setStore (
574+ "message" ,
575+ sessionID ,
576+ produce ( ( draft ) => {
577+ draft . splice ( match . index , 1 )
578+ } ) ,
579+ )
580+ }
581+ setStore (
582+ "part" ,
583+ produce ( ( draft ) => {
584+ delete draft [ messageID ]
585+ } ) ,
586+ )
587+ } )
588+ } ,
523589 async sync ( sessionID : string ) {
524590 if ( fullSyncedSessions . has ( sessionID ) ) return
525591 const [ session , messages , todo , diff ] = await Promise . all ( [
@@ -531,12 +597,21 @@ export const { use: useSync, provider: SyncProvider } = createSimpleContext({
531597 setStore (
532598 produce ( ( draft ) => {
533599 const match = Binary . search ( draft . session , sessionID , ( s ) => s . id )
600+ const merged = mergeFetchedMessages ( {
601+ currentMessages : draft . message [ sessionID ] ?? [ ] ,
602+ currentParts : draft . part ,
603+ fetched : messages . data ! ,
604+ optimisticMessages,
605+ } )
534606 if ( match . found ) draft . session [ match . index ] = session . data !
535607 if ( ! match . found ) draft . session . splice ( match . index , 0 , session . data ! )
536608 draft . todo [ sessionID ] = todo . data ?? [ ]
537- draft . message [ sessionID ] = messages . data ! . map ( ( x ) => x . info )
538- for ( const message of messages . data ! ) {
539- draft . part [ message . info . id ] = message . parts
609+ draft . message [ sessionID ] = merged . messages
610+ for ( const messageID of merged . resolved ) {
611+ optimisticMessages . delete ( messageID )
612+ }
613+ for ( const [ messageID , parts ] of merged . parts ) {
614+ draft . part [ messageID ] = parts
540615 }
541616 draft . session_diff [ sessionID ] = diff . data ?? [ ]
542617 } ) ,
0 commit comments