diff --git a/FrameworkTests/AgentReplayPersistenceTests.swift b/FrameworkTests/AgentReplayPersistenceTests.swift index 9c719b0e..6e7e9f2b 100644 --- a/FrameworkTests/AgentReplayPersistenceTests.swift +++ b/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()) + let runRows = try container.mainContext.fetch(FetchDescriptor()) + let replayRows = try container.mainContext.fetch(FetchDescriptor()) + #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 diff --git a/FrameworkTests/AssistantConversationOperationsTests.swift b/FrameworkTests/AssistantConversationOperationsTests.swift index f17f2714..8a1d2f85 100644 --- a/FrameworkTests/AssistantConversationOperationsTests.swift +++ b/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 { diff --git a/FrameworkTests/AssistantExecutionControllerTests.swift b/FrameworkTests/AssistantExecutionControllerTests.swift index 48e42649..fef4d96e 100644 --- a/FrameworkTests/AssistantExecutionControllerTests.swift +++ b/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] = [] private var saveReleaseWaiters: [CheckedContinuation] = [] + private var snapshotWaiters: [(state: AgentRunState, continuation: CheckedContinuation)] = [] private var truncationEntered = false private var truncationReleased = false private var truncationEntryWaiters: [CheckedContinuation] = [] @@ -620,6 +660,15 @@ private actor RecordingAgentRunRepository: AgentRunPersisting { func save(_ snapshot: AgentRunSnapshot) async { snapshots.append(snapshot) + var pendingSnapshotWaiters: [(state: AgentRunState, continuation: CheckedContinuation)] = [] + 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 diff --git a/Frameworks/TermindApplication/Agent/AssistantConversationOperations.swift b/Frameworks/TermindApplication/Agent/AssistantConversationOperations.swift index 39d82313..a1ef8e4e 100644 --- a/Frameworks/TermindApplication/Agent/AssistantConversationOperations.swift +++ b/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) } } diff --git a/Frameworks/TermindApplication/Agent/AssistantExecutionController.swift b/Frameworks/TermindApplication/Agent/AssistantExecutionController.swift index 94d38bcb..b5432284 100644 --- a/Frameworks/TermindApplication/Agent/AssistantExecutionController.swift +++ b/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 } diff --git a/Frameworks/TermindApplication/Agent/ConversationModel.swift b/Frameworks/TermindApplication/Agent/ConversationModel.swift index 4a680a0e..74ea0208 100644 --- a/Frameworks/TermindApplication/Agent/ConversationModel.swift +++ b/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[.. targetSummary.startedAt ? summary.id : nil + } + return summary.id.uuidString > targetSummary.id.uuidString ? summary.id : nil + }) + } + let hasEarlierVisibleUser = + visibleTarget.map { index in + messages[..= 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[.. 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() diff --git a/Frameworks/TermindPersistence/SwiftDataAgentRunRepository.swift b/Frameworks/TermindPersistence/SwiftDataAgentRunRepository.swift index 31e7aae1..6632ff25 100644 --- a/Frameworks/TermindPersistence/SwiftDataAgentRunRepository.swift +++ b/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(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(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( - 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( + 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( - 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( + 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( + 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()).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()).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()) 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()) let runs = allRuns.filter { $0.conversationID == id } - .sorted { $0.startedAt < $1.startedAt } + .sorted(by: Self.runPrecedes) let allMessages = try context.fetch(FetchDescriptor()) + 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[..()) 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[..()) + where messageIDs.contains(row.messageID) { context.delete(row) } + for row in try context.fetch(FetchDescriptor()) + 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()) + 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()) + where row.runID == runID && invocationIDs.contains(row.invocationID) { context.delete(row) } + for row in try context.fetch(FetchDescriptor()) + 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(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()) + guard messages.filter({ $0.runID == id }).count == 1 else { return false } + let invocations = try context.fetch(FetchDescriptor()) + let results = try context.fetch(FetchDescriptor()) + let reviews = try context.fetch(FetchDescriptor()) + 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)) diff --git a/TFITests/MobileAssistantTests.swift b/TFITests/MobileAssistantTests.swift index 42a66c19..511cce54 100644 --- a/TFITests/MobileAssistantTests.swift +++ b/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() diff --git a/TFMTests/AgentPersistenceTests.swift b/TFMTests/AgentPersistenceTests.swift index 00071731..3c169e87 100644 --- a/TFMTests/AgentPersistenceTests.swift +++ b/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()) + let replayRows = try container.mainContext.fetch( + FetchDescriptor()) + #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()) == 1) + #expect(try container.mainContext.fetchCount(FetchDescriptor()) == 1) + #expect(try container.mainContext.fetchCount(FetchDescriptor()) == 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( + predicate: #Predicate { $0.id == replacementID } + ) + ).first?.runID == realRunID) + #expect( + try container.mainContext.fetch( + FetchDescriptor( + 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( diff --git a/TFMTests/ChatComposerStateTests.swift b/TFMTests/ChatComposerStateTests.swift index a3430438..d4bb3185 100644 --- a/TFMTests/ChatComposerStateTests.swift +++ b/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 {