Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions Wisp/ViewModels/ChatViewModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ final class ChatViewModel {
private var usedResume = false
var queuedPrompt: String?
var queuedAttachments: [AttachedFile] = []
var stashedDraft: String?
private var retriedAfterTimeout = false
private var turnHasMutations = false
private var pendingForkContext: String?
Expand Down Expand Up @@ -210,6 +211,19 @@ final class ChatViewModel {
try? modelContext.save()
}

func stashDraft() {
let text = inputText.trimmingCharacters(in: .whitespacesAndNewlines)
guard !text.isEmpty else { return }
stashedDraft = text
inputText = ""
}

private func restoreStash() {
guard let stash = stashedDraft else { return }
stashedDraft = nil
inputText = stash
}

func fetchRemoteSessions(apiClient: SpritesAPIClient, existingSessionIds: Set<String>) {
guard !isLoadingRemoteSessions else { return }
isLoadingRemoteSessions = true
Expand Down Expand Up @@ -479,12 +493,14 @@ final class ChatViewModel {
queuedPrompt = text
queuedAttachments = attachedFiles
attachedFiles = []
restoreStash()
return
}

// Build prompt with attached file paths prepended
let prompt = buildPrompt(text: text, attachments: attachedFiles)
attachedFiles = []
restoreStash()

let isFirstMessage = messages.isEmpty
let userMessage = ChatMessage(role: .user, content: [.text(prompt)])
Expand Down
8 changes: 8 additions & 0 deletions Wisp/Views/SpriteDetail/Chat/ChatInputBar.swift
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ struct ChatInputBar: View {
var attachedFiles: [AttachedFile] = []
var onRemoveAttachment: ((AttachedFile) -> Void)? = nil
var lastUploadedFileName: String? = nil
var onStash: (() -> Void)? = nil
var isFocused: FocusState<Bool>.Binding

@State private var showStopConfirmation = false
Expand Down Expand Up @@ -93,6 +94,13 @@ struct ChatInputBar: View {
.tint(isEmpty || hasQueuedMessage ? .gray : Color("AccentColor"))
.disabled(isEmpty || hasQueuedMessage)
.buttonStyle(.glass)
.contextMenu {
if let onStash, !isEmpty {
Button("Stash Draft", systemImage: "tray.and.arrow.down") {
onStash()
}
}
}
}
}
.animation(.easeInOut(duration: 0.2), value: attachedFiles.count)
Expand Down
15 changes: 15 additions & 0 deletions Wisp/Views/SpriteDetail/Chat/ChatView.swift
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ struct ChatView: View {
viewModel.attachedFiles.removeAll { $0.id == file.id }
},
lastUploadedFileName: viewModel.lastUploadedFileName,
onStash: { viewModel.stashDraft() },
isFocused: $isInputFocused
)
}
Expand Down Expand Up @@ -244,3 +245,17 @@ struct ChatView: View {
.background(.bar)
}
}

#Preview {
let viewModel = ChatViewModel(
spriteName: "my-sprite",
chatId: UUID(),
currentServiceName: nil,
workingDirectory: "/home/sprite/project"
)
NavigationStack {
ChatView(viewModel: viewModel)
.environment(SpritesAPIClient())
.modelContainer(for: [SpriteChat.self, SpriteSession.self], inMemory: true)
}
}
56 changes: 56 additions & 0 deletions WispTests/ChatViewModelTests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -601,4 +601,60 @@ struct ChatViewModelTests {

#expect(vm.currentAssistantMessageId == nil)
}

// MARK: - stashDraft

@Test func stashDraft_movesInputTextToStash() throws {
let ctx = try makeModelContext()
let (vm, _) = makeChatViewModel(modelContext: ctx)

vm.inputText = "my long prompt"
vm.stashDraft()

#expect(vm.stashedDraft == "my long prompt")
#expect(vm.inputText == "")
}

@Test func stashDraft_doesNothingWhenEmpty() throws {
let ctx = try makeModelContext()
let (vm, _) = makeChatViewModel(modelContext: ctx)

vm.inputText = " "
vm.stashDraft()

#expect(vm.stashedDraft == nil)
#expect(vm.inputText == " ")
}

@Test func stashDraft_overwritesPreviousStash() throws {
let ctx = try makeModelContext()
let (vm, _) = makeChatViewModel(modelContext: ctx)

vm.inputText = "first draft"
vm.stashDraft()
vm.inputText = "second draft"
vm.stashDraft()

#expect(vm.stashedDraft == "second draft")
#expect(vm.inputText == "")
}

@Test func stashDraft_leavesInputReadyForNextMessage() throws {
let ctx = try makeModelContext()
let (vm, _) = makeChatViewModel(modelContext: ctx)

vm.inputText = "long prompt I want to come back to"
vm.stashDraft()

// After stashing, the field is clear and stash holds the draft
#expect(vm.inputText == "")
#expect(vm.stashedDraft == "long prompt I want to come back to")

// Manually simulate the restore (as sendMessage would do)
vm.inputText = vm.stashedDraft!
vm.stashedDraft = nil

#expect(vm.inputText == "long prompt I want to come back to")
#expect(vm.stashedDraft == nil)
}
}
Loading