Fix Assistant edits at message boundaries

FrameworkTests/AgentReplayPersistenceTests.swift

@@ -29,6 +29,8 @@ struct AgentReplayPersistenceTests {
         try await repository.save(snapshot)
         // A subsequent run must neither overwrite an earlier snapshot nor duplicate its backing.
         snapshot.id = UUID()
+        snapshot.startedAt = snapshot.startedAt.addingTimeInterval(1)
+        snapshot.finishedAt = snapshot.startedAt
         snapshot.messages.append(.init(role: .user, text: "continue", contextSnapshot: .init()))
         try await repository.save(snapshot)
         let restored = try #require(try await repository.loadConversations().first)
@@ -58,9 +60,55 @@ struct AgentReplayPersistenceTests {
         var result = AgentToolResult(callID: call.id, content: "full captured result")
         result.modelContent = "frozen preview"
         try await repository.record(.toolFinished(invocation, result), runID: snapshot.id)
-        let restored = try #require(try await repository.loadConversations().first?.messages.last)
+        let conversation = try #require(try await repository.loadConversations().first)
+        let restored = try #require(conversation.messages.last)
+        let resultID = AgentTranscriptMessage.toolResultID(runID: snapshot.id, callID: call.id)
         #expect(restored.modelText == "frozen preview")
-        #expect(restored.id == AgentTranscriptMessage.toolResultID(runID: snapshot.id, callID: call.id))
+        #expect(restored.id == resultID)
+        #expect(conversation.messageRunIDs[resultID] == snapshot.id)
+    }
+
+    @Test(arguments: [AgentRunState.failed, .interrupted])
+    func aRealRunAtomicallyAdoptsTheEditedStagingMessageAndReplayBacking(
+        state: AgentRunState
+    ) async throws {
+        let container = try TermindPersistence.makeContainer(inMemory: true)
+        let repository = SwiftDataAgentRunRepository(context: container.mainContext)
+        let conversationID = UUID()
+        let original = AgentTranscriptMessage(role: .user, text: "before", authoredText: "before")
+        let originalRun = AgentRunSnapshot(
+            conversationID: conversationID, state: .completed,
+            permissionMode: .askForApproval, messages: [original])
+        try await repository.save(originalRun)
+        let replacement = AgentTranscriptMessage(
+            role: .user, text: "after", authoredText: "after",
+            contextSnapshot: .init(sessionContext: "captured host"))
+        try await repository.truncateConversation(
+            conversationID, fromUserMessageID: original.id, replacingWith: replacement)
+
+        let realRun = AgentRunSnapshot(
+            conversationID: conversationID, state: state,
+            permissionMode: .askForApproval, messages: [replacement],
+            failure: state == .failed ? "preparation failed" : nil)
+        try await repository.save(realRun)
+
+        let messageRows = try container.mainContext.fetch(FetchDescriptor<PersistentAgentMessage>())
+        let runRows = try container.mainContext.fetch(FetchDescriptor<PersistentAgentRun>())
+        let replayRows = try container.mainContext.fetch(FetchDescriptor<PersistentAgentMessageReplay>())
+        #expect(messageRows.count == 1)
+        #expect(messageRows.first?.id == replacement.id)
+        #expect(messageRows.first?.runID == realRun.id)
+        #expect(runRows.map(\.id) == [realRun.id])
+        #expect(replayRows.map(\.messageID) == [replacement.id])
+
+        let freshRepository = SwiftDataAgentRunRepository(context: ModelContext(container))
+        let restored = try #require(try await freshRepository.loadConversations().first)
+        #expect(restored.messages.map(\.id) == [replacement.id])
+        #expect(restored.messages.first?.contextSnapshot == replacement.contextSnapshot)
+        #expect(restored.messageRunIDs[replacement.id] == realRun.id)
+        #expect(restored.runs.map(\.id) == [realRun.id])
+        #expect(restored.runs.first?.state == state)
+        #expect(restored.runs.first?.failure == (state == .failed ? "preparation failed" : nil))
     }
 
     @Test

FrameworkTests/AssistantConversationOperationsTests.swift

@@ -224,6 +224,67 @@ struct AssistantConversationOperationsTests {
         try await operations.delete(conversation)
         #expect(await repository.loadConversations().isEmpty)
     }
+
+    @Test func invalidReplacementNeverCallsPersistence() async throws {
+        let repository = TruncationSpyAgentRunRepository()
+        let operations = AssistantConversationOperations(repository: repository)
+        let original = AgentTranscriptMessage(role: .user, text: "Before", authoredText: "Before")
+        let conversation = ConversationModel(
+            messages: [ChatMessage(id: original.id, role: .user, text: original.text)],
+            transcript: [original])
+        let invalid = AgentTranscriptMessage(role: .assistant, text: "Not a user turn")
+
+        #expect(
+            try await operations.edit(
+                conversation, messageID: original.id, replacement: invalid) == nil)
+        #expect(await repository.truncationCount == 0)
+        #expect(conversation.transcript == [original])
+        #expect(!conversation.isHistoryChanging)
+    }
+
+    @Test func truncationFailureLeavesInMemoryHistoryUnchanged() async throws {
+        let original = AgentTranscriptMessage(role: .user, text: "Before", authoredText: "Before")
+        let answer = AgentTranscriptMessage(role: .assistant, text: "Existing answer")
+        let conversation = ConversationModel(
+            messages: [
+                ChatMessage(id: original.id, role: .user, text: original.text),
+                ChatMessage(id: answer.id, role: .assistant, text: answer.text),
+            ],
+            transcript: [original, answer])
+        let originalMessages = conversation.messages
+        let originalTranscript = conversation.transcript
+        let replacement = AgentTranscriptMessage(role: .user, text: "After", authoredText: "After")
+        let operations = AssistantConversationOperations(repository: RejectingConversationRepository())
+
+        await #expect(throws: CocoaError.self) {
+            try await operations.edit(
+                conversation, messageID: original.id, replacement: replacement)
+        }
+
+        #expect(conversation.messages == originalMessages)
+        #expect(conversation.transcript == originalTranscript)
+        #expect(!conversation.isHistoryChanging)
+    }
+}
+
+private actor TruncationSpyAgentRunRepository: AgentRunPersisting {
+    private(set) var truncationCount = 0
+
+    func save(_ snapshot: AgentRunSnapshot) {}
+    func record(_ event: AgentEvent, runID: UUID) {}
+    func persistCompaction(_ record: AgentCompactionRecord, conversationID: UUID) {}
+    func loadConversations() -> [StoredAgentConversation] { [] }
+    func deleteConversation(_ id: UUID) {}
+    func setTitle(_ title: String, conversationID: UUID) {}
+    func setConversationMetadata(_ metadata: StoredAgentConversationMetadata, conversationID: UUID) {}
+    func reorderConversations(_ updates: [StoredAgentConversationOrder]) {}
+    func truncateConversation(
+        _ id: UUID, fromUserMessageID messageID: UUID,
+        replacingWith message: AgentTranscriptMessage
+    ) {
+        truncationCount += 1
+    }
+    func createConversation(_ conversation: StoredAgentConversation) {}
 }
 
 private actor RejectingConversationRepository: AgentRunPersisting {

FrameworkTests/AssistantExecutionControllerTests.swift

@@ -81,6 +81,7 @@ struct AssistantExecutionControllerTests {
 
         try await waitUntil { !conversation.isPreparingAgentRun }
         #expect(conversation.runState == .failed)
+        await repository.waitUntilSnapshot(state: .failed)
         #expect(await repository.snapshots.last?.state == .failed)
     }
 
@@ -249,6 +250,7 @@ struct AssistantExecutionControllerTests {
         await memoryStore.releaseSnapshot()
         try await waitUntil { !preparing.isPreparingAgentRun }
         #expect(preparing.runState == .failed)
+        await repository.waitUntilSnapshot(state: .failed)
         #expect(await repository.snapshots.last?.state == .failed)
     }
 
@@ -492,6 +494,43 @@ struct AssistantExecutionControllerTests {
         #expect(edited.contextSnapshot?.sessionContext?.contains("captured-host") == true)
         #expect(edited.contextSnapshot?.planContract == TermindAgentPolicy.planContract)
         #expect(conversation.runState == .failed)
+        let failed = try #require(conversation.runSummaries.values.first)
+        #expect(conversation.runSummaries.count == 1)
+        #expect(failed.state == .failed)
+        #expect(failed.id != edited.id)
+        #expect(conversation.messageRunIDs[edited.id] == failed.id)
+        #expect(conversation.runSummaries[edited.id] == nil)
+    }
+
+    @Test
+    func cancellationAfterSuccessfulEditAdoptsTheInterruptedRunProvenance() async throws {
+        let memoryStore = SnapshotBarrierMemoryStore()
+        let memory = MemoryModel(storage: memoryStore)
+        await memory.setEnabled(true)
+        let repository = RecordingAgentRunRepository()
+        let controller = makeController(repository: repository, memory: memory)
+        let operations = AssistantConversationOperations(repository: repository)
+        let message = AgentTranscriptMessage(role: .user, text: "before", authoredText: "before")
+        let conversation = ConversationModel(
+            messages: [ChatMessage(id: message.id, role: .user, text: "before")], transcript: [message])
+        let prepared = try #require(
+            controller.prepareAgentEdit(conversation, messageID: message.id, text: "after"))
+
+        #expect(try await controller.performAgentEdit(prepared, in: conversation, operations: operations))
+        let edited = try #require(conversation.transcript.first)
+        #expect(conversation.messageRunIDs[edited.id] == edited.id)
+        await memoryStore.waitUntilSnapshotEntered()
+
+        let cancellation = controller.cancelPendingRuns()
+        await memoryStore.releaseSnapshot()
+        await cancellation.value
+
+        let interrupted = try #require(await repository.snapshots.last)
+        #expect(interrupted.state == .interrupted)
+        #expect(conversation.runState == .interrupted)
+        #expect(conversation.messageRunIDs[edited.id] == interrupted.id)
+        #expect(conversation.runSummaries[edited.id] == nil)
+        #expect(conversation.runSummaries[interrupted.id]?.state == .interrupted)
     }
 
     private func makeController(
@@ -604,6 +643,7 @@ private actor RecordingAgentRunRepository: AgentRunPersisting {
     private var saveReleased = false
     private var saveEntryWaiters: [CheckedContinuation<Void, Never>] = []
     private var saveReleaseWaiters: [CheckedContinuation<Void, Never>] = []
+    private var snapshotWaiters: [(state: AgentRunState, continuation: CheckedContinuation<Void, Never>)] = []
     private var truncationEntered = false
     private var truncationReleased = false
     private var truncationEntryWaiters: [CheckedContinuation<Void, Never>] = []
@@ -620,6 +660,15 @@ private actor RecordingAgentRunRepository: AgentRunPersisting {
 
     func save(_ snapshot: AgentRunSnapshot) async {
         snapshots.append(snapshot)
+        var pendingSnapshotWaiters: [(state: AgentRunState, continuation: CheckedContinuation<Void, Never>)] = []
+        for waiter in snapshotWaiters {
+            if waiter.state == snapshot.state {
+                waiter.continuation.resume()
+            } else {
+                pendingSnapshotWaiters.append(waiter)
+            }
+        }
+        snapshotWaiters = pendingSnapshotWaiters
         guard blocksFirstSave, snapshots.count == 1 else { return }
         saveEntered = true
         let entryWaiters = saveEntryWaiters
@@ -635,6 +684,11 @@ private actor RecordingAgentRunRepository: AgentRunPersisting {
         await withCheckedContinuation { saveEntryWaiters.append($0) }
     }
 
+    func waitUntilSnapshot(state: AgentRunState) async {
+        if snapshots.contains(where: { $0.state == state }) { return }
+        await withCheckedContinuation { snapshotWaiters.append((state, $0)) }
+    }
+
     func releaseSave() {
         saveReleased = true
         let waiters = saveReleaseWaiters

Frameworks/TermindApplication/Agent/AssistantConversationOperations.swift

@@ -131,7 +131,7 @@ public final class AssistantConversationOperations {
         replacement: AgentTranscriptMessage
     ) async throws -> ConversationModel.ComposedUserTurn? {
         try Task.checkCancellation()
-        guard
+        guard conversation.canApplyUserMessageEdit(messageID, replacement: replacement),
             pending.insert(conversation.id).inserted
         else { return nil }
         conversation.isHistoryChanging = true
@@ -143,6 +143,6 @@ public final class AssistantConversationOperations {
         try await repository.truncateConversation(
             conversation.id, fromUserMessageID: messageID, replacingWith: replacement)
         conversation.isHistoryChanging = false
-        return conversation.editUserMessage(messageID, replacement: replacement)
+        return conversation.applyPersistedUserMessageEdit(messageID, replacement: replacement)
     }
 }

Frameworks/TermindApplication/Agent/AssistantExecutionController.swift

@@ -723,12 +723,13 @@ public final class AssistantExecutionController {
         if state.recordsInterruption {
             let now = Date.now
             do {
-                try await agentRunRepository.save(
-                    AgentRunSnapshot(
-                        conversationID: conversation.id, state: .interrupted,
-                        permissionMode: capture.permissionMode,
-                        messages: conversation.transcript, startedAt: now, finishedAt: now))
+                let snapshot = AgentRunSnapshot(
+                    conversationID: conversation.id, state: .interrupted,
+                    permissionMode: capture.permissionMode,
+                    messages: conversation.transcript, startedAt: now, finishedAt: now)
+                try await agentRunRepository.save(snapshot)
                 guard preparations[conversation.id]?.id == capture.id else { return }
+                conversation.adoptStagedReplacementRun(snapshot)
                 conversation.handle(.runState(.interrupted))
             } catch {
                 guard preparations[conversation.id]?.id == capture.id else { return }
@@ -745,12 +746,13 @@ public final class AssistantExecutionController {
         if state.recordsInterruption {
             let now = Date.now
             do {
-                try await agentRunRepository.save(
-                    AgentRunSnapshot(
-                        conversationID: conversation.id, state: .interrupted,
-                        permissionMode: state.permissionMode,
-                        messages: conversation.transcript, startedAt: now, finishedAt: now))
+                let snapshot = AgentRunSnapshot(
+                    conversationID: conversation.id, state: .interrupted,
+                    permissionMode: state.permissionMode,
+                    messages: conversation.transcript, startedAt: now, finishedAt: now)
+                try await agentRunRepository.save(snapshot)
                 guard preparations[conversation.id]?.id == preparationID else { return }
+                conversation.adoptStagedReplacementRun(snapshot)
                 conversation.handle(.runState(.interrupted))
             } catch {
                 guard preparations[conversation.id]?.id == preparationID else { return }

Frameworks/TermindApplication/Agent/ConversationModel.swift

@@ -1067,20 +1067,79 @@ public final class ConversationModel: Identifiable {
         )
     }
 
+    public func canApplyUserMessageEdit(
+        _ id: UUID, replacement: AgentTranscriptMessage
+    ) -> Bool {
+        guard canEditUserMessage(id), replacement.role == .user,
+            !replacement.isCompaction, replacement.toolCalls.isEmpty,
+            replacement.toolCallID == nil, replacement.toolName == nil,
+            !replacement.isError, replacement.id != id,
+            !transcript.contains(where: { $0.id == replacement.id })
+        else { return false }
+        return true
+    }
+
     public func editUserMessage(
         _ id: UUID, replacement: AgentTranscriptMessage
     ) -> ComposedUserTurn? {
-        guard canEditUserMessage(id),
+        guard canApplyUserMessageEdit(id, replacement: replacement) else { return nil }
+        return applyPersistedUserMessageEdit(id, replacement: replacement)
+    }
+
+    func applyPersistedUserMessageEdit(
+        _ id: UUID, replacement: AgentTranscriptMessage
+    ) -> ComposedUserTurn? {
+        guard replacement.role == .user, !replacement.isCompaction,
+            replacement.toolCalls.isEmpty, replacement.toolCallID == nil,
+            replacement.toolName == nil, !replacement.isError,
+            replacement.id != id,
+            !transcript.contains(where: { $0.id == replacement.id }),
             let target = transcript.firstIndex(where: { $0.id == id })
         else { return nil }
         let original = transcript[target]
-        let removedMessageIDs = Set(transcript[target...].map(\.id))
-        let removedRunIDs = Set(removedMessageIDs.compactMap { messageRunIDs[$0] })
+        guard original.role == .user, !original.isCompaction else { return nil }
+        let prefix = Array(transcript[..<target])
+        let visibleTarget = messages.firstIndex(where: { $0.id == id })
+        let visiblePrefix = visibleTarget.map { Array(messages[..<$0]) } ?? []
+        let visibleRemoved = visibleTarget.map { Array(messages[$0...]) } ?? []
+        let removedMessageIDs = Set(transcript[target...].map(\.id)).union(visibleRemoved.map(\.id))
+        let keptMessageIDs = Set(prefix.map(\.id)).union(visiblePrefix.map(\.id))
+        var removedRunIDs = Set(removedMessageIDs.compactMap { messageRunIDs[$0] })
+        let prefixRunIDs = Set(prefix.compactMap { messageRunIDs[$0.id] })
+            .union(visiblePrefix.compactMap { messageRunIDs[$0.id] })
+        let targetRunID = messageRunIDs[id]
+        if let targetRunID, let targetSummary = runSummaries[targetRunID] {
+            removedRunIDs.formUnion(
+                runSummaries.values.compactMap { summary in
+                    if summary.startedAt != targetSummary.startedAt {
+                        return summary.startedAt > targetSummary.startedAt ? summary.id : nil
+                    }
+                    return summary.id.uuidString > targetSummary.id.uuidString ? summary.id : nil
+                })
+        }
+        let hasEarlierVisibleUser =
+            visibleTarget.map { index in
+                messages[..<index].contains { $0.role == .user }
+            } ?? false
+        let keptCallIDs = Set(prefix.flatMap { $0.toolCalls.map(\.id) })
+        let removedCallIDs = Set(transcript[target...].flatMap { $0.toolCalls.map(\.id) })
+        let whollyRemovedRunIDs = removedRunIDs.subtracting(prefixRunIDs)
+        let lastKeptRun = runSummaries.values.filter { !whollyRemovedRunIDs.contains($0.id) }
+            .max {
+                if $0.startedAt != $1.startedAt { return $0.startedAt < $1.startedAt }
+                return $0.id.uuidString < $1.id.uuidString
+            }
+        let stagingStartedAt =
+            if let lastKeptRun, replacement.createdAt <= lastKeptRun.startedAt {
+                lastKeptRun.startedAt.addingTimeInterval(0.001)
+            } else {
+                replacement.createdAt
+            }
         transcript.replaceSubrange(target..., with: [replacement])
         contextEstimateTracker.remove(removedMessageIDs)
         contextEstimateTracker.upsert(replacement)
 
-        if let visibleIndex = messages.firstIndex(where: { $0.id == id }) {
+        if let visibleIndex = visibleTarget {
             messages.replaceSubrange(
                 visibleIndex...,
                 with: [
@@ -1095,22 +1154,63 @@ public final class ConversationModel: Identifiable {
                 ])
         }
         toolCards.removeAll {
-            removedRunIDs.contains($0.invocation.runID)
-                || $0.invocation.createdAt >= original.createdAt
+            if whollyRemovedRunIDs.contains($0.invocation.runID) { return true }
+            if let source = $0.invocation.sourceMessageID {
+                if removedMessageIDs.contains(source) { return true }
+                if keptMessageIDs.contains(source) { return false }
+            }
+            if removedCallIDs.contains($0.invocation.call.id) { return true }
+            if keptCallIDs.contains($0.invocation.call.id) { return false }
+            let resultMessageID = AgentTranscriptMessage.toolResultID(
+                runID: $0.invocation.runID, callID: $0.invocation.call.id)
+            if messageRunIDs[resultMessageID] == $0.invocation.runID,
+                !removedMessageIDs.contains(resultMessageID)
+            {
+                return false
+            }
+            if $0.invocation.runID != targetRunID {
+                if prefixRunIDs.contains($0.invocation.runID) { return false }
+                if removedRunIDs.contains($0.invocation.runID) { return true }
+            }
+            // Legacy rows with no recoverable source, call, or result ownership
+            // belong to the edited suffix at the boundary as well as after it.
+            return $0.invocation.createdAt >= original.createdAt
         }
         for messageID in removedMessageIDs { messageRunIDs[messageID] = nil }
-        for runID in removedRunIDs { runSummaries[runID] = nil }
+        for runID in removedRunIDs {
+            guard prefixRunIDs.contains(runID) else {
+                runSummaries[runID] = nil
+                continue
+            }
+            let lastPrefixDate =
+                visiblePrefix.last(where: { messageRunIDs[$0.id] == runID })?.createdAt
+                ?? prefix.last(where: {
+                    messageRunIDs[$0.id] == runID && $0.role != .tool
+                })?.createdAt
+                ?? prefix.last(where: { messageRunIDs[$0.id] == runID })?.createdAt
+                ?? original.createdAt
+            let startedAt = runSummaries[runID]?.startedAt ?? lastPrefixDate
+            runSummaries[runID] = AgentRunSummary(
+                id: runID, state: .completed, startedAt: startedAt,
+                finishedAt: lastPrefixDate)
+        }
+        messageRunIDs[replacement.id] = replacement.id
+        runSummaries[replacement.id] = AgentRunSummary(
+            id: replacement.id, state: .completed,
+            startedAt: stagingStartedAt, finishedAt: stagingStartedAt)
         plans.removeAll { $0.createdAt >= original.createdAt }
-        compactions.removeAll { $0.createdAt >= original.createdAt }
+        compactionRecords.removeAll { record in
+            guard let boundary = record.compactedThroughMessageID else { return false }
+            return removedMessageIDs.contains(boundary)
+        }
+        compactions = compactionRecords.map(AgentCompactionMark.init)
         tasks.removeAll()
         usage = nil
         contextEstimateTracker.setMeasured(nil)
         publishContextEstimateImmediately()
         reportedRunState = .idle
         updatedAt = replacement.createdAt
-        if !transcript[..<target].contains(where: { $0.role == .user && !$0.isCompaction }) {
-            title = nil
-        }
+        if !hasEarlierVisibleUser { title = nil }
         return ComposedUserTurn(
             id: replacement.id, text: replacement.text,
             authoredText: replacement.authoredText ?? "", images: replacement.images,
@@ -1941,18 +2041,57 @@ public final class ConversationModel: Identifiable {
         guard snapshot.conversationID == id, snapshot.state == .failed,
             let failure = snapshot.failure
         else { return }
-        if let prompt = snapshot.messages.last(where: { $0.role == .user }) {
-            messageRunIDs[prompt.id] = snapshot.id
-        }
-        runSummaries[snapshot.id] = AgentRunSummary(
+        var summary = AgentRunSummary(
             id: snapshot.id, state: .failed, startedAt: snapshot.startedAt,
             finishedAt: snapshot.finishedAt, failure: failure
         )
+        if let prompt = snapshot.messages.last(where: { $0.role == .user }),
+            messageRunIDs[prompt.id] == nil
+                || adoptStagedReplacementRun(promptID: prompt.id, summary: &summary)
+        {
+            messageRunIDs[prompt.id] = snapshot.id
+        }
+        runSummaries[snapshot.id] = summary
         reportedRunState = .failed
         updatedAt = snapshot.finishedAt ?? .now
         runDidFinish?()
     }
 
+    public func adoptStagedReplacementRun(_ snapshot: AgentRunSnapshot) {
+        guard snapshot.conversationID == id,
+            let prompt = snapshot.messages.last(where: { $0.role == .user })
+        else { return }
+        var summary = AgentRunSummary(
+            id: snapshot.id, state: snapshot.state,
+            startedAt: snapshot.startedAt, finishedAt: snapshot.finishedAt,
+            failure: snapshot.failure)
+        adoptStagedReplacementRun(
+            promptID: prompt.id,
+            summary: &summary)
+    }
+
+    @discardableResult
+    private func adoptStagedReplacementRun(
+        promptID: UUID, summary: inout AgentRunSummary
+    ) -> Bool {
+        guard messageRunIDs[promptID] == promptID,
+            let staging = runSummaries[promptID], staging.state == .completed,
+            staging.failure == nil,
+            messageRunIDs.filter({ $0.value == promptID }).count == 1,
+            !toolCards.contains(where: { $0.invocation.runID == promptID })
+        else { return false }
+        if summary.startedAt < staging.startedAt {
+            summary.startedAt = staging.startedAt
+        }
+        if let finishedAt = summary.finishedAt, finishedAt < summary.startedAt {
+            summary.finishedAt = summary.startedAt
+        }
+        runSummaries[promptID] = nil
+        messageRunIDs[promptID] = summary.id
+        runSummaries[summary.id] = summary
+        return true
+    }
+
     /// Retains the ordered event consumer until the cancellation snapshot is durable.
     public func stopAndWait() async {
         await cancelPendingAgentPreparationAndWait()
@@ -2484,14 +2623,23 @@ public final class ConversationModel: Identifiable {
         // them — but they are listed rather than swept into a `default` so
         // adding an event stays a compile error here.
         case .runStarted(let run):
-            activeRunID = run.id
-            runSummaries[run.id] = run
+            var adoptedRun = run
+            activeRunID = adoptedRun.id
             if let prompt = transcript.last, prompt.role == .user,
                 messageRunIDs[prompt.id] == nil
+                    || adoptStagedReplacementRun(promptID: prompt.id, summary: &adoptedRun)
             {
-                messageRunIDs[prompt.id] = run.id
+                messageRunIDs[prompt.id] = adoptedRun.id
             }
+            runSummaries[adoptedRun.id] = adoptedRun
         case .runFinished(let run):
+            var run = run
+            if let current = runSummaries[run.id], run.startedAt < current.startedAt {
+                run.startedAt = current.startedAt
+            }
+            if let finishedAt = run.finishedAt, finishedAt < run.startedAt {
+                run.finishedAt = run.startedAt
+            }
             publishAllStagedTools()
             settleStreamingPresentation(for: run.id)
             runSummaries[run.id] = run
@@ -2645,6 +2793,7 @@ public final class ConversationModel: Identifiable {
                 role: .tool, text: result.content, toolCallID: result.callID,
                 toolName: invocation.call.name, isError: result.isError, modelText: result.modelContent
             )
+            messageRunIDs[message.id] = invocation.runID
             transcript.append(message)
             contextEstimateTracker.upsert(message)
             publishContextEstimateImmediately()

Frameworks/TermindPersistence/SwiftDataAgentRunRepository.swift

@@ -19,60 +19,102 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
             deletedRunIDs.insert(snapshot.id)
             return
         }
-        let runID = snapshot.id
-        let runDescriptor = FetchDescriptor<PersistentAgentRun>(predicate: #Predicate { $0.id == runID })
-        if let row = try context.fetch(runDescriptor).first {
-            row.update(snapshot)
-        } else {
-            context.insert(PersistentAgentRun(snapshot))
-        }
+        guard !deletedRunIDs.contains(snapshot.id) else { return }
+        let previousDeletedRunIDs = deletedRunIDs
+        do {
+            let runID = snapshot.id
+            let runDescriptor = FetchDescriptor<PersistentAgentRun>(predicate: #Predicate { $0.id == runID })
+            let runRow: PersistentAgentRun
+            if let row = try context.fetch(runDescriptor).first {
+                let persistedStartedAt = row.startedAt
+                row.update(snapshot)
+                if row.startedAt < persistedStartedAt {
+                    row.startedAt = persistedStartedAt
+                }
+                if let finishedAt = row.finishedAt, finishedAt < row.startedAt {
+                    row.finishedAt = row.startedAt
+                }
+                runRow = row
+            } else {
+                let row = PersistentAgentRun(snapshot)
+                context.insert(row)
+                runRow = row
+            }
 
-        let conversationID = snapshot.conversationID
-        let conversationDescriptor = FetchDescriptor<PersistentAgentConversation>(
-            predicate: #Predicate { $0.id == conversationID }
-        )
-        // No title seeding here. `title` means "the generated summary", and
-        // empty means there is not one yet — which is what lets the UI decide to
-        // fall back to the opening line rather than having that line frozen into
-        // storage as though it were a name someone chose.
-        if let conversation = try context.fetch(conversationDescriptor).first {
-            conversation.updatedAt = .now
-        } else {
-            context.insert(
-                PersistentAgentConversation(
-                    id: conversationID, title: "", createdAt: snapshot.startedAt
-                ))
-        }
+            let conversationID = snapshot.conversationID
+            let conversationDescriptor = FetchDescriptor<PersistentAgentConversation>(
+                predicate: #Predicate { $0.id == conversationID }
+            )
+            // No title seeding here. `title` means "the generated summary", and
+            // empty means there is not one yet — which is what lets the UI decide to
+            // fall back to the opening line rather than having that line frozen into
+            // storage as though it were a name someone chose.
+            if let conversation = try context.fetch(conversationDescriptor).first {
+                conversation.updatedAt = .now
+            } else {
+                context.insert(
+                    PersistentAgentConversation(
+                        id: conversationID, title: "", createdAt: snapshot.startedAt
+                    ))
+            }
 
-        // Matched on message id, not on this run's id. A conversation's earlier
-        // turns are replayed into every later run, so filtering by `runID` made
-        // each save re-insert the whole history — colliding with the unique
-        // constraint and rewriting rows that were already correct.
-        let durableMessages = snapshot.messages.filter { !$0.isCompaction }
-        let ids = Set(durableMessages.map(\.id))
-        let stored = try context.fetch(
-            FetchDescriptor<PersistentAgentMessage>(
-                predicate: #Predicate { ids.contains($0.id) }
-            ))
-        let storedIDs = Set(stored.map(\.id))
-        var inserted = false
-        for (sequence, message) in durableMessages.enumerated()
-        where !storedIDs.contains(message.id) {
-            context.insert(
-                PersistentAgentMessage(
-                    message, runID: runID, sequence: sequence
+            // Matched on message id, not on this run's id. A conversation's earlier
+            // turns are replayed into every later run, so filtering by `runID` made
+            // each save re-insert the whole history — colliding with the unique
+            // constraint and rewriting rows that were already correct.
+            let durableMessages = snapshot.messages.filter { !$0.isCompaction }
+            let ids = Set(durableMessages.map(\.id))
+            let stored = try context.fetch(
+                FetchDescriptor<PersistentAgentMessage>(
+                    predicate: #Predicate { ids.contains($0.id) }
                 ))
-            insertContextAttachments(for: message)
-            inserted = true
+            var storedByID = Dictionary(uniqueKeysWithValues: stored.map { ($0.id, $0) })
+            if let prompt = durableMessages.last(where: { $0.role == .user }),
+                let stagedMessage = storedByID[prompt.id],
+                stagedMessage.runID == prompt.id, prompt.id != runID,
+                try isStagingReplacementRun(prompt.id, message: stagedMessage, conversationID: conversationID)
+            {
+                stagedMessage.runID = runID
+                stagedMessage.sequence = durableMessages.firstIndex(where: { $0.id == prompt.id }) ?? 0
+                let stagingID = prompt.id
+                let stagingDescriptor = FetchDescriptor<PersistentAgentRun>(
+                    predicate: #Predicate { $0.id == stagingID }
+                )
+                if let stagingRun = try context.fetch(stagingDescriptor).first {
+                    if runRow.startedAt < stagingRun.startedAt {
+                        runRow.startedAt = stagingRun.startedAt
+                    }
+                    if let finishedAt = runRow.finishedAt, finishedAt < runRow.startedAt {
+                        runRow.finishedAt = runRow.startedAt
+                    }
+                    context.delete(stagingRun)
+                }
+                deletedRunIDs.insert(stagingID)
+            }
+            let storedIDs = Set(storedByID.keys)
+            var inserted = false
+            for (sequence, message) in durableMessages.enumerated()
+            where !storedIDs.contains(message.id) {
+                let row = PersistentAgentMessage(message, runID: runID, sequence: sequence)
+                context.insert(row)
+                storedByID[message.id] = row
+                insertContextAttachments(for: message)
+                inserted = true
+            }
+            // A user row may already exist (editing or interrupted preparation) before
+            // the run freezes its context. Add its backing once without rewriting history.
+            let replayIDs = Set(
+                try context.fetch(FetchDescriptor<PersistentAgentMessageReplay>()).map(\.messageID))
+            for message in durableMessages where !replayIDs.contains(message.id) {
+                insertReplay(for: message)
+            }
+            guard inserted || context.hasChanges else { return }
+            try context.save()
+        } catch {
+            context.rollback()
+            deletedRunIDs = previousDeletedRunIDs
+            throw error
         }
-        // A user row may already exist (editing or interrupted preparation) before
-        // the run freezes its context. Add its backing once without rewriting history.
-        let replayIDs = Set(try context.fetch(FetchDescriptor<PersistentAgentMessageReplay>()).map(\.messageID))
-        for message in durableMessages where !replayIDs.contains(message.id) {
-            insertReplay(for: message)
-        }
-        guard inserted || context.hasChanges else { return }
-        try context.save()
     }
 
     @MainActor
@@ -174,7 +216,7 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
         let reviews = try context.fetch(FetchDescriptor<PersistentAgentReviewRecord>())
         return conversations.map { conversation in
             let matchingRuns = runs.filter { $0.conversationID == conversation.id }
-                .sorted { $0.startedAt < $1.startedAt }
+                .sorted(by: Self.runPrecedes)
             let runIDs = Set(matchingRuns.map(\.id))
             // Ordered by run first, then by time within it. Sorting the whole
             // conversation on `createdAt` alone let two runs interleave whenever
@@ -184,24 +226,7 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
                 uniqueKeysWithValues: matchingRuns.enumerated().map { ($0.element.id, $0.offset) }
             )
             let transcript = messages.filter { runIDs.contains($0.runID) }
-                .sorted { lhs, rhs in
-                    let left = runOrder[lhs.runID] ?? 0
-                    let right = runOrder[rhs.runID] ?? 0
-                    if left != right { return left < right }
-                    if lhs.sequence != rhs.sequence { return lhs.sequence < rhs.sequence }
-                    if lhs.createdAt != rhs.createdAt { return lhs.createdAt < rhs.createdAt }
-                    // Transcript grammar gives equal-position, equal-time rows
-                    // a stable and useful final ordering.
-                    let rank = { (role: String) -> Int in
-                        switch AgentTranscriptRole(rawValue: role) {
-                        case .user: 0
-                        case .assistant: 1
-                        case .tool: 2
-                        case nil: 3
-                        }
-                    }
-                    return rank(lhs.role) < rank(rhs.role)
-                }
+                .sorted { lhs, rhs in Self.messagePrecedes(lhs, rhs, runOrder: runOrder) }
                 .compactMap { row -> AgentTranscriptMessage? in
                     guard let role = AgentTranscriptRole(rawValue: row.role) else { return nil }
                     let calls =
@@ -230,22 +255,48 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
                         createdAt: row.createdAt
                     )
                 }
-            let messageRunIDs = Dictionary(
+            var messageRunIDs = Dictionary(
                 uniqueKeysWithValues: messages.filter { runIDs.contains($0.runID) }
                     .map { ($0.id, $0.runID) }
             )
             let state =
                 matchingRuns.last.flatMap { AgentRunState(rawValue: $0.state) }
                 ?? .completed
+            var sourceMessageIDsByRun: [UUID: [String: UUID]] = [:]
+            for message in transcript where message.role == .assistant {
+                guard let runID = messageRunIDs[message.id] else { continue }
+                for call in message.toolCalls {
+                    sourceMessageIDsByRun[runID, default: [:]][call.id] = message.id
+                }
+            }
             let cards = invocations.filter { runIDs.contains($0.runID) }
                 .sorted { $0.createdAt < $1.createdAt }
-                .compactMap { storedToolCard($0, results: results, reviews: reviews) }
+                .compactMap { row -> StoredAgentToolCard? in
+                    guard var card = storedToolCard(row, results: results, reviews: reviews) else {
+                        return nil
+                    }
+                    // Older invocation rows predate source-message ownership. The
+                    // raw transcript still has the exact run-local call grammar,
+                    // so project that identity on restore instead of guessing by time.
+                    if card.invocation.sourceMessageID == nil {
+                        card.invocation.sourceMessageID = sourceMessageIDsByRun[row.runID]?[row.callID]
+                    }
+                    return card
+                }
+            let matchingResults = results.filter { runIDs.contains($0.runID) }
+            let recoveredTranscript = Self.recovering(
+                transcript, results: matchingResults,
+                invocations: invocations.filter { runIDs.contains($0.runID) })
+            let recoveredMessageIDs = Set(recoveredTranscript.map(\.id))
+            for result in matchingResults {
+                let messageID = AgentTranscriptMessage.toolResultID(
+                    runID: result.runID, callID: result.callID)
+                if recoveredMessageIDs.contains(messageID) {
+                    messageRunIDs[messageID] = result.runID
+                }
+            }
             let projection = AgentCompactionProjection.restore(
-                messages: Self.recovering(
-                    transcript,
-                    results: results.filter { runIDs.contains($0.runID) },
-                    invocations: invocations.filter { runIDs.contains($0.runID) }
-                ),
+                messages: recoveredTranscript,
                 records: compactionRows.filter { $0.conversationID == conversation.id }
                     .map(\.value)
             )
@@ -355,28 +406,65 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
         replacingWith replacement: AgentTranscriptMessage
     ) async throws {
         try Task.checkCancellation()
+        guard replacement.role == .user, !replacement.isCompaction,
+            replacement.toolCalls.isEmpty, replacement.toolCallID == nil,
+            replacement.toolName == nil, !replacement.isError,
+            replacement.id != messageID
+        else { throw AgentConversationPersistenceError.messageNotFound }
         let allRuns = try context.fetch(FetchDescriptor<PersistentAgentRun>())
         let runs = allRuns.filter { $0.conversationID == id }
-            .sorted { $0.startedAt < $1.startedAt }
+            .sorted(by: Self.runPrecedes)
         let allMessages = try context.fetch(FetchDescriptor<PersistentAgentMessage>())
+        let runIDs = Set(runs.map(\.id))
+        let runOrder = Dictionary(
+            uniqueKeysWithValues: runs.enumerated().map { ($0.element.id, $0.offset) })
+        let conversationMessages = allMessages.filter { runIDs.contains($0.runID) }
+            .sorted { Self.messagePrecedes($0, $1, runOrder: runOrder) }
         guard
-            let target = allMessages.first(where: { message in
-                message.id == messageID && runs.contains { $0.id == message.runID }
-            }),
-            let targetIndex = runs.firstIndex(where: { $0.id == target.runID })
+            !allMessages.contains(where: { $0.id == replacement.id }),
+            !allRuns.contains(where: { $0.id == replacement.id }),
+            let boundary = conversationMessages.firstIndex(where: { $0.id == messageID }),
+            conversationMessages[boundary].role == AgentTranscriptRole.user.rawValue,
+            !conversationMessages[boundary].isCompaction,
+            let targetRunIndex = runs.firstIndex(where: { $0.id == conversationMessages[boundary].runID })
         else { throw AgentConversationPersistenceError.messageNotFound }
 
-        let removedRuns = Array(runs[targetIndex...])
-        let removedRunIDs = Set(removedRuns.map(\.id))
-        let removedMessageIDs = Set(allMessages.filter { removedRunIDs.contains($0.runID) }.map(\.id))
-        deletedRunIDs.formUnion(removedRunIDs)
-        let earlierRunIDs = Set(runs[..<targetIndex].map(\.id))
-        let removesOpeningTurn = !allMessages.contains {
-            earlierRunIDs.contains($0.runID) && $0.role == AgentTranscriptRole.user.rawValue
-                && !$0.isCompaction
+        let target = conversationMessages[boundary]
+        let keptMessages = Array(conversationMessages[..<boundary])
+        let removedMessages = Array(conversationMessages[boundary...])
+        let removedMessageIDs = Set(removedMessages.map(\.id))
+        let laterRuns = Array(runs.dropFirst(targetRunIndex + 1))
+        let laterRunIDs = Set(laterRuns.map(\.id))
+        let keptTargetMessages = keptMessages.filter { $0.runID == target.runID }
+        let removesTargetRun = keptTargetMessages.isEmpty
+        let runsToDelete = removesTargetRun ? [runs[targetRunIndex]] + laterRuns : laterRuns
+        let runIDsToDelete = Set(runsToDelete.map(\.id))
+        let retiredRunIDs = laterRunIDs.union([target.runID])
+        let previousDeletedRunIDs = deletedRunIDs
+        deletedRunIDs.formUnion(retiredRunIDs)
+        let removesOpeningTurn = !keptMessages.contains {
+            $0.role == AgentTranscriptRole.user.rawValue && !$0.isCompaction
         }
         do {
-            try deleteGraph(runIDs: removedRunIDs, runs: removedRuns)
+            try deleteGraph(runIDs: runIDsToDelete, runs: runsToDelete)
+            if !removesTargetRun,
+                let lastKept = keptTargetMessages.last(where: {
+                    $0.role != AgentTranscriptRole.tool.rawValue
+                }) ?? keptTargetMessages.last
+            {
+                try deleteTargetRunTools(
+                    runID: target.runID, keptMessages: keptTargetMessages,
+                    removedMessages: removedMessages.filter { $0.runID == target.runID },
+                    boundary: target.createdAt)
+                try deleteMessageBacking(removedMessages.filter { $0.runID == target.runID })
+                let targetRun = runs[targetRunIndex]
+                targetRun.state = AgentRunState.completed.rawValue
+                // Tool results are omitted from the visible transcript after
+                // restore. End the neutral prefix summary at the last retained
+                // non-tool message so live editing and reopening agree.
+                targetRun.finishedAt = lastKept.createdAt
+                targetRun.failure = nil
+            }
             for row in try context.fetch(FetchDescriptor<PersistentAgentCompaction>())
             where row.conversationID == id
                 && row.compactedThroughMessageID.map(removedMessageIDs.contains) == true
@@ -384,10 +472,17 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
                 context.delete(row)
             }
 
+            let lastKeptRun = removesTargetRun ? runs[..<targetRunIndex].last : runs[targetRunIndex]
+            let replacementStartedAt =
+                if let lastKeptRun, replacement.createdAt <= lastKeptRun.startedAt {
+                    lastKeptRun.startedAt.addingTimeInterval(0.001)
+                } else {
+                    replacement.createdAt
+                }
             let replacementRun = AgentRunSnapshot(
-                conversationID: id, state: .completed,
+                id: replacement.id, conversationID: id, state: .completed,
                 permissionMode: .askForApproval, messages: [replacement],
-                startedAt: replacement.createdAt, finishedAt: replacement.createdAt
+                startedAt: replacementStartedAt, finishedAt: replacementStartedAt
             )
             context.insert(PersistentAgentRun(replacementRun))
             context.insert(
@@ -407,7 +502,7 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
             try context.save()
         } catch {
             context.rollback()
-            deletedRunIDs.subtract(removedRunIDs)
+            deletedRunIDs = previousDeletedRunIDs
             throw error
         }
     }
@@ -535,6 +630,92 @@ public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
         for row in runs { context.delete(row) }
     }
 
+    private func deleteMessageBacking(_ messages: [PersistentAgentMessage]) throws {
+        let messageIDs = Set(messages.map(\.id))
+        for row in try context.fetch(FetchDescriptor<PersistentAgentContextAttachments>())
+        where messageIDs.contains(row.messageID) { context.delete(row) }
+        for row in try context.fetch(FetchDescriptor<PersistentAgentMessageReplay>())
+        where messageIDs.contains(row.messageID) { context.delete(row) }
+        for row in messages { context.delete(row) }
+    }
+
+    private func deleteTargetRunTools(
+        runID: UUID, keptMessages: [PersistentAgentMessage],
+        removedMessages: [PersistentAgentMessage], boundary: Date
+    ) throws {
+        let removedMessageIDs = Set(removedMessages.map(\.id))
+        let keptMessageIDs = Set(keptMessages.map(\.id))
+        let keptCallIDs = Set(keptMessages.flatMap(Self.toolCallIDs))
+        let removedCallIDs = Set(removedMessages.flatMap(Self.toolCallIDs))
+        let invocations = try context.fetch(FetchDescriptor<PersistentAgentToolInvocation>())
+        let removedInvocations = invocations.filter { invocation in
+            guard invocation.runID == runID else { return false }
+            if let source = invocation.sourceMessageID {
+                if removedMessageIDs.contains(source) { return true }
+                if keptMessageIDs.contains(source) { return false }
+            }
+            if removedCallIDs.contains(invocation.callID) { return true }
+            if keptCallIDs.contains(invocation.callID) { return false }
+            // Legacy rows with no recoverable source or call ownership belong
+            // to the edited suffix at the boundary as well as after it.
+            return invocation.createdAt >= boundary
+        }
+        let invocationIDs = Set(removedInvocations.map(\.id))
+        for row in try context.fetch(FetchDescriptor<PersistentAgentToolResult>())
+        where row.runID == runID && invocationIDs.contains(row.invocationID) { context.delete(row) }
+        for row in try context.fetch(FetchDescriptor<PersistentAgentReviewRecord>())
+        where row.runID == runID && invocationIDs.contains(row.invocationID) { context.delete(row) }
+        for row in removedInvocations { context.delete(row) }
+    }
+
+    private func isStagingReplacementRun(
+        _ id: UUID, message: PersistentAgentMessage, conversationID: UUID
+    ) throws -> Bool {
+        let runDescriptor = FetchDescriptor<PersistentAgentRun>(predicate: #Predicate { $0.id == id })
+        guard let run = try context.fetch(runDescriptor).first,
+            run.conversationID == conversationID, run.state == AgentRunState.completed.rawValue,
+            run.failure == nil, message.id == id, message.runID == id,
+            message.role == AgentTranscriptRole.user.rawValue, !message.isCompaction
+        else { return false }
+        let messages = try context.fetch(FetchDescriptor<PersistentAgentMessage>())
+        guard messages.filter({ $0.runID == id }).count == 1 else { return false }
+        let invocations = try context.fetch(FetchDescriptor<PersistentAgentToolInvocation>())
+        let results = try context.fetch(FetchDescriptor<PersistentAgentToolResult>())
+        let reviews = try context.fetch(FetchDescriptor<PersistentAgentReviewRecord>())
+        return !invocations.contains { $0.runID == id }
+            && !results.contains { $0.runID == id }
+            && !reviews.contains { $0.runID == id }
+    }
+
+    private static func runPrecedes(_ lhs: PersistentAgentRun, _ rhs: PersistentAgentRun) -> Bool {
+        if lhs.startedAt != rhs.startedAt { return lhs.startedAt < rhs.startedAt }
+        return lhs.id.uuidString < rhs.id.uuidString
+    }
+
+    private static func messagePrecedes(
+        _ lhs: PersistentAgentMessage, _ rhs: PersistentAgentMessage, runOrder: [UUID: Int]
+    ) -> Bool {
+        let left = runOrder[lhs.runID] ?? 0
+        let right = runOrder[rhs.runID] ?? 0
+        if left != right { return left < right }
+        if lhs.sequence != rhs.sequence { return lhs.sequence < rhs.sequence }
+        if lhs.createdAt != rhs.createdAt { return lhs.createdAt < rhs.createdAt }
+        let rank: (String) -> Int = { role in
+            switch AgentTranscriptRole(rawValue: role) {
+            case .user: 0
+            case .assistant: 1
+            case .tool: 2
+            case nil: 3
+            }
+        }
+        if rank(lhs.role) != rank(rhs.role) { return rank(lhs.role) < rank(rhs.role) }
+        return lhs.id.uuidString < rhs.id.uuidString
+    }
+
+    private static func toolCallIDs(_ message: PersistentAgentMessage) -> [String] {
+        ((try? JSONDecoder().decode([AgentToolCall].self, from: message.toolCallsData)) ?? []).map(\.id)
+    }
+
     private func insertReplay(for message: AgentTranscriptMessage) {
         guard message.contextSnapshot != nil || message.modelText != nil else { return }
         context.insert(PersistentAgentMessageReplay(message))

TFITests/MobileAssistantTests.swift

@@ -5,6 +5,7 @@ import SwiftUI
 import TermindApplication
 import TermindDomain
 import TermindInfrastructure
+import TermindPersistence
 import TermindPresentation
 import TermindRemoteInfrastructure
 import Testing
@@ -150,6 +151,53 @@ struct MobileAssistantTests {
         }
     }
 
+    @Test func swiftDataRestoreAndEditPreserveASameRunPrefixOnMobile() async throws {
+        let container = try TermindPersistence.makeContainer(inMemory: true)
+        let repository = SwiftDataAgentRunRepository(context: container.mainContext)
+        let instant = Date(timeIntervalSinceReferenceDate: 100)
+        let runID = try #require(UUID(uuidString: "FFFFFFFF-FFFF-FFFF-FFFF-FFFFFFFFFFFF"))
+        let replacementID = try #require(UUID(uuidString: "00000000-0000-0000-0000-000000000001"))
+        let opening = AgentTranscriptMessage(
+            role: .user, text: "Opening", authoredText: "Opening", createdAt: instant)
+        let prefix = AgentTranscriptMessage(
+            role: .assistant, text: "Prefix", createdAt: instant)
+        let steering = AgentTranscriptMessage(
+            role: .user, text: "Steer", authoredText: "Steer", createdAt: instant)
+        let suffix = AgentTranscriptMessage(
+            role: .assistant, text: "Suffix", createdAt: instant)
+        try await repository.save(
+            AgentRunSnapshot(
+                id: runID, conversationID: UUID(), state: .failed,
+                permissionMode: .askForApproval,
+                messages: [opening, prefix, steering, suffix],
+                startedAt: instant, finishedAt: instant, failure: "old failure"))
+        let store = AssistantConversationStore()
+        store.restore(try await repository.loadConversations())
+        let conversation = try #require(store.conversations.first)
+        let replacement = AgentTranscriptMessage(
+            id: replacementID, role: .user, text: "Replacement", authoredText: "Replacement",
+            createdAt: instant)
+
+        let operations = AssistantConversationOperations(repository: repository)
+        #expect(
+            try await operations.edit(
+                conversation, messageID: steering.id, replacement: replacement) != nil)
+
+        #expect(conversation.messages.map(\.text) == ["Opening", "Prefix", "Replacement"])
+        #expect(conversation.messageRunIDs[opening.id] == runID)
+        #expect(conversation.messageRunIDs[prefix.id] == runID)
+        #expect(conversation.messageRunIDs[replacement.id] == replacement.id)
+        #expect(try #require(conversation.runSummaries[replacement.id]).startedAt > instant)
+        let reopened = try #require(try await repository.loadConversations().first)
+        #expect(reopened.messages.map(\.text) == ["Opening", "Prefix", "Replacement"])
+        #expect(reopened.messageRunIDs[opening.id] == runID)
+        #expect(reopened.messageRunIDs[prefix.id] == runID)
+        #expect(reopened.messageRunIDs[replacement.id] == replacement.id)
+        #expect(reopened.runs.map(\.id) == [runID, replacement.id])
+        #expect(reopened.runs.first { $0.id == runID }?.state == .completed)
+        #expect(reopened.runs.first { $0.id == runID }?.failure == nil)
+    }
+
     @Test func changingConversationsPreservesTheirDraftsAndArchiveNavigation() async throws {
         let app = try MobileAppModel(inMemory: true)
         let model = app.assistant.makeSceneModel()

TFMTests/AgentPersistenceTests.swift

@@ -868,6 +868,266 @@ struct AgentPersistenceTests {
         #expect(restored.messages.map(\.text) == ["One", "Answer one", "Edited"])
     }
 
+    @Test
+    func sameRunTruncationPreservesThePrefixAndAtomicallyAdoptsTheReplacement() async throws {
+        let configuration = localModelConfiguration(
+            "AgentSameRunTruncationTests", schema: HostPersistence.schema,
+            isStoredInMemoryOnly: true
+        )
+        let container = try ModelContainer(
+            for: HostPersistence.schema, configurations: [configuration]
+        )
+        let repository = SwiftDataAgentRunRepository(context: container.mainContext)
+        let conversationID = UUID()
+        let oldRunID = try #require(UUID(uuidString: "FFFFFFFF-FFFF-FFFF-FFFF-FFFFFFFFFFFF"))
+        let replacementID = try #require(UUID(uuidString: "00000000-0000-0000-0000-000000000001"))
+        let realRunID = try #require(UUID(uuidString: "00000000-0000-0000-0000-000000000002"))
+        let instant = Date(timeIntervalSinceReferenceDate: 100)
+        let prefixAttachment = AgentContextAttachment(
+            kind: .file, label: "prefix.txt", contentTypeIdentifier: "public.plain-text",
+            suggestedFilename: "prefix.txt", data: Data("prefix".utf8))
+        let removedAttachment = AgentContextAttachment(
+            kind: .file, label: "removed.txt", contentTypeIdentifier: "public.plain-text",
+            suggestedFilename: "removed.txt", data: Data("removed".utf8))
+        let replacementAttachment = AgentContextAttachment(
+            kind: .file, label: "replacement.txt", contentTypeIdentifier: "public.plain-text",
+            suggestedFilename: "replacement.txt", data: Data("replacement".utf8))
+        let opening = AgentTranscriptMessage(
+            role: .user, text: "Opening", authoredText: "Opening",
+            contextAttachments: [prefixAttachment],
+            contextSnapshot: .init(sessionContext: "prefix context"), createdAt: instant)
+        let prefixCall = AgentToolCall(
+            id: "prefix-call", name: "inspect", arguments: .object([:]))
+        let prefixAssistant = AgentTranscriptMessage(
+            role: .assistant, text: "Prefix", toolCalls: [prefixCall], createdAt: instant)
+        let prefixResult = AgentTranscriptMessage(
+            role: .tool, text: "prefix result", toolCallID: prefixCall.id,
+            toolName: prefixCall.name, modelText: "prefix preview", createdAt: instant)
+        let steering = AgentTranscriptMessage(
+            role: .user, text: "Steer", authoredText: "Steer",
+            contextAttachments: [removedAttachment],
+            contextSnapshot: .init(sessionContext: "removed context"), createdAt: instant)
+        let suffixCall = AgentToolCall(
+            id: "suffix-call", name: "inspect", arguments: .object([:]))
+        let suffixAssistant = AgentTranscriptMessage(
+            role: .assistant, text: "Suffix", toolCalls: [suffixCall], createdAt: instant)
+        let suffixResult = AgentTranscriptMessage(
+            role: .tool, text: "suffix result", toolCallID: suffixCall.id,
+            toolName: suffixCall.name, modelText: "suffix preview", createdAt: instant)
+        let oldMessages = [
+            opening, prefixAssistant, prefixResult, steering, suffixAssistant, suffixResult,
+        ]
+        let replacement = AgentTranscriptMessage(
+            id: replacementID, role: .user, text: "Replacement", authoredText: "Replacement",
+            contextAttachments: [replacementAttachment],
+            contextSnapshot: .init(sessionContext: "replacement context"), createdAt: instant)
+        let oldSummary = AgentRunSummary(
+            id: oldRunID, state: .failed, startedAt: instant,
+            finishedAt: instant, failure: "old failure")
+        let live = ConversationModel(
+            messages: [
+                ChatMessage(id: opening.id, role: .user, text: opening.text, createdAt: instant),
+                ChatMessage(
+                    id: prefixAssistant.id, role: .assistant, text: prefixAssistant.text,
+                    createdAt: instant),
+                ChatMessage(id: steering.id, role: .user, text: steering.text, createdAt: instant),
+                ChatMessage(
+                    id: suffixAssistant.id, role: .assistant, text: suffixAssistant.text,
+                    createdAt: instant),
+            ], transcript: oldMessages, runState: .failed,
+            messageRunIDs: Dictionary(uniqueKeysWithValues: oldMessages.map { ($0.id, oldRunID) }),
+            runs: [oldSummary])
+        #expect(live.editUserMessage(steering.id, replacement: replacement) != nil)
+        let livePrefixFinishedAt = try #require(live.runSummaries[oldRunID]?.finishedAt)
+
+        let oldSnapshot = AgentRunSnapshot(
+            id: oldRunID, conversationID: conversationID, state: .failed,
+            permissionMode: .askForApproval, messages: oldMessages,
+            startedAt: instant, finishedAt: instant, failure: "old failure")
+        try await repository.save(oldSnapshot)
+        try await repository.setTitle("Preserved title", conversationID: conversationID)
+
+        let descriptor = AgentToolDescriptor(
+            name: "inspect", summary: "Inspect", inputSchema: .object([:]),
+            target: .local, approvalPolicy: .approve)
+        let prefixInvocation = AgentToolInvocation(
+            runID: oldRunID, call: prefixCall, sourceMessageID: nil,
+            createdAt: instant.addingTimeInterval(1))
+        let suffixInvocation = AgentToolInvocation(
+            runID: oldRunID, call: suffixCall,
+            sourceMessageID: suffixAssistant.id, createdAt: instant)
+        let review = GuardianDecision(
+            verdict: .approve, risk: .low, userAuthorization: .high, reason: "Allowed")
+        for invocation in [prefixInvocation, suffixInvocation] {
+            try await repository.record(.toolProposed(invocation, descriptor), runID: oldRunID)
+            try await repository.record(
+                .toolFinished(
+                    invocation,
+                    AgentToolResult(callID: invocation.call.id, content: "recorded result")),
+                runID: oldRunID)
+            let request = AgentApprovalRequest(
+                invocation: invocation, descriptor: descriptor, userIntent: "Inspect")
+            try await repository.record(.reviewFinished(request, review), runID: oldRunID)
+        }
+        // Provider-native tools can finish durably before their synthetic assistant
+        // source message reaches a snapshot. Its unknown source must fall through
+        // to the conservative timestamp rule when this steering suffix is edited.
+        let orphanInvocation = AgentToolInvocation(
+            runID: oldRunID,
+            call: AgentToolCall(id: "native-call", name: "web_search", arguments: .object([:])),
+            sourceMessageID: UUID(), createdAt: instant)
+        try await repository.record(.toolProposed(orphanInvocation, descriptor), runID: oldRunID)
+        try await repository.record(
+            .toolFinished(
+                orphanInvocation,
+                AgentToolResult(callID: orphanInvocation.call.id, content: "native result")),
+            runID: oldRunID)
+
+        let nilBoundary = AgentCompactionRecord(
+            compactedThroughMessageID: nil, sequence: 0, summary: "Legacy base",
+            summarizedTokens: 4, keptTokens: 1, displayAnchor: instant)
+        let keptCompaction = AgentCompactionRecord(
+            compactedThroughMessageID: prefixResult.id, sequence: 1, summary: "Kept prefix",
+            summarizedTokens: 8, keptTokens: 2, displayAnchor: instant)
+        for record in [nilBoundary, keptCompaction] {
+            try await repository.persistCompaction(record, conversationID: conversationID)
+        }
+
+        let preEditStore = AssistantConversationStore()
+        preEditStore.restore(try await repository.loadConversations())
+        let persistedLive = try #require(preEditStore.conversations.first)
+        #expect(
+            Set(persistedLive.toolCards.map(\.invocation.id)) == [
+                prefixInvocation.id, suffixInvocation.id, orphanInvocation.id,
+            ])
+        let operations = AssistantConversationOperations(repository: repository)
+        #expect(
+            try await operations.edit(
+                persistedLive, messageID: steering.id, replacement: replacement) != nil)
+        #expect(persistedLive.toolCards.map(\.invocation.id) == [prefixInvocation.id])
+        #expect(persistedLive.runSummaries[oldRunID]?.finishedAt == livePrefixFinishedAt)
+
+        let staged = try #require(try await repository.loadConversations().first)
+        #expect(staged.messages.map(\.id) == [opening.id, prefixAssistant.id, prefixResult.id, replacement.id])
+        #expect(staged.compactions.map(\.id) == [nilBoundary.id, keptCompaction.id])
+        #expect(staged.title == "Preserved title")
+        #expect(staged.toolCards.map(\.invocation.id) == [prefixInvocation.id])
+        #expect(staged.toolCards.first?.invocation.sourceMessageID == prefixAssistant.id)
+        #expect(staged.runs.map(\.id) == [oldRunID, replacement.id])
+        let stagedPrefix = try #require(staged.runs.first { $0.id == oldRunID })
+        #expect(stagedPrefix.state == .completed)
+        #expect(stagedPrefix.failure == nil)
+        #expect(stagedPrefix.finishedAt == livePrefixFinishedAt)
+        #expect(staged.messageRunIDs[opening.id] == oldRunID)
+        #expect(staged.messageRunIDs[prefixResult.id] == oldRunID)
+        #expect(staged.messageRunIDs[replacement.id] == replacement.id)
+
+        let attachmentRows = try container.mainContext.fetch(
+            FetchDescriptor<PersistentAgentContextAttachments>())
+        let replayRows = try container.mainContext.fetch(
+            FetchDescriptor<PersistentAgentMessageReplay>())
+        #expect(Set(attachmentRows.map(\.messageID)) == [opening.id, replacement.id])
+        #expect(Set(replayRows.map(\.messageID)) == [opening.id, prefixResult.id, replacement.id])
+        #expect(try container.mainContext.fetchCount(FetchDescriptor<PersistentAgentToolInvocation>()) == 1)
+        #expect(try container.mainContext.fetchCount(FetchDescriptor<PersistentAgentToolResult>()) == 1)
+        #expect(try container.mainContext.fetchCount(FetchDescriptor<PersistentAgentReviewRecord>()) == 1)
+
+        var realSnapshot = AgentRunSnapshot(
+            id: realRunID, conversationID: conversationID, state: .running,
+            permissionMode: .askForApproval,
+            messages: [opening, prefixAssistant, prefixResult, replacement],
+            startedAt: instant)
+        try await repository.save(realSnapshot)
+        let adopted = try #require(try await repository.loadConversations().first)
+        #expect(adopted.runs.map(\.id) == [oldRunID, realRunID])
+        #expect(try #require(adopted.runs.first { $0.id == realRunID }).startedAt > instant)
+        realSnapshot.state = .completed
+        realSnapshot.finishedAt = instant.addingTimeInterval(1)
+        try await repository.save(realSnapshot)
+        try await repository.save(oldSnapshot)
+        let lateInvocation = AgentToolInvocation(
+            runID: oldRunID,
+            call: AgentToolCall(id: "late-call", name: "inspect", arguments: .object([:])),
+            sourceMessageID: suffixAssistant.id, createdAt: instant)
+        try await repository.record(.toolProposed(lateInvocation, descriptor), runID: oldRunID)
+
+        let reopened = try #require(try await repository.loadConversations().first)
+        #expect(reopened.messages.map(\.id) == [opening.id, prefixAssistant.id, prefixResult.id, replacement.id])
+        #expect(reopened.messageRunIDs[replacement.id] == realRunID)
+        #expect(reopened.runs.map(\.id) == [oldRunID, realRunID])
+        #expect(reopened.runs.first { $0.id == oldRunID }?.finishedAt == livePrefixFinishedAt)
+        #expect(reopened.lastRunState == .completed)
+        #expect(try #require(reopened.runs.first { $0.id == realRunID }).startedAt > instant)
+        #expect(reopened.toolCards.map(\.invocation.id) == [prefixInvocation.id])
+        #expect(
+            try container.mainContext.fetch(
+                FetchDescriptor<PersistentAgentMessage>(
+                    predicate: #Predicate { $0.id == replacementID }
+                )
+            ).first?.runID == realRunID)
+        #expect(
+            try container.mainContext.fetch(
+                FetchDescriptor<PersistentAgentRun>(
+                    predicate: #Predicate { $0.id == replacementID }
+                )
+            ).isEmpty)
+    }
+
+    @Test
+    func editingALaterRunOpeningKeepsAnEarlierEmptyRunBeforeTheReplacement() async throws {
+        let configuration = localModelConfiguration(
+            "AgentEmptyPrefixRunTruncationTests", schema: HostPersistence.schema,
+            isStoredInMemoryOnly: true
+        )
+        let container = try ModelContainer(
+            for: HostPersistence.schema, configurations: [configuration]
+        )
+        let repository = SwiftDataAgentRunRepository(context: container.mainContext)
+        let conversationID = UUID()
+        let emptyRunID = try #require(UUID(uuidString: "FFFFFFFF-FFFF-FFFF-FFFF-FFFFFFFFFFFE"))
+        let targetRunID = try #require(UUID(uuidString: "FFFFFFFF-FFFF-FFFF-FFFF-FFFFFFFFFFFF"))
+        let replacementID = try #require(UUID(uuidString: "00000000-0000-0000-0000-000000000001"))
+        let messageClock = Date(timeIntervalSinceReferenceDate: 100)
+        let emptyRunClock = Date(timeIntervalSinceReferenceDate: 200)
+        let targetRunClock = Date(timeIntervalSinceReferenceDate: 300)
+        try await repository.save(
+            AgentRunSnapshot(
+                id: emptyRunID, conversationID: conversationID, state: .completed,
+                permissionMode: .askForApproval, messages: [],
+                startedAt: emptyRunClock, finishedAt: emptyRunClock))
+        let target = AgentTranscriptMessage(
+            role: .user, text: "Later opening", authoredText: "Later opening",
+            createdAt: messageClock)
+        try await repository.save(
+            AgentRunSnapshot(
+                id: targetRunID, conversationID: conversationID, state: .failed,
+                permissionMode: .askForApproval, messages: [target],
+                startedAt: targetRunClock, finishedAt: targetRunClock, failure: "old failure"))
+        let store = AssistantConversationStore()
+        store.restore(try await repository.loadConversations())
+        let conversation = try #require(store.conversations.first)
+        let replacement = AgentTranscriptMessage(
+            id: replacementID, role: .user, text: "Replacement", authoredText: "Replacement",
+            createdAt: messageClock)
+
+        let operations = AssistantConversationOperations(repository: repository)
+        #expect(
+            try await operations.edit(
+                conversation, messageID: target.id, replacement: replacement) != nil)
+
+        let liveStaging = try #require(conversation.runSummaries[replacement.id])
+        #expect(liveStaging.startedAt > emptyRunClock)
+        #expect(
+            conversation.runSummaries.values.sorted {
+                if $0.startedAt != $1.startedAt { return $0.startedAt < $1.startedAt }
+                return $0.id.uuidString < $1.id.uuidString
+            }.map(\.id) == [emptyRunID, replacement.id])
+        let reopened = try #require(try await repository.loadConversations().first)
+        #expect(reopened.runs.map(\.id) == [emptyRunID, replacement.id])
+        #expect(reopened.runs.last?.startedAt == liveStaging.startedAt)
+        #expect(reopened.messages.map(\.id) == [replacement.id])
+    }
+
     @Test
     func truncationRejectsAMessageOutsideThePersistedConversation() async throws {
         let configuration = localModelConfiguration(

TFMTests/ChatComposerStateTests.swift

@@ -2257,7 +2257,97 @@ struct ChatComposerStateTests {
 
         #expect(conversation.workGroups.isEmpty)
         #expect(conversation.toolCards.isEmpty)
-        #expect(conversation.runSummaries.isEmpty)
+        #expect(conversation.runSummaries.count == 1)
+        let staging = try #require(conversation.runSummaries.values.first)
+        #expect(staging.id == conversation.transcript.last?.id)
+        #expect(staging.state == .completed)
+        #expect(staging.failure == nil)
+    }
+
+    @Test
+    func editingSteeringKeepsCompactedSameRunPrefixAndLegacyToolOwnership() throws {
+        let instant = Date(timeIntervalSinceReferenceDate: 100)
+        let runID = UUID()
+        let opening = AgentTranscriptMessage(
+            role: .user, text: "Opening", authoredText: "Opening", createdAt: instant)
+        let call = AgentToolCall(id: "prefix-call", name: "inspect", arguments: .object([:]))
+        let firstReply = AgentTranscriptMessage(
+            role: .assistant, text: "First reply", toolCalls: [call], createdAt: instant)
+        let steering = AgentTranscriptMessage(
+            role: .user, text: "Steer", authoredText: "Steer", createdAt: instant)
+        let secondReply = AgentTranscriptMessage(
+            role: .assistant, text: "Second reply", createdAt: instant)
+        let keptCompaction = AgentCompactionRecord(
+            compactedThroughMessageID: firstReply.id, sequence: 0, summary: "Earlier context",
+            summarizedTokens: 10, keptTokens: 2, displayAnchor: instant)
+        let removedCompaction = AgentCompactionRecord(
+            compactedThroughMessageID: secondReply.id, sequence: 1, summary: "Later context",
+            summarizedTokens: 12, keptTokens: 2, displayAnchor: instant)
+        let conversation = ConversationModel(
+            messages: [
+                ChatMessage(id: opening.id, role: .user, text: opening.text, createdAt: instant),
+                ChatMessage(id: firstReply.id, role: .assistant, text: firstReply.text, createdAt: instant),
+                ChatMessage(id: steering.id, role: .user, text: steering.text, createdAt: instant),
+                ChatMessage(id: secondReply.id, role: .assistant, text: secondReply.text, createdAt: instant),
+            ],
+            transcript: [steering, secondReply], runState: .failed,
+            messageRunIDs: [
+                opening.id: runID, firstReply.id: runID, steering.id: runID, secondReply.id: runID,
+            ],
+            runs: [
+                AgentRunSummary(
+                    id: runID, state: .failed, startedAt: instant,
+                    finishedAt: instant, failure: "old failure")
+            ],
+            compactionRecords: [keptCompaction, removedCompaction])
+        conversation.title = "Keep this title"
+        let invocation = AgentToolInvocation(
+            runID: runID, call: call, sourceMessageID: firstReply.id,
+            createdAt: instant.addingTimeInterval(1))
+        conversation.handle(
+            .toolProposed(
+                invocation,
+                AgentToolDescriptor(
+                    name: call.name, summary: "Inspect", inputSchema: .object([:]),
+                    target: .local, approvalPolicy: .approve)))
+        conversation.handle(
+            .planPresented(
+                AgentPlan(
+                    title: "Same-time plan", body: "Must be removed", scratchPath: "/tmp/plan.md",
+                    createdAt: instant)))
+
+        let replacement = try #require(
+            conversation.editUserMessage(steering.id, authoredText: "Replacement"))
+
+        #expect(conversation.messages.map(\.text) == ["Opening", "First reply", "Replacement"])
+        #expect(conversation.transcript.map(\.text) == ["Replacement"])
+        #expect(conversation.toolCards.map(\.id) == [invocation.id])
+        #expect(conversation.compactionRecords == [keptCompaction])
+        #expect(conversation.compactions.map(\.id) == [keptCompaction.id])
+        #expect(conversation.plans.isEmpty)
+        #expect(conversation.title == "Keep this title")
+        let prefix = try #require(conversation.runSummaries[runID])
+        #expect(prefix.state == .completed)
+        #expect(prefix.finishedAt == instant)
+        #expect(prefix.failure == nil)
+        #expect(conversation.messageRunIDs[opening.id] == runID)
+        #expect(conversation.messageRunIDs[firstReply.id] == runID)
+        #expect(conversation.messageRunIDs[replacement.id] == replacement.id)
+
+        let stagingStartedAt = try #require(conversation.runSummaries[replacement.id]).startedAt
+        let realRunID = UUID()
+        conversation.handle(
+            .runStarted(
+                AgentRunSummary(id: realRunID, state: .running, startedAt: instant)))
+        #expect(conversation.messageRunIDs[replacement.id] == realRunID)
+        #expect(conversation.runSummaries[realRunID]?.startedAt == stagingStartedAt)
+        conversation.handle(
+            .runFinished(
+                AgentRunSummary(
+                    id: realRunID, state: .completed, startedAt: instant,
+                    finishedAt: instant)))
+        #expect(conversation.runSummaries[realRunID]?.startedAt == stagingStartedAt)
+        #expect(conversation.runSummaries[realRunID]?.finishedAt == stagingStartedAt)
     }
 
     private func waitForTerminal(_ conversation: ConversationModel) async {