diff --git a/apps/swift-ios/App/NativeConversationRewind.swift b/apps/swift-ios/App/NativeConversationRewind.swift new file mode 100644 index 000000000000..351e74557a03 --- /dev/null +++ b/apps/swift-ios/App/NativeConversationRewind.swift @@ -0,0 +1,93 @@ +import Foundation + +enum NativeConversationRewind { + /// Use checkpoint counts, not visible message indexes. Pages and steering messages + /// do not each represent one completed provider turn. + static func turnCount(before messageID: String, in thread: OrchestrationThread) -> Int? { + guard let index = thread.messages.firstIndex(where: { $0.id == messageID }), + thread.messages[index].role == "user" else { return nil } + let checkpoints = Dictionary( + thread.checkpoints.compactMap { checkpoint in + checkpoint.assistantMessageId.map { ($0, checkpoint.checkpointTurnCount) } + }, + uniquingKeysWith: max + ) + for message in thread.messages.dropFirst(index + 1) { + if message.role == "user" { return nil } + if let count = checkpoints[message.id] { return max(0, count - 1) } + } + return nil + } + + static func isComplete(_ thread: OrchestrationThread, messageID: String, turnCount: Int) -> Bool { + !thread.messages.contains(where: { $0.id == messageID }) + && thread.checkpoints.allSatisfy { $0.checkpointTurnCount <= turnCount } + && (turnCount == 0 + ? thread.latestTurn == nil + : thread.checkpoints.contains { $0.turnId == thread.latestTurn?.turnId }) + } + + /// Command acceptance precedes provider rollback. Wait for its completion event + /// or an authoritative replacement snapshot, including while another thread is open. + static func waitForCompletion( + batches: AsyncThrowingStream<[ThreadStreamItem], Error>, + threadID: String, + messageID: String, + turnCount: Int, + afterSequence: Int, + previousFailureIDs: Set, + timeout: Duration = .seconds(120) + ) async throws -> Int { + try await withThrowingTaskGroup(of: Int.self) { group in + group.addTask { + for try await batch in batches { + for item in batch { + switch item { + case .synchronized: + continue + case let .snapshot(snapshot): + guard snapshot.thread.id == threadID, + snapshot.snapshotSequence > afterSequence else { continue } + if isComplete(snapshot.thread, messageID: messageID, turnCount: turnCount) { + return snapshot.snapshotSequence + } + if let failure = snapshot.thread.activities.last(where: { + $0.kind == "checkpoint.revert.failed" && !previousFailureIDs.contains($0.id) + && $0.payload["turnCount"] == .number(Double(turnCount)) + }) { + throw FeatureConversationRewindError( + message: failure.payload["detail"]?.stringValue ?? failure.summary, + didNotRevert: true + ) + } + case let .event(event): + guard event["payload"]?["threadId"]?.stringValue == threadID, + case let .number(sequence)? = event["sequence"], + sequence > Double(afterSequence) else { continue } + if event["type"]?.stringValue == "thread.activity-appended", + let activity = event["payload"]?["activity"], + activity["kind"]?.stringValue == "checkpoint.revert.failed", + activity["payload"]?["turnCount"] == .number(Double(turnCount)) { + throw FeatureConversationRewindError( + message: activity["payload"]?["detail"]?.stringValue + ?? activity["summary"]?.stringValue ?? "Conversation rewind failed.", + didNotRevert: true + ) + } + if event["type"]?.stringValue == "thread.reverted", + event["payload"]?["turnCount"] == .number(Double(turnCount)) { return Int(sequence) } + } + } + } + throw FeatureConversationRewindError(message: "The connection closed before rewind finished. Reload the thread before trying again.") + } + group.addTask { + try await Task.sleep(for: timeout) + throw FeatureConversationRewindError(message: "Timed out waiting for rewind. Reload the thread before trying again.") + } + defer { group.cancelAll() } + guard let sequence = try await group.next() else { throw CancellationError() } + return sequence + } + } +} diff --git a/apps/swift-ios/App/NativeFeatureClient.swift b/apps/swift-ios/App/NativeFeatureClient.swift index e835e340630f..907ccc65084a 100644 --- a/apps/swift-ios/App/NativeFeatureClient.swift +++ b/apps/swift-ios/App/NativeFeatureClient.swift @@ -2317,6 +2317,156 @@ final class NativeFeatureClient: FeatureClient, FeatureDeviceManaging, try? await refresh(client: route.client) } + func canRewindConversation(threadID: String, messageID: String) -> Bool { + guard let route = try? threadRoute(for: threadID), + let thread = activeThreadID == route.uiID + ? activeRawThread : threadResumeStates[route.uiID]?.thread, + let provider = serverConfigsByEnvironmentID[route.environmentID]?.providers.first(where: { + $0.instanceId == (thread.session?.providerInstanceId ?? thread.modelSelection.instanceId) + }), + provider.supportsConversationRollback != false, + provider.driver != "cursor", provider.driver != "grok" else { return false } + return NativeConversationRewind.turnCount(before: messageID, in: thread) != nil + } + + func rewindConversation( + threadID: String, messageID: String, + prepareRecovery: @MainActor (FeatureRevertedMessage) async throws -> Void + ) async throws { + let route = try threadRoute(for: threadID) + let generation = environmentGeneration + // The visible page can omit checkpoints. Validate against the whole thread. + let snapshot = try await route.client.threadSnapshot(id: route.wireID) + guard let provider = serverConfigsByEnvironmentID[route.environmentID]?.providers.first(where: { + $0.instanceId == (snapshot.thread.session?.providerInstanceId ?? snapshot.thread.modelSelection.instanceId) + }), provider.supportsConversationRollback != false, + provider.driver != "cursor", provider.driver != "grok" else { + throw FeatureCapabilityUnavailable("Conversation rewind for this provider") + } + guard snapshot.thread.session?.status != "running", + snapshot.thread.session?.status != "starting", + let turnCount = NativeConversationRewind.turnCount(before: messageID, in: snapshot.thread), + let original = snapshot.thread.messages.first(where: { $0.id == messageID }) else { + throw FeatureConversationRewindError(message: "Wait for this turn to finish before rewinding.") + } + let message = mapMessage(original, environmentID: route.environmentID) + let fileStore = ManagedAttachmentFileStore() + var attachments: [FeatureDraftAttachment] = [] + var recoveryIsStored = false + defer { + if !recoveryIsStored { + for attachment in attachments { + if let file = attachment.ownedFile { + try? fileStore.removeOwnedFile(fileName: file.fileName) + } + } + } + } + // Rewind removes the old server assets, even when workspace files are kept. + for attachment in message.attachments { + let url = try await attachmentAssetURL(threadID: route.uiID, attachment: attachment) + let (temporaryURL, response) = try await URLSession.shared.download(from: url) + defer { try? FileManager.default.removeItem(at: temporaryURL) } + guard let response = response as? HTTPURLResponse, + (200..<300).contains(response.statusCode) else { + throw FeatureConversationRewindError(message: "Could not save \(attachment.name) before rewind.") + } + let id = UUID() + let file = try await Task.detached { + try fileStore.copyOwnedFile( + from: temporaryURL, attachmentID: id, originalFileName: attachment.name + ) + }.value + attachments.append(FeatureDraftAttachment( + id: id, ownedFile: file, filename: attachment.name, mimeType: attachment.mimeType, + source: attachment.source + )) + } + try await prepareRecovery(FeatureRevertedMessage(message: message, attachments: attachments)) + recoveryIsStored = true + let subscription: (events: AsyncThrowingStream<[ThreadStreamItem], Error>, connectionID: UUID) + do { + try Task.checkCancellation() + guard isKnownClient(route.client, environmentID: route.environmentID, generation: generation) else { + throw CancellationError() + } + subscription = try await route.client.threadEventBatches( + threadID: route.wireID, after: snapshot.snapshotSequence + ) + } catch { + throw FeatureConversationRewindError(message: error.localizedDescription, didNotRevert: true) + } + // Keep events that arrive before dispatch replies. The pump owns socket + // cancellation even when dispatch rejects before completion tracking starts. + let buffered = AsyncThrowingStream<[ThreadStreamItem], Error>.makeStream(bufferingPolicy: .bufferingOldest(64)) + let pump = Task { + do { + for try await batch in subscription.events { + if case .dropped = buffered.continuation.yield(batch) { + throw FeatureConversationRewindError(message: "Rewind updates fell behind. Reload the thread to check its history.") + } + } + buffered.continuation.finish() + } catch { + buffered.continuation.finish(throwing: error) + } + } + defer { + pump.cancel() + buffered.continuation.finish() + } + do { + let accepted: DispatchResult + do { + accepted = try await route.client.dispatch(OrchestrationCommands.revertConversation( + threadID: route.wireID, turnCount: turnCount + )) + } catch let error as RPCError { + if case .remote = error { + throw FeatureConversationRewindError(message: error.localizedDescription, didNotRevert: true) + } + throw error + } catch let error as HTTPError { + if case let .status(code, _, _) = error, (400..<500).contains(code) { + throw FeatureConversationRewindError(message: error.localizedDescription, didNotRevert: true) + } + throw error + } + let receiptSequence = try await NativeConversationRewind.waitForCompletion( + batches: buffered.stream, + threadID: route.wireID, + messageID: messageID, + turnCount: turnCount, + afterSequence: accepted.sequence, + previousFailureIDs: Set(snapshot.thread.activities.filter { + $0.kind == "checkpoint.revert.failed" + }.map(\.id)) + ) + let current = try await route.client.threadSnapshot(id: route.wireID) + guard current.snapshotSequence >= receiptSequence, + NativeConversationRewind.isComplete(current.thread, messageID: messageID, turnCount: turnCount) else { + throw FeatureConversationRewindError(message: "Could not confirm the final rewind state. The prompt remains saved for recovery.") + } + } catch { + if (error as? FeatureConversationRewindError)?.didNotRevert == true { throw error } + // A lost socket does not mean rollback failed. An HTTP read can + // confirm completion without submitting the destructive command again. + guard let current = try? await route.client.threadSnapshot(id: route.wireID), + current.snapshotSequence > snapshot.snapshotSequence, + NativeConversationRewind.isComplete(current.thread, messageID: messageID, turnCount: turnCount) else { + throw error + } + } + // Do not turn a failed refresh into a failed rewind. The receipt confirms + // that history changed, so the recovered prompt must still reach the draft. + threadResumeStates[route.uiID] = nil + do { + try await refreshThread(id: route.uiID, client: route.client) + } catch { + continuation.yield(.threadSync(id: route.uiID, state: .failed(error.localizedDescription))) + } + } + func resolveApproval(id: String, decision: FeatureApprovalDecision) async throws { guard let request = approvalRoutes[id] else { throw NativeFeatureClientError.approvalNotFound diff --git a/apps/swift-ios/Core/ServerConfigModels.swift b/apps/swift-ios/Core/ServerConfigModels.swift index 8abc432de721..bd0186c1a46f 100644 --- a/apps/swift-ios/Core/ServerConfigModels.swift +++ b/apps/swift-ios/Core/ServerConfigModels.swift @@ -137,6 +137,7 @@ public struct ServerProviderSnapshot: Codable, Identifiable, Equatable, Sendable public var workspaceSnapshots: [ServerProviderWorkspaceSnapshot]? = nil public var setup: ProviderSetupCapabilities? = nil public var usageLimits: ServerProviderUsageLimits? = nil + public var supportsConversationRollback: Bool? = nil } public enum ServerThreadEnvironmentMode: String, Codable, Equatable, Sendable { diff --git a/apps/swift-ios/Core/T3Client.swift b/apps/swift-ios/Core/T3Client.swift index 235302e9d3d9..1a1c03a671ba 100644 --- a/apps/swift-ios/Core/T3Client.swift +++ b/apps/swift-ios/Core/T3Client.swift @@ -2065,6 +2065,22 @@ public enum RPCMethod: String, Sendable { } public enum OrchestrationCommands { + /// A distinct command makes older servers reject this action without restoring files. + public static func revertConversation( + threadID: String, + turnCount: Int, + commandID: String = UUID().uuidString, + createdAt: String = now() + ) -> JSONValue { + .object([ + "type": .string("thread.conversation.revert"), + "commandId": .string(commandID), + "threadId": .string(threadID), + "turnCount": .number(Double(turnCount)), + "createdAt": .string(createdAt), + ]) + } + public static func createThread( threadID: String = UUID().uuidString, projectID: String, diff --git a/apps/swift-ios/Features/Chat/FeatureComposerView.swift b/apps/swift-ios/Features/Chat/FeatureComposerView.swift index 4319f9cf3f8d..389e1974641d 100644 --- a/apps/swift-ios/Features/Chat/FeatureComposerView.swift +++ b/apps/swift-ios/Features/Chat/FeatureComposerView.swift @@ -29,6 +29,7 @@ struct FeatureComposerUploadStatus { struct FeatureComposerView: View { @SwiftUI.Environment(\.scenePhase) private var scenePhase + @SwiftUI.Environment(\.isEnabled) private var isEnabled @State private var isManuallyExpanded = false @State private var isAttachmentFlowActive = false @State private var isModelPickerPresented = false @@ -61,6 +62,7 @@ struct FeatureComposerView: View { private let onRefreshModels: (() async throws -> Void)? private let draftSaveError: String? private let onRetryDraftSave: (() -> Void)? + private let onInputPreparationChange: ((Bool) -> Void)? private let threadSelection: FeatureSelection? private let materializesDefaultSelection: Bool private let isSending: Bool @@ -112,7 +114,8 @@ struct FeatureComposerView: View { onRefreshModels: (() async throws -> Void)? = nil, draftSaveError: String? = nil, onRetryDraftSave: (() -> Void)? = nil, - context: Binding = .constant(nil) + context: Binding = .constant(nil), + onInputPreparationChange: ((Bool) -> Void)? = nil ) { _text = text _selection = selection @@ -124,6 +127,7 @@ struct FeatureComposerView: View { self.environmentIsConnected = environmentIsConnected self.attachmentUploads = attachmentUploads self.attachmentPreferences = attachmentPreferences + self.onInputPreparationChange = onInputPreparationChange self.onRefreshModels = onRefreshModels self.draftSaveError = draftSaveError self.onRetryDraftSave = onRetryDraftSave @@ -302,6 +306,9 @@ struct FeatureComposerView: View { .stroke(T3Colors.inputBorder, lineWidth: 1) } .clipShape(composerShape) + .onChange(of: attachmentPreparation.isPreparing || isAttachmentFlowActive || voiceInputController.isBusy, initial: true) { _, busy in + onInputPreparationChange?(busy) + } .modifier( FeatureComposerImageDrop( isEnabled: imagesAllowed && !voiceInputController.isBusy, @@ -365,7 +372,7 @@ struct FeatureComposerView: View { focused: $focused, placeholder: composerPlaceholder, acceptsImages: imagesAllowed, - isReadOnly: voiceInputController.isBusy, + isReadOnly: voiceInputController.isBusy || !isEnabled, skills: powerFeatures.enabledSkills, selectionRequest: textSelectionRequest, onSelectionChange: handleTextSelectionChange, diff --git a/apps/swift-ios/Features/Chat/ThreadDetailView.swift b/apps/swift-ios/Features/Chat/ThreadDetailView.swift index 15cb7eaaf74f..14c6973d4959 100644 --- a/apps/swift-ios/Features/Chat/ThreadDetailView.swift +++ b/apps/swift-ios/Features/Chat/ThreadDetailView.swift @@ -20,6 +20,9 @@ public struct ThreadDetailView: View { @State private var selection: FeatureSelection? @State private var attachments: [FeatureDraftAttachment] = [] @State private var isSending = false + @State private var pendingRewindMessageID: String? + @State private var isPreparingRewind = false + @State private var isPreparingInput = false @State private var submittingCompaction = false @State private var isLoading = true @State private var sendFailed = false @@ -56,7 +59,7 @@ public struct ThreadDetailView: View { self.draftStore = draftStore } - public var body: some View { + private var threadContent: some View { Group { if let detail { timeline(detail) @@ -104,6 +107,7 @@ public struct ThreadDetailView: View { .task(id: thread.id) { // A cached thread can already show its composer while the server // is catching up. Local drafts must not wait for that request. + await model.checkRewindRecovery(for: currentThread) guard !didRestoreDraft else { return } await restoreDraft(from: composerDraft, key: draftKey) } @@ -121,6 +125,9 @@ public struct ThreadDetailView: View { }) .onChange(of: draft) { scheduleDraftSave() } .onChange(of: selection) { scheduleDraftSave() } + .onChange(of: model.recoveredRewindDrafts[thread.id]) { _, recovered in + if recovered != nil { restoreRewindDraft() } + } .onChange(of: threadConnectionState) { _, state in if state == .connected, case .failed = model.detailLoadStates[thread.id], @@ -185,6 +192,10 @@ public struct ThreadDetailView: View { .presentationDragIndicator(.visible) .t3CodeSizing(steps: codeSizeSteps) } + } + + public var body: some View { + threadContent .alert("Message not sent", isPresented: $sendFailed) { // Refocusing happens here rather than when the send fails: the // alert takes first responder from the composer, so a refocus @@ -193,6 +204,19 @@ public struct ThreadDetailView: View { } message: { Text("Your draft is still here. Check your connection and try again.") } + .confirmationDialog("Edit from here?", isPresented: Binding( + get: { pendingRewindMessageID != nil }, + set: { if !$0 { pendingRewindMessageID = nil } } + ), titleVisibility: .visible) { + Button("Revert and keep changes") { + guard let messageID = pendingRewindMessageID else { return } + pendingRewindMessageID = nil + rewindConversation(before: messageID) + } + Button("Cancel", role: .cancel) { pendingRewindMessageID = nil } + } message: { + Text("Rewind chat to before this message. Your prompt and attachments return to the composer. File changes stay as they are.") + } .alert( feedbackIdentifier == nil ? "Could not send feedback" : "Feedback sent to OpenAI", isPresented: Binding( @@ -259,6 +283,53 @@ public struct ThreadDetailView: View { submittingCompaction || detail?.isCompacting == true } + private var isRewinding: Bool { + isPreparingRewind || model.rewindingThreadIDs.contains(thread.id) + } + + private func canRewind(_ messageID: String) -> Bool { + !isSending && !isRewinding && !isPreparingInput && didRestoreDraft + && model.canRewindConversation(threadID: thread.id, messageID: messageID) + } + + private func rewindConversation(before messageID: String) { + guard canRewind(messageID) else { return } + isPreparingRewind = true + dismissKeyboard() + let pendingSave = draftSaveTask + pendingSave?.cancel() + draftSaveTask = nil + let saved = composerDraft + Task { + await pendingSave?.value + await model.rewindConversation(threadID: thread.id, messageID: messageID, draft: saved) + isPreparingRewind = false + } + } + + private func restoreRewindDraft() { + guard let recovered = model.consumeRewindDraft(threadID: thread.id) else { return } + draft = recovered.text + attachments = recovered.attachments + selection = recovered.selection + composerContext = recovered.context + didRestoreDraft = true + } + + private func recoverSavedRewind() { + guard !isRewinding && !isSending && !isPreparingInput else { return } + isPreparingRewind = true + let pendingSave = draftSaveTask + pendingSave?.cancel() + draftSaveTask = nil + let saved = composerDraft + Task { + await pendingSave?.value + await model.recoverSavedRewind(threadID: thread.id, draft: saved) + isPreparingRewind = false + } + } + private var currentSelection: FeatureSelection? { guard let providerID = detail?.thread.providerID ?? thread.providerID, let modelID = detail?.thread.modelID ?? thread.modelID else { return nil } @@ -468,6 +539,12 @@ public struct ThreadDetailView: View { Button(action: reloadThread) { Label("Reload", systemImage: "arrow.clockwise") } + if let message = detail?.messages.last(where: { $0.role == .user }), + canRewind(message.id) { + Button { pendingRewindMessageID = message.id } label: { + Label("Edit last prompt", systemImage: "arrow.uturn.backward") + } + } } Section("Workspace") { Button { toolSurface = .files } label: { @@ -712,13 +789,39 @@ public struct ThreadDetailView: View { onLoadEarlier: { Task { await model.loadEarlierTurns(for: thread.id) } }, - onDismissKeyboard: dismissKeyboard + onDismissKeyboard: dismissKeyboard, + canEditMessage: canRewind, + onEditMessage: { pendingRewindMessageID = $0 } ) } } .safeAreaInset(edge: .bottom, spacing: 0) { VStack(spacing: 0) { refreshStatus + if isRewinding { + Text("Rewinding conversation") + .font(T3Typography.supporting) + .foregroundStyle(T3Colors.textPrimary) + .padding(.vertical, 8) + .accessibilityIdentifier("conversation-rewind-status") + } + if let error = model.rewindErrors[thread.id] { + Text(error) + .font(T3Typography.supporting) + .foregroundStyle(T3Colors.danger) + .padding(.horizontal, 18) + .padding(.vertical, 8) + } + if model.pendingRewindRecoveryIDs.contains(thread.id), !isRewinding { + VStack(spacing: 4) { + Text("A prompt is saved from an unconfirmed rewind. Reload the thread to check its history.") + .font(T3Typography.supporting) + Button("Recover saved prompt", action: recoverSavedRewind) + .disabled(isSending || isPreparingInput || !didRestoreDraft) + } + .padding(.horizontal, 18) + .padding(.vertical, 8) + } FeatureComposerView( text: $draft, selection: $selection, @@ -734,7 +837,7 @@ public struct ThreadDetailView: View { providers: threadProviders, threadSelection: currentSelection, materializesDefaultSelection: false, - isSending: isSending, + isSending: isSending || isRewinding, isWorking: detail.thread.state == .working || detail.thread.state == .queued || isCompacting, focused: $composerFocused, @@ -760,8 +863,10 @@ public struct ThreadDetailView: View { onRefreshModels: refreshThreadEnvironmentModels, draftSaveError: draftSaveError, onRetryDraftSave: persistDraftImmediately, - context: contextBinding + context: contextBinding, + onInputPreparationChange: { isPreparingInput = $0 } ) + .disabled(isRewinding) } .background(T3Colors.background) } @@ -923,6 +1028,7 @@ public struct ThreadDetailView: View { } private func send() { + guard !isRewinding else { return } let message = draft let pendingContext = composerContext let pendingAttachments = currentThread.environmentID.map { @@ -1110,8 +1216,16 @@ public struct ThreadDetailView: View { @MainActor private func restoreDraft(from baseline: FeatureComposerDraft, key: String) async { + if model.recoveredRewindDrafts[thread.id] != nil { + restoreRewindDraft() + return + } let saved = try? await draftStore.draft(for: key) guard !Task.isCancelled else { return } + if model.recoveredRewindDrafts[thread.id] != nil { + restoreRewindDraft() + return + } let liveDraft = composerDraft var restored = FeatureComposerDraftRestoration.merge( @@ -1144,7 +1258,7 @@ public struct ThreadDetailView: View { } private func scheduleDraftSave() { - guard didRestoreDraft else { return } + guard didRestoreDraft, !isRewinding else { return } let previousSave = draftSaveTask previousSave?.cancel() let snapshot = composerDraft @@ -1175,7 +1289,7 @@ public struct ThreadDetailView: View { } private func persistDraftImmediately() { - guard didRestoreDraft else { return } + guard didRestoreDraft, !isRewinding else { return } let previousSave = draftSaveTask previousSave?.cancel() let snapshot = composerDraft @@ -1465,6 +1579,8 @@ private struct FeatureTranscriptCollectionView: UIViewRepresentable { let isLoadingEarlier: Bool let onLoadEarlier: () -> Void let onDismissKeyboard: () -> Void + let canEditMessage: (String) -> Bool + let onEditMessage: (String) -> Void func makeCoordinator() -> Coordinator { Coordinator() @@ -1488,6 +1604,8 @@ private struct FeatureTranscriptCollectionView: UIViewRepresentable { func updateUIView(_ collectionView: UICollectionView, context: Context) { context.coordinator.currentOpenURL = openURL + context.coordinator.canEditMessage = canEditMessage + context.coordinator.onEditMessage = onEditMessage context.coordinator.update( threadID: threadID, messages: messages, @@ -1547,6 +1665,8 @@ private struct FeatureTranscriptCollectionView: UIViewRepresentable { private var orderedIDs: [String] = [] private var currentThreadID: String? var currentOpenURL: OpenURLAction? + var canEditMessage: ((String) -> Bool)? + var onEditMessage: ((String) -> Void)? private var currentImageContext: MarkdownImageContext? private var currentAttachmentContext: FeatureAttachmentContext? private var currentSkills: [FeatureProviderSkill] = [] @@ -1568,6 +1688,24 @@ private struct FeatureTranscriptCollectionView: UIViewRepresentable { markdownPrefetches.values.forEach { $0.task.cancel() } } + func collectionView( + _ collectionView: UICollectionView, + contextMenuConfigurationForItemAt indexPath: IndexPath, + point: CGPoint + ) -> UIContextMenuConfiguration? { + guard let messageID = dataSource?.itemIdentifier(for: indexPath), + messagesByID[messageID]?.role == .user, + canEditMessage?(messageID) == true else { return nil } + return UIContextMenuConfiguration(identifier: nil, previewProvider: nil) { [weak self] _ in + UIMenu(children: [UIAction( + title: "Edit from here", image: UIImage(systemName: "arrow.uturn.backward") + ) { [weak self] _ in + guard self?.canEditMessage?(messageID) == true else { return } + self?.onEditMessage?(messageID) + }]) + } + } + func connect(to collectionView: UICollectionView) { let registration = UICollectionView.CellRegistration { [weak self] cell, _, messageID in diff --git a/apps/swift-ios/Features/Root/FeatureRootModel.swift b/apps/swift-ios/Features/Root/FeatureRootModel.swift index ddcf34a91acd..cb91af596570 100644 --- a/apps/swift-ios/Features/Root/FeatureRootModel.swift +++ b/apps/swift-ios/Features/Root/FeatureRootModel.swift @@ -70,6 +70,10 @@ public final class FeatureRootModel { public private(set) var isLoading = true public private(set) var isPerformingAction = false private(set) var isArrangingThreads = false + public private(set) var rewindingThreadIDs: Set = [] + public private(set) var recoveredRewindDrafts: [String: FeatureComposerDraft] = [:] + public private(set) var rewindErrors: [String: String] = [:] + public private(set) var pendingRewindRecoveryIDs: Set = [] /// Approval and question IDs with a response in flight. Views disable /// only that request, not every request in every thread. public private(set) var resolvingRequestIDs: Set = [] @@ -86,6 +90,7 @@ public final class FeatureRootModel { draftStore: draftStore ) private var pendingSubmissionsByID: [String: FeatureQueuedSubmission] = [:] + private var activeSubmissionCounts: [String: Int] = [:] private var pendingThreadsByID: [String: FeatureThread] = [:] private var pendingSettlementMutations: [String: PendingSettlementMutation] = [:] private var pendingCompletionSubmissionIDs: Set = [] @@ -800,6 +805,12 @@ public final class FeatureRootModel { } public func sendMessage(_ submission: FeatureMessageSubmission) async -> Bool { + guard !rewindingThreadIDs.contains(submission.threadID) else { return false } + activeSubmissionCounts[submission.threadID, default: 0] += 1 + defer { + let remaining = (activeSubmissionCounts[submission.threadID] ?? 1) - 1 + activeSubmissionCounts[submission.threadID] = remaining > 0 ? remaining : nil + } let trimmed = submission.text.trimmingCharacters(in: .whitespacesAndNewlines) guard !trimmed.isEmpty || !submission.attachments.isEmpty else { return false } @@ -884,6 +895,117 @@ public final class FeatureRootModel { } } + public func canRewindConversation(threadID: String, messageID: String) -> Bool { + guard !rewindingThreadIDs.contains(threadID), + !pendingRewindRecoveryIDs.contains(threadID), + activeSubmissionCounts[threadID] == nil, + let detail = details[threadID], + FeatureConversationRewind.canStart(in: detail), + let environmentID = detail.thread.environmentID, + isEnvironmentConnected(environmentID), + !pendingSubmissionsByID.values.contains(where: { $0.threadID == threadID }), + detail.messages.contains(where: { + $0.id == messageID && $0.role == .user && $0.state == .complete + }) else { return false } + return client.canRewindConversation(threadID: threadID, messageID: messageID) + } + + /// The lock and saved recovery live outside the view so navigation cannot + /// enable sends or lose a restored prompt while provider rollback is pending. + public func rewindConversation( + threadID: String, messageID: String, draft: FeatureComposerDraft + ) async { + guard canRewindConversation(threadID: threadID, messageID: messageID), + let detail = details[threadID], + let message = detail.messages.first(where: { $0.id == messageID }) else { return } + guard draft.attachments.count + message.attachments.count <= 8 else { + rewindErrors[threadID] = "Make room for this message's attachments before rewinding." + return + } + rewindingThreadIDs.insert(threadID) + rewindErrors[threadID] = nil + recoveredRewindDrafts[threadID] = nil + defer { rewindingThreadIDs.remove(threadID) } + let key = FeatureComposerDraftStore.threadKey(detail.thread) + let recoveryKey = FeatureComposerDraftStore.rewindRecoveryKey(for: key) + do { + guard try await !draftStore.hasRewindRecovery(for: key) else { + pendingRewindRecoveryIDs.insert(threadID) + rewindErrors[threadID] = "Recover the saved prompt before starting another rewind." + return + } + try await draftStore.setDraft(draft, for: key) + try await client.rewindConversation(threadID: threadID, messageID: messageID) { reverted in + guard self.snapshot.environments.contains(where: { $0.id == detail.thread.environmentID }) else { + throw FeatureConversationRewindError(message: "The computer was removed before rewind started.", didNotRevert: true) + } + let recovery = try FeatureConversationRewind.recover(reverted, draft: FeatureComposerDraft()) + // Check the combined record limit before the server changes history. + _ = try FeatureComposerContext.merge(draft.context, recovery.context) + try await self.draftStore.setDraft(recovery, for: recoveryKey) + self.pendingRewindRecoveryIDs.insert(threadID) + } + try await finishRewindRecovery(thread: detail.thread) + } catch { + rewindErrors[threadID] = error.localizedDescription + if (error as? FeatureConversationRewindError)?.didNotRevert == true { + do { + try await draftStore.discardRewindRecovery(for: key) + pendingRewindRecoveryIDs.remove(threadID) + } catch { + rewindErrors[threadID] = "The saved prompt could not be cleared. \(error.localizedDescription)" + } + } + } + } + + public func checkRewindRecovery(for thread: FeatureThread) async { + guard !rewindingThreadIDs.contains(thread.id) else { return } + do { + let key = FeatureComposerDraftStore.threadKey(thread) + let saved = try await draftStore.hasRewindRecovery(for: key) + guard !rewindingThreadIDs.contains(thread.id) else { return } + if saved { pendingRewindRecoveryIDs.insert(thread.id) } + else { pendingRewindRecoveryIDs.remove(thread.id) } + } catch { + rewindErrors[thread.id] = "Could not read the saved rewind prompt. \(error.localizedDescription)" + } + } + + /// Recovery changes only the local composer. It never sends another rewind. + public func recoverSavedRewind(threadID: String, draft: FeatureComposerDraft) async { + guard let thread = details[threadID]?.thread ?? snapshot.threads.first(where: { $0.id == threadID }), + rewindingThreadIDs.insert(threadID).inserted else { return } + defer { rewindingThreadIDs.remove(threadID) } + do { + try await draftStore.setDraft(draft, for: FeatureComposerDraftStore.threadKey(thread)) + try await finishRewindRecovery(thread: thread) + } catch { + rewindErrors[threadID] = error.localizedDescription + } + } + + private func finishRewindRecovery(thread: FeatureThread) async throws { + guard snapshot.environments.contains(where: { $0.id == thread.environmentID }) else { return } + let key = FeatureComposerDraftStore.threadKey(thread) + if let recovered = try await draftStore.consumeRewindRecovery(for: key) { + guard snapshot.environments.contains(where: { $0.id == thread.environmentID }) else { + try await draftStore.removeDraft(for: key) + return + } + recoveredRewindDrafts[thread.id] = recovered + if let environmentID = thread.environmentID { + attachmentUploads.syncOwner(draftKey: key, environmentID: environmentID, attachments: recovered.attachments) + } + } + pendingRewindRecoveryIDs.remove(thread.id) + rewindErrors[thread.id] = nil + } + + public func consumeRewindDraft(threadID: String) -> FeatureComposerDraft? { + recoveredRewindDrafts.removeValue(forKey: threadID) + } + public func cancelTurn(threadID: String) async { if pendingSubmissionsByID.values.contains(where: { $0.threadID == threadID && $0.creation != nil @@ -1212,6 +1334,12 @@ public final class FeatureRootModel { private func install(_ value: FeatureSnapshot) { var value = value + let environmentIDs = Set(value.environments.map(\.id)) + for thread in snapshot.threads where thread.environmentID.map({ !environmentIDs.contains($0) }) == true { + recoveredRewindDrafts[thread.id] = nil + pendingRewindRecoveryIDs.remove(thread.id) + rewindErrors[thread.id] = nil + } if settingsWriteTask != nil { // A shell refresh can still contain the settings from before a // queued write. Keep both the visible choice and its rollback point. diff --git a/apps/swift-ios/Features/Shared/FeatureClient.swift b/apps/swift-ios/Features/Shared/FeatureClient.swift index d938c1f47f09..5b791cca92cf 100644 --- a/apps/swift-ios/Features/Shared/FeatureClient.swift +++ b/apps/swift-ios/Features/Shared/FeatureClient.swift @@ -84,6 +84,12 @@ public protocol FeatureClient: AnyObject { context: OrchestrationMessageContext? ) async throws func cancelTurn(threadID: String) async throws + func canRewindConversation(threadID: String, messageID: String) -> Bool + /// Returns only after provider history is rewound. Attachment bytes are copied first. + func rewindConversation( + threadID: String, messageID: String, + prepareRecovery: @MainActor (FeatureRevertedMessage) async throws -> Void + ) async throws func resolveApproval(id: String, decision: FeatureApprovalDecision) async throws func resolveUserInput( id: String, answers: [String: FeatureInputAnswer], @@ -237,6 +243,15 @@ public extension FeatureClient { func setThreadPullRequest(id: String, url: String, linked: Bool) async throws { throw FeatureCapabilityUnavailable("Pull request linking") } + func canRewindConversation(threadID: String, messageID: String) -> Bool { false } + + func rewindConversation( + threadID: String, messageID: String, + prepareRecovery: @MainActor (FeatureRevertedMessage) async throws -> Void + ) async throws { + throw FeatureCapabilityUnavailable("Conversation rewind") + } + func serverPreferences(environmentID: String) async throws -> ServerSettingsSnapshot { throw FeatureCapabilityUnavailable("Server preferences") } diff --git a/apps/swift-ios/Features/Shared/FeatureComposerDraftStore.swift b/apps/swift-ios/Features/Shared/FeatureComposerDraftStore.swift index 413a179591d5..5998a7a2599a 100644 --- a/apps/swift-ios/Features/Shared/FeatureComposerDraftStore.swift +++ b/apps/swift-ios/Features/Shared/FeatureComposerDraftStore.swift @@ -337,18 +337,75 @@ public actor FeatureComposerDraftStore { loadedDrafts = drafts } + public static func rewindRecoveryKey(for threadKey: String) -> String { + "rewind-recovery:" + threadKey + } + + public func hasRewindRecovery(for threadKey: String) throws -> Bool { + try loadIfNeeded()[Self.rewindRecoveryKey(for: threadKey)] != nil + } + + /// Moving recovery into the composer is one disk write. A crash cannot + /// leave the same recovered message available to append a second time. + public func consumeRewindRecovery(for threadKey: String) throws -> FeatureComposerDraft? { + var drafts = try loadIfNeeded() + let recoveryKey = Self.rewindRecoveryKey(for: threadKey) + guard let savedRecovery = drafts[recoveryKey] else { return nil } + let recovery = savedRecovery.featureValue(fileStore: attachmentFileStore) + let current = drafts[threadKey]?.featureValue(fileStore: attachmentFileStore) ?? FeatureComposerDraft() + let files = (recovery.attachments + current.attachments).compactMap(\.ownedFile) + let filesAreReadable = files.allSatisfy { file in + guard FileManager.default.isReadableFile(atPath: file.url.path), + let attributes = try? file.url.resourceValues(forKeys: [.isRegularFileKey, .fileSizeKey]) else { + return false + } + return attributes.isRegularFile == true && attributes.fileSize == file.byteCount + } + guard recovery.attachments.count == savedRecovery.attachments.count, + current.attachments.count == (drafts[threadKey]?.attachments.count ?? 0), + filesAreReadable else { + throw FeatureConversationRewindError(message: "Some saved attachments could not be read. The recovery copy is kept.") + } + guard current.attachments.count + recovery.attachments.count <= 8 else { + throw FeatureConversationRewindError(message: "Make room for the saved prompt's attachments before recovering it.") + } + let recovered = try FeatureConversationRewind.merge(recovery: recovery, into: current) + var persisted = PersistedDraft(recovered) + persisted.importedShareIDs = drafts[threadKey]?.importedShareIDs + drafts[threadKey] = persisted + drafts.removeValue(forKey: recoveryKey) + try persist(drafts) + loadedDrafts = drafts + return recovered + } + + /// Call only when the command was rejected or never sent. Recovery copies + /// have fresh IDs and are not shared with the outbox or existing drafts. + public func discardRewindRecovery(for threadKey: String) throws { + let recoveryKey = Self.rewindRecoveryKey(for: threadKey) + let attachments = try loadIfNeeded()[recoveryKey]?.attachments ?? [] + try removeDraft(for: recoveryKey) + for attachment in attachments { + if let fileName = attachment.ownedFileName { + try? attachmentFileStore.removeOwnedFile(fileName: fileName) + } + } + } + public func removeDrafts( environmentID: String, logicalProjectIDs: Set = [] ) throws { var drafts = try loadIfNeeded() let environmentPrefix = "environment:\(environmentID):" + let rewindPrefix = Self.rewindRecoveryKey(for: environmentPrefix) let questionPrefix = FeatureQuestionAttachmentDraft.key( inputID: FeatureScopedID.input(environmentID: environmentID, wireID: "") ) let logicalKeys = Set(logicalProjectIDs.map(Self.newTaskKey(logicalProjectID:))) drafts = drafts.filter { - !$0.key.hasPrefix(environmentPrefix) && !$0.key.hasPrefix(questionPrefix) + !$0.key.hasPrefix(environmentPrefix) && !$0.key.hasPrefix(rewindPrefix) + && !$0.key.hasPrefix(questionPrefix) && !logicalKeys.contains($0.key) } try persist(drafts) diff --git a/apps/swift-ios/Features/Shared/FeatureConversationRewind.swift b/apps/swift-ios/Features/Shared/FeatureConversationRewind.swift new file mode 100644 index 000000000000..1683fbd34e32 --- /dev/null +++ b/apps/swift-ios/Features/Shared/FeatureConversationRewind.swift @@ -0,0 +1,57 @@ +import Foundation + +public struct FeatureRevertedMessage: Sendable { + public let message: FeatureMessage + public let attachments: [FeatureDraftAttachment] + + public init(message: FeatureMessage, attachments: [FeatureDraftAttachment]) { + self.message = message + self.attachments = attachments + } +} + +struct FeatureConversationRewindError: LocalizedError { + let message: String + var didNotRevert = false + var errorDescription: String? { message } +} + +enum FeatureConversationRewind { + static func canStart(in detail: FeatureThreadDetail) -> Bool { + switch detail.thread.state { + case .idle, .completed, .failed: + detail.isCompacting != true && !detail.backgroundWorkIsActive + && detail.approvals.isEmpty && detail.userInputs.isEmpty + default: + false + } + } + + /// Append recovered input without changing the current model or workspace. + static func recover(_ reverted: FeatureRevertedMessage, draft: FeatureComposerDraft) throws -> FeatureComposerDraft { + let original = reverted.message.text + let prompt = !reverted.attachments.isEmpty + && original == "[User attached one or more files without additional text. Respond using the conversation context and the attached files.]" + ? "" : original + // Copies keep the original attachment order but receive fresh local IDs. + // Keep the context links stable and point their records at those copies. + let attachmentIDs = Dictionary(zip(reverted.message.attachments, reverted.attachments).map { original, copy in + (original.id, copy.id.uuidString) + }, uniquingKeysWith: { first, _ in first }) + let context = ComposerContextReferences.rebind(reverted.message.context, attachmentIDs: attachmentIDs) + return try merge( + recovery: FeatureComposerDraft(text: prompt, attachments: reverted.attachments, context: context), + into: draft + ) + } + + static func merge(recovery: FeatureComposerDraft, into draft: FeatureComposerDraft) throws -> FeatureComposerDraft { + var recovered = draft + if !recovery.text.isEmpty { + recovered.text = draft.text.isEmpty ? recovery.text : draft.text + "\n\n" + recovery.text + } + recovered.attachments.append(contentsOf: recovery.attachments) + recovered.context = try FeatureComposerContext.merge(draft.context, recovery.context) + return recovered + } +} diff --git a/apps/swift-ios/Tests/FeatureTests/FeatureRootModelTests.swift b/apps/swift-ios/Tests/FeatureTests/FeatureRootModelTests.swift index dc38067eefc1..241354476b03 100644 --- a/apps/swift-ios/Tests/FeatureTests/FeatureRootModelTests.swift +++ b/apps/swift-ios/Tests/FeatureTests/FeatureRootModelTests.swift @@ -9,6 +9,212 @@ import XCTest @MainActor @Suite("Feature root model") struct FeatureRootModelTests { + @Test + func rewindLocksSendingAndSavesRecoveredInputAfterLeavingTheThread() async throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: directory) } + let drafts = FeatureComposerDraftStore(fileURL: directory.appendingPathComponent("drafts.json")) + let outbox = FeatureOutboxStore(fileURL: directory.appendingPathComponent("outbox.json")) + let thread = FeatureThread(id: "rewind-thread", projectID: "project", environmentID: "environment", title: "Task") + let message = FeatureMessage(id: "original", role: .user, text: "Original prompt") + let client = FeatureClientStub() + client.snapshot = FeatureSnapshot( + environments: [.init(id: "environment", name: "Computer", endpoint: "https://example.test", connectionState: .connected)], + threads: [thread] + ) + client.threadDetail = .init(thread: thread, messages: [message]) + let model = FeatureRootModel(client: client, outboxStore: outbox, draftStore: drafts) + await model.reload() + _ = await model.detail(for: thread.id) + let selection = FeatureSelection(providerID: "chosen-provider", modelID: "chosen-model") + let draft = FeatureComposerDraft(text: "Existing draft", selection: selection) + client.rewindHandler = { threadID, messageID in + #expect(threadID == thread.id) + #expect(messageID == message.id) + #expect(model.rewindingThreadIDs.contains(thread.id)) + let savedDraft = try await drafts.draft(for: FeatureComposerDraftStore.threadKey(thread)) + let savedRecovery = try await drafts.draft(for: FeatureComposerDraftStore.rewindRecoveryKey( + for: FeatureComposerDraftStore.threadKey(thread) + )) + #expect(savedDraft == draft) + #expect(savedRecovery?.text == "Original prompt") + #expect(!(await model.sendMessage(.init(threadID: thread.id, text: "Do not send", selection: nil, attachments: [])))) + #expect(client.sendMessageCallCount == 0) + model.releaseThread(thread.id) + return FeatureRevertedMessage(message: message, attachments: []) + } + + await model.rewindConversation(threadID: thread.id, messageID: message.id, draft: draft) + + let saved = try await drafts.draft(for: FeatureComposerDraftStore.threadKey(thread)) + #expect(saved?.text == "Existing draft\n\nOriginal prompt") + #expect(saved?.selection == selection) + #expect(model.recoveredRewindDrafts[thread.id] == saved) + #expect(model.rewindingThreadIDs.isEmpty) + #expect(try await outbox.submissions().isEmpty) + #expect(model.rewindErrors[thread.id] == nil) + } + + @Test + func rejectedRewindKeepsTheDraftAndReportsTheServerError() async throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: directory) } + let drafts = FeatureComposerDraftStore(fileURL: directory.appendingPathComponent("drafts.json")) + let thread = FeatureThread(id: "thread", projectID: "project", environmentID: "environment", title: "Task") + let client = FeatureClientStub() + client.snapshot = FeatureSnapshot( + environments: [.init(id: "environment", name: "Computer", endpoint: "https://example.test", connectionState: .connected)], + threads: [thread] + ) + client.threadDetail = .init(thread: thread, messages: [.init(id: "user", role: .user, text: "Original")]) + client.rewindHandler = { _, _ in + throw FeatureConversationRewindError(message: "Unknown command thread.conversation.revert", didNotRevert: true) + } + let model = FeatureRootModel( + client: client, outboxStore: .init(fileURL: directory.appendingPathComponent("outbox.json")), draftStore: drafts + ) + await model.reload() + _ = await model.detail(for: thread.id) + let draft = FeatureComposerDraft(text: "Keep this draft") + + await model.rewindConversation(threadID: thread.id, messageID: "user", draft: draft) + + #expect(try await drafts.draft(for: FeatureComposerDraftStore.threadKey(thread)) == draft) + #expect(model.details[thread.id]?.messages.map(\.id) == ["user"]) + #expect(model.rewindErrors[thread.id] == "Unknown command thread.conversation.revert") + #expect(model.recoveredRewindDrafts.isEmpty) + #expect(model.rewindingThreadIDs.isEmpty) + #expect(try await drafts.draft(for: FeatureComposerDraftStore.rewindRecoveryKey( + for: FeatureComposerDraftStore.threadKey(thread) + )) == nil) + } + + @Test + func lostRewindReceiptKeepsPromptAndAttachmentBytesAcrossRestart() async throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: directory) } + let draftURL = directory.appendingPathComponent("drafts.json") + let drafts = FeatureComposerDraftStore(fileURL: draftURL) + let thread = FeatureThread(id: "thread", projectID: "project", environmentID: "environment", title: "Task") + let attachment = FeatureDraftAttachment(data: Data([4, 5, 6]), filename: "saved.txt", mimeType: "text/plain", source: .pastedText) + let fileContext = ComposerContextRecord(contextId: "saved-file", label: "Pasted text", payload: .file(.init( + attachmentId: "server-file", name: "saved.txt", mimeType: "text/plain", sizeBytes: 3 + ))) + let originalPrompt = "Original prompt " + ComposerContextReferences.format(fileContext) + let client = FeatureClientStub() + client.snapshot = FeatureSnapshot( + environments: [.init(id: "environment", name: "Computer", endpoint: "https://example.test", connectionState: .connected)], + threads: [thread] + ) + client.threadDetail = .init(thread: thread, messages: [.init( + id: "user", role: .user, text: originalPrompt, attachments: [.init( + id: "server-file", name: "saved.txt", mimeType: "text/plain", sizeBytes: 3, source: .pastedText + )], context: .init(records: [fileContext]) + )]) + client.rewindAttachments = [attachment] + client.rewindHandler = { _, _ in throw RPCError.disconnected } + let model = FeatureRootModel( + client: client, outboxStore: .init(fileURL: directory.appendingPathComponent("outbox.json")), draftStore: drafts + ) + await model.reload() + _ = await model.detail(for: thread.id) + await model.rewindConversation(threadID: thread.id, messageID: "user", draft: .init(text: "Current draft")) + let key = FeatureComposerDraftStore.threadKey(thread) + let recoveryKey = FeatureComposerDraftStore.rewindRecoveryKey(for: key) + #expect(try await drafts.draft(for: key)?.text == "Current draft") + #expect(try await drafts.draft(for: recoveryKey)?.attachments.first?.data == Data([4, 5, 6])) + #expect(try await drafts.draft(for: recoveryKey)?.context?.records.first?.attachment?.attachmentId == attachment.id.uuidString) + #expect(model.pendingRewindRecoveryIDs.contains(thread.id)) + #expect(!model.canRewindConversation(threadID: thread.id, messageID: "user")) + + let restartedDrafts = FeatureComposerDraftStore(fileURL: draftURL) + let restarted = FeatureRootModel( + client: client, outboxStore: .init(fileURL: directory.appendingPathComponent("restart-outbox.json")), + draftStore: restartedDrafts + ) + await restarted.reload() + await restarted.checkRewindRecovery(for: thread) + #expect(restarted.pendingRewindRecoveryIDs.contains(thread.id)) + let newContext = ComposerContextRecord(contextId: "new-source", label: "New source", payload: .mention(.init(path: "src/new.swift"))) + let newPrompt = "Newer draft " + ComposerContextReferences.format(newContext) + await restarted.recoverSavedRewind(threadID: thread.id, draft: .init(text: newPrompt, context: .init(records: [newContext]))) + let recovered = try await restartedDrafts.draft(for: key) + #expect(recovered?.text == newPrompt + "\n\n" + originalPrompt) + #expect(recovered?.attachments.first?.data == Data([4, 5, 6])) + #expect(recovered?.attachments.first?.uploadedReference == nil) + #expect(recovered?.attachments.first?.source == .pastedText) + #expect(recovered?.context?.records.map(\.contextId) == [newContext.contextId, fileContext.contextId]) + #expect(recovered?.context?.records.last?.attachment?.attachmentId == attachment.id.uuidString) + #expect(try await restartedDrafts.draft(for: recoveryKey) == nil) + #expect(try await restartedDrafts.consumeRewindRecovery(for: key) == nil) + #expect(restarted.pendingRewindRecoveryIDs.isEmpty) + } + + @Test + func rewindRejectsContextOverflowBeforeDispatch() async throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: directory) } + let drafts = FeatureComposerDraftStore(fileURL: directory.appendingPathComponent("drafts.json")) + let thread = FeatureThread(id: "thread", projectID: "project", environmentID: "environment", title: "Task") + let record = ComposerContextRecord(contextId: "recovered", label: "Saved", payload: .mention(.init(path: "saved"))) + let message = FeatureMessage(id: "user", role: .user, text: ComposerContextReferences.format(record), context: .init(records: [record])) + let client = FeatureClientStub() + client.snapshot = FeatureSnapshot( + environments: [.init(id: "environment", name: "Computer", endpoint: "https://example.test", connectionState: .connected)], + threads: [thread] + ) + client.threadDetail = .init(thread: thread, messages: [message]) + var dispatched = false + client.rewindHandler = { _, _ in + dispatched = true + return FeatureRevertedMessage(message: message, attachments: []) + } + let model = FeatureRootModel( + client: client, outboxStore: .init(fileURL: directory.appendingPathComponent("outbox.json")), draftStore: drafts + ) + await model.reload() + _ = await model.detail(for: thread.id) + let current = FeatureComposerDraft(text: "Current draft", context: .init(records: (0..<200).map { + ComposerContextRecord(contextId: "existing-\($0)", label: "Source \($0)", payload: .mention(.init(path: "source-\($0)"))) + })) + await model.rewindConversation(threadID: thread.id, messageID: message.id, draft: current) + #expect(!dispatched) + #expect(model.rewindErrors[thread.id] == FeatureComposerContext.MergeError.tooManyRecords.localizedDescription) + #expect(try await drafts.draft(for: FeatureComposerDraftStore.threadKey(thread)) == current) + #expect(try await !drafts.hasRewindRecovery(for: FeatureComposerDraftStore.threadKey(thread))) + } + + @Test + func rewindDoesNotPublishRecoveredContentAfterItsEnvironmentWasRemoved() async throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: directory) } + let drafts = FeatureComposerDraftStore(fileURL: directory.appendingPathComponent("drafts.json")) + let thread = FeatureThread(id: "thread", projectID: "project", environmentID: "environment", title: "Task") + let message = FeatureMessage(id: "user", role: .user, text: "Original prompt") + let client = FeatureClientStub() + client.snapshot = FeatureSnapshot( + environments: [.init(id: "environment", name: "Computer", endpoint: "https://example.test", connectionState: .connected)], + threads: [thread] + ) + client.threadDetail = .init(thread: thread, messages: [message]) + let model = FeatureRootModel( + client: client, outboxStore: .init(fileURL: directory.appendingPathComponent("outbox.json")), draftStore: drafts + ) + await model.reload() + _ = await model.detail(for: thread.id) + client.rewindHandler = { _, _ in + try await drafts.removeDrafts(environmentID: "environment") + client.snapshot = FeatureSnapshot() + await model.reload() + return FeatureRevertedMessage(message: message, attachments: []) + } + await model.rewindConversation(threadID: thread.id, messageID: "user", draft: .init(text: "Current draft")) + #expect(try await drafts.draft(for: FeatureComposerDraftStore.threadKey(thread)) == nil) + #expect(try await !drafts.hasRewindRecovery(for: FeatureComposerDraftStore.threadKey(thread))) + #expect(model.recoveredRewindDrafts.isEmpty) + #expect(model.pendingRewindRecoveryIDs.isEmpty) + } + @Test func transcriptSkillPillsUseTheThreadWorkspaceCatalog() async { let skill = FeatureProviderSkill(name: "project-only", displayName: "Project only") @@ -3552,6 +3758,20 @@ private func orchestrationThread( @MainActor private final class FeatureClientStub: FeatureClient, T3ConnectCapable { + var rewindHandler: ((String, String) async throws -> FeatureRevertedMessage)? + var rewindAttachments: [FeatureDraftAttachment] = [] + func canRewindConversation(threadID: String, messageID: String) -> Bool { rewindHandler != nil } + func rewindConversation( + threadID: String, messageID: String, + prepareRecovery: @MainActor (FeatureRevertedMessage) async throws -> Void + ) async throws { + guard let rewindHandler, + let message = threadDetail?.messages.first(where: { $0.id == messageID }) else { + throw FeatureCapabilityUnavailable("Conversation rewind") + } + try await prepareRecovery(.init(message: message, attachments: rewindAttachments)) + _ = try await rewindHandler(threadID, messageID) + } var foregroundReconnects: [Bool] = [] func resumeAfterBackground(reconnect: Bool) async { foregroundReconnects.append(reconnect) } private let eventStream: AsyncStream diff --git a/apps/swift-ios/Tests/FeatureTests/NativeConversationRewindTests.swift b/apps/swift-ios/Tests/FeatureTests/NativeConversationRewindTests.swift new file mode 100644 index 000000000000..aebb8ffab6cb --- /dev/null +++ b/apps/swift-ios/Tests/FeatureTests/NativeConversationRewindTests.swift @@ -0,0 +1,246 @@ +import Foundation +import Testing +@testable import T3Code + +@Suite("Conversation rewind") +struct NativeConversationRewindTests { + @Test + func keepFilesUsesACommandOlderServersCannotTreatAsFileRestore() { + let command = OrchestrationCommands.revertConversation( + threadID: "thread", turnCount: 4, commandID: "command", createdAt: "2026-09-13T00:00:00Z" + ) + #expect(command == .object([ + "type": .string("thread.conversation.revert"), + "threadId": .string("thread"), "turnCount": .number(4), + "commandId": .string("command"), "createdAt": .string("2026-09-13T00:00:00Z"), + ])) + } + + @Test + func paginatedAndSteeredMessagesUseCheckpointBoundaries() { + let thread = thread(messages: [ + message("user", role: "user"), message("steering", role: "user"), + message("assistant", role: "assistant"), + ], checkpoints: [.init( + turnId: "turn", checkpointTurnCount: 8, checkpointRef: "ref", + status: "ready", files: [], assistantMessageId: "assistant", completedAt: "2026-09-13T00:00:00Z" + )]) + #expect(NativeConversationRewind.turnCount(before: "steering", in: thread) == 7) + #expect(NativeConversationRewind.turnCount(before: "user", in: thread) == nil) + #expect(NativeConversationRewind.turnCount(before: "assistant", in: thread) == nil) + #expect(NativeConversationRewind.turnCount(before: "missing", in: thread) == nil) + } + + @Test + func draftRecoveryKeepsExistingInputAndSameNameAttachments() throws { + let existing = FeatureDraftAttachment(data: Data([1]), filename: "same.txt", mimeType: "text/plain") + let recovered = FeatureDraftAttachment(data: Data([2]), filename: "same.txt", mimeType: "text/plain") + let draft = FeatureComposerDraft( + text: "Current draft", attachments: [existing], + selection: .init(providerID: "selected-provider", modelID: "selected-model"), + workspace: .init(mode: .worktree, branch: "feature", worktreePath: "/worktree", startFromOrigin: true) + ) + let result = try FeatureConversationRewind.recover(.init( + message: .init(id: "message", role: .user, text: "Original prompt"), attachments: [recovered] + ), draft: draft) + #expect(result.text == "Current draft\n\nOriginal prompt") + #expect(result.attachments.map(\.id) == [existing.id, recovered.id]) + #expect(result.attachments.map(\.data) == [Data([1]), Data([2])]) + #expect(result.selection == draft.selection) + #expect(result.workspace == draft.workspace) + } + + @Test + func recoveredContextKeepsItsLinksAndUsesNewAttachmentIDsOnResend() throws { + let existing = ComposerContextRecord(contextId: "existing", label: "Sources", payload: .mention(.init(path: "src"))) + let shared = ComposerContextRecord(contextId: "shared", label: "Build", payload: .skill(.init(name: "build"))) + let file = ComposerContextRecord(contextId: "file", label: "Pasted text", payload: .file(.init( + attachmentId: "old-server-file", name: "paste.txt", mimeType: "text/plain", sizeBytes: 3 + ))) + let terminal = FeatureComposerContext.terminalRecord(text: "Original output", terminalID: "terminal", label: "Terminal") + let future = ComposerContextRecord(contextId: "future", label: "Captured input", payload: .unknown( + kind: "future", payload: .object(["value": .string("Keep this")]) + )) + let restoredRecords = [shared, file, terminal, future] + let prompt = restoredRecords.map(ComposerContextReferences.format).joined(separator: " ") + let copied = FeatureDraftAttachment(data: Data([1, 2, 3]), filename: "paste.txt", mimeType: "text/plain", source: .pastedText) + let result = try FeatureConversationRewind.recover(.init( + message: .init(id: "user", role: .user, text: prompt, attachments: [.init( + id: "old-server-file", name: "paste.txt", mimeType: "text/plain", sizeBytes: 3, source: .pastedText + )], context: .init(records: restoredRecords)), + attachments: [copied] + ), draft: .init(text: ComposerContextReferences.format(existing), context: .init(records: [existing, shared]))) + + #expect(result.text.hasSuffix(prompt)) + #expect(result.context?.records.map(\.contextId) == ["existing", "shared", "file", terminal.contextId, "future"]) + #expect(result.context?.records.first(where: { $0.contextId == "file" })?.attachment?.attachmentId == copied.id.uuidString) + #expect(result.context?.records.last == future) + #expect(result.attachments.first?.source == .pastedText) + + let upload = try UploadChatAttachment( + id: copied.id, data: copied.data, name: copied.filename, mimeType: copied.mimeType, contextSource: copied.source + ) + let resent = T3Client.prepareMessageContext( + text: result.text + " Edit this", context: result.context, attachments: [upload], + uploadedAttachments: [.object(["id": .string("new-server-file")])], supportsContext: true + ) + #expect(resent.context?.records.count == result.context?.records.count) + #expect(resent.context?.records.first(where: { $0.contextId == "file" })?.attachment?.attachmentId == "new-server-file") + #expect(resent.text == result.text + " Edit this") + #expect(upload.contextSource == .pastedText) + } + + @Test + func busyThreadsCannotRewind() { + for state in [FeatureThreadState.working, .queued, .monitoring, .waitingForApproval, .waitingForInput] { + let detail = FeatureThreadDetail(thread: .init(id: "thread", projectID: "project", title: "Task", state: state)) + #expect(!FeatureConversationRewind.canStart(in: detail)) + } + } + + @Test + func attachmentOnlyRewindDoesNotRestoreGeneratedBootstrapText() throws { + let result = try FeatureConversationRewind.recover(.init( + message: .init(id: "user", role: .user, text: "[User attached one or more files without additional text. Respond using the conversation context and the attached files.]"), + attachments: [.init(data: Data([1]), filename: "input.txt", mimeType: "text/plain")] + ), draft: .init(text: "Existing draft")) + #expect(result.text == "Existing draft") + #expect(result.attachments.count == 1) + } + + @Test + func literalBootstrapSentenceWithoutAttachmentsIsPreserved() throws { + let message = FeatureMessage(id: "user", role: .user, text: "[User attached one or more files without additional text. Respond using the conversation context and the attached files.]") + let result = try FeatureConversationRewind.recover(.init(message: message, attachments: []), draft: .init()) + #expect(result.text == message.text) + } + + @Test(arguments: [0, 2, 3]) + func recoveryRequiresReadableFilesWithTheSavedSize(actualBytes: Int) async throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: directory) } + let attachmentRoot = directory.appendingPathComponent("owned") + let fileName = UUID().uuidString + ".txt" + let fileURL = attachmentRoot.appendingPathComponent(fileName) + if actualBytes > 0 { + try FileManager.default.createDirectory(at: attachmentRoot, withIntermediateDirectories: true) + try Data(repeating: 1, count: actualBytes).write(to: fileURL) + } + let store = FeatureComposerDraftStore( + fileURL: directory.appendingPathComponent("drafts.json"), + attachmentStorageRootURL: attachmentRoot + ) + let recoveryKey = FeatureComposerDraftStore.rewindRecoveryKey(for: "thread") + try await store.setDraft(.init(attachments: [.init( + ownedFile: .init(fileName: fileName, url: fileURL, byteCount: 3), + filename: "input.txt", mimeType: "text/plain" + )]), for: recoveryKey) + do { + let recovered = try await store.consumeRewindRecovery(for: "thread") + #expect(actualBytes == 3) + #expect(recovered?.attachments.first?.ownedFile?.url == fileURL) + } catch { + #expect(actualBytes != 3) + #expect(error.localizedDescription.contains("recovery copy is kept")) + } + #expect(try await store.hasRewindRecovery(for: "thread") == (actualBytes != 3)) + #expect(try await (store.draft(for: "thread") == nil) == (actualBytes != 3)) + } + + @Test + func recoveryKeysDoNotOverwriteAnotherThreadsDraft() async throws { + let directory = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) + defer { try? FileManager.default.removeItem(at: directory) } + let store = FeatureComposerDraftStore(fileURL: directory.appendingPathComponent("drafts.json")) + let key = "environment:one:thread:foo" + let otherKey = key + ":rewind-recovery" + try await store.setDraft(.init(text: "Other thread's draft"), for: otherKey) + try await store.setDraft(.init(text: "Recovered prompt"), for: FeatureComposerDraftStore.rewindRecoveryKey(for: key)) + let recovered = try await store.consumeRewindRecovery(for: key) + #expect(recovered?.text == "Recovered prompt") + #expect(try await store.draft(for: otherKey)?.text == "Other thread's draft") + } + + @Test + func completionIgnoresOldAndUnrelatedEvents() async throws { + let stream = AsyncThrowingStream<[ThreadStreamItem], Error>.makeStream() + stream.continuation.yield([ + .event(event("thread.reverted", sequence: 10, threadID: "thread", turnCount: 0)), + .event(event("thread.reverted", sequence: 11, threadID: "other", turnCount: 0)), + .event(event("thread.reverted", sequence: 12, threadID: "thread", turnCount: 2)), + ]) + stream.continuation.finish() + do { + try await wait(stream.stream) + Issue.record("Unrelated events must not complete this rewind") + } catch { + #expect(error.localizedDescription.contains("connection closed")) + } + } + + @Test + func completionAcceptsTheRevertedEvent() async throws { + let stream = AsyncThrowingStream<[ThreadStreamItem], Error>.makeStream() + stream.continuation.yield([.event(event("thread.reverted", sequence: 11, threadID: "thread", turnCount: 0))]) + try await wait(stream.stream) + } + + @Test + func providerFailureDoesNotBecomeSuccessfulRecovery() async throws { + let stream = AsyncThrowingStream<[ThreadStreamItem], Error>.makeStream() + stream.continuation.yield([.event(.object([ + "type": .string("thread.activity-appended"), "sequence": .number(11), + "payload": .object([ + "threadId": .string("thread"), + "activity": .object([ + "kind": .string("checkpoint.revert.failed"), + "payload": .object(["detail": .string("History boundary is unavailable"), "turnCount": .number(0)]), + ]), + ]), + ]))]) + do { + try await wait(stream.stream) + Issue.record("Provider failure must reject the rewind") + } catch { + #expect(error.localizedDescription == "History boundary is unavailable") + } + } + + @Test + func replacementSnapshotConfirmsHistoryWasRemoved() async throws { + let stream = AsyncThrowingStream<[ThreadStreamItem], Error>.makeStream() + stream.continuation.yield([.snapshot(.init(snapshotSequence: 11, thread: thread(), page: nil))]) + try await wait(stream.stream) + } + + private func wait(_ events: AsyncThrowingStream<[ThreadStreamItem], Error>) async throws { + _ = try await NativeConversationRewind.waitForCompletion( + batches: events, threadID: "thread", messageID: "user", turnCount: 0, + afterSequence: 10, previousFailureIDs: [] + ) + } + + private func event(_ type: String, sequence: Int, threadID: String, turnCount: Int) -> JSONValue { + .object([ + "type": .string(type), "sequence": .number(Double(sequence)), + "payload": .object(["threadId": .string(threadID), "turnCount": .number(Double(turnCount))]), + ]) + } + + private func message(_ id: String, role: String) -> OrchestrationMessage { + .init(id: id, role: role, text: id, attachments: [], turnId: "turn", streaming: false, + createdAt: "2026-09-13T00:00:00Z", updatedAt: "2026-09-13T00:00:00Z") + } + + private func thread(messages: [OrchestrationMessage] = [], checkpoints: [CheckpointSummary] = []) -> OrchestrationThread { + .init( + id: "thread", projectId: "project", title: "Task", + modelSelection: .init(instanceId: "codex", model: "test-model"), + runtimeMode: .fullAccess, interactionMode: .default, branch: nil, worktreePath: nil, + latestTurn: nil, createdAt: "2026-09-13T00:00:00Z", updatedAt: "2026-09-13T00:00:00Z", + archivedAt: nil, settledOverride: nil, settledAt: nil, snoozedUntil: nil, + snoozedAt: nil, pinnedAt: nil, deletedAt: nil, messages: messages, activities: [], + checkpoints: checkpoints, session: nil + ) + } +}