@@ -594,7 +594,7 @@ final class ChatViewModel {
594594 // Cancel the stale stream and reconnect via service logs
595595 streamTask? . cancel ( )
596596 streamTask = Task {
597- await reconnectToServiceLogs ( apiClient: apiClient, modelContext: modelContext, serviceIsRunning : true )
597+ await reconnectToServiceLogs ( apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped : false )
598598 }
599599 }
600600
@@ -642,19 +642,41 @@ final class ChatViewModel {
642642 // If the last session completed cleanly, content is already loaded from
643643 // persistence — no need to hit the network at all.
644644 if let chat = fetchChat ( modelContext: modelContext) , chat. lastSessionComplete {
645+ // Edge case: app killed between persistMessages and saveSession(isComplete:false).
646+ // The session was previously complete but a new user message was appended and
647+ // persisted before the service was created. Restore it as a draft.
648+ restoreUndeliveredDraft ( modelContext: modelContext)
645649 return
646650 }
647651
648652 streamTask = Task {
649653 // Only reconnect if the service exists (running or stopped with logs)
650654 guard let serviceInfo = try ? await apiClient. getServiceStatus ( spriteName: spriteName, serviceName: serviceName)
651- else { return }
655+ else {
656+ // Service was never created (e.g. app killed before the service call).
657+ // Restore any trailing user message as a draft rather than leaving a
658+ // stale bubble with no response.
659+ restoreUndeliveredDraft ( modelContext: modelContext)
660+ return
661+ }
652662
653- let isRunning = serviceInfo. state. status = =" running "
654- await reconnectToServiceLogs ( apiClient: apiClient, modelContext: modelContext, serviceIsRunning : isRunning )
663+ let alreadyStopped = serviceInfo. state. status ! =" running "
664+ await reconnectToServiceLogs ( apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped : alreadyStopped )
655665 }
656666 }
657667
668+ /// If the last message is a user message with no response, remove it from history
669+ /// and restore its text to the input box as a draft.
670+ func restoreUndeliveredDraft( modelContext: ModelContext ) {
671+ guard let last = messages. last, last. role == . user else { return }
672+ let text = last. textContent
673+ messages. removeLast ( )
674+ persistMessages ( modelContext: modelContext)
675+ guard !text. isEmpty, inputText. isEmpty else { return }
676+ inputText = text
677+ saveDraft ( modelContext: modelContext)
678+ }
679+
658680 // MARK: - Private
659681
660682 private func executeClaudeCommand(
@@ -796,7 +818,7 @@ final class ChatViewModel {
796818 // Attempt reconnection on disconnect
797819 if case . disconnected = streamResult {
798820 logger. info ( " [Chat] Disconnected mid-stream, attempting reconnect via service logs " )
799- await reconnectToServiceLogs ( apiClient: apiClient, modelContext: modelContext, serviceIsRunning : true )
821+ await reconnectToServiceLogs ( apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped : false )
800822 return
801823 }
802824
@@ -933,6 +955,7 @@ final class ChatViewModel {
933955 receivedData = true
934956 timeoutTask. cancel ( )
935957 if case . connecting = status { status = . streaming }
958+ else if case . reconnecting = status { status = . streaming }
936959
937960 // Two-level NDJSON: ServiceLogEvent.data contains Claude NDJSON.
938961 // The logs endpoint prefixes each line with a timestamp
@@ -961,6 +984,7 @@ final class ChatViewModel {
961984 receivedData = true
962985 timeoutTask. cancel ( )
963986 if case . connecting = status { status = . streaming }
987+ else if case . reconnecting = status { status = . streaming }
964988 if let text = event. data {
965989 logger. warning ( " Service stderr: \( text. prefix ( 500 ) , privacy: . public) " )
966990 }
@@ -994,6 +1018,7 @@ final class ChatViewModel {
9941018
9951019 case . started:
9961020 if case . connecting = status { status = . streaming }
1021+ else if case . reconnecting = status { status = . streaming }
9971022
9981023 case . stopping, . stopped:
9991024 break
@@ -1038,10 +1063,10 @@ final class ChatViewModel {
10381063 func runReconnectLoop(
10391064 apiClient: some ServiceLogsProvider ,
10401065 modelContext: ModelContext ,
1041- serviceIsRunning : Bool = false
1066+ serviceAlreadyStopped : Bool = false
10421067 ) async {
10431068 isReplaying = true
1044- isReplayingLiveService = serviceIsRunning
1069+ isReplayingLiveService = !serviceAlreadyStopped
10451070 defer {
10461071 applyReplayBuffer ( )
10471072 isReplaying = false
@@ -1060,7 +1085,11 @@ final class ChatViewModel {
10601085 if let existing = currentAssistantMessage {
10611086 assistantMessage = existing
10621087 if !hasPriorEvents { assistantMessage. content = [ ] }
1063- } else if let last = messages. last ( where: { $0. role == . assistant } ) {
1088+ } else if let last = messages. last, last. role == . assistant {
1089+ // Only reuse an assistant message if it's already at the tail — meaning
1090+ // the last exchange ended with a partial response that needs replaying.
1091+ // If the last message is a user message, fall through to create a new
1092+ // assistant message after it (each service is scoped to one user message).
10641093 assistantMessage = last
10651094 if !hasPriorEvents { assistantMessage. content = [ ] }
10661095 currentAssistantMessage = last
@@ -1108,10 +1137,12 @@ final class ChatViewModel {
11081137 // If we got a result event, Claude is done
11091138 if receivedResultEvent { break }
11101139
1111- // Check if service is still running
1112- let isRunning = ( try ? await apiClient . getServiceStatus ( spriteName : spriteName , serviceName : serviceName ) ) ? . state . status == " running "
1140+ // If we already knew the service was stopped before entering, no need to re-check
1141+ if serviceAlreadyStopped { break }
11131142
1114- if isRunning {
1143+ // Check if service is still running before retrying
1144+ if let serviceInfo = try ? await apiClient. getServiceStatus ( spriteName: spriteName, serviceName: serviceName) ,
1145+ serviceInfo. state. status == " running " {
11151146 logger. info ( " [Chat] Service still running, will re-poll after delay " )
11161147 try ? await Task . sleep ( for: . seconds( 2 ) )
11171148 continue
@@ -1147,9 +1178,9 @@ final class ChatViewModel {
11471178 private func reconnectToServiceLogs(
11481179 apiClient: SpritesAPIClient ,
11491180 modelContext: ModelContext ,
1150- serviceIsRunning : Bool = false
1181+ serviceAlreadyStopped : Bool = false
11511182 ) async {
1152- await runReconnectLoop ( apiClient: apiClient, modelContext: modelContext, serviceIsRunning : serviceIsRunning )
1183+ await runReconnectLoop ( apiClient: apiClient, modelContext: modelContext, serviceAlreadyStopped : serviceAlreadyStopped )
11531184
11541185 if let queued = queuedPrompt, !Task. isCancelled {
11551186 let prompt = buildPrompt ( text: queued, attachments: queuedAttachments)
0 commit comments