220 lines
13 KiB
Swift
220 lines
13 KiB
Swift
import CryptoKit
|
|
import Foundation
|
|
import Observation
|
|
|
|
struct AgentDraftItem: Codable, Equatable, Identifiable, Sendable {
|
|
enum State: String, Codable, Sendable { case loading, uploading, ready, failed }
|
|
var id: String = UUID().uuidString // Stable multipart uploadId, including across relaunch.
|
|
var filename: String
|
|
var mimeType: String?
|
|
var sizeBytes: Int?
|
|
var state: State = .loading
|
|
var attachment: AgentAttachment?
|
|
var issue: AgentAttachmentIssue?
|
|
}
|
|
struct AgentDraftSnapshot: Codable, Sendable {
|
|
var caption = ""
|
|
var items: [AgentDraftItem] = []
|
|
var id: String? = nil // Optional for drafts saved before durable outbox handoff tracking.
|
|
}
|
|
struct AgentDraftStorage: Sendable {
|
|
var read: @Sendable () throws -> AgentDraftSnapshot
|
|
var write: @Sendable (AgentDraftSnapshot) throws -> Void
|
|
var stage: @Sendable (String, Data) throws -> Void
|
|
var bytes: @Sendable (String) throws -> Data
|
|
var remove: @Sendable (String) -> Void
|
|
|
|
static func file(baseURL: URL, root: URL = .applicationSupportDirectory) -> Self {
|
|
let key = SHA256.hash(data: Data(baseURL.absoluteString.utf8)).map { String(format: "%02x", $0) }.joined()
|
|
let directory = root.appendingPathComponent("AgentAttachmentDrafts/\(key)", isDirectory: true)
|
|
let manifest = directory.appendingPathComponent("draft.json")
|
|
let file: @Sendable (String) throws -> URL = { id in
|
|
guard UUID(uuidString: id) != nil else { throw AgentAttachmentIssue.local("Invalid staged upload identity.") }
|
|
return directory.appendingPathComponent(id + ".bytes")
|
|
}
|
|
let write: @Sendable (Data, URL) throws -> Void = { data, destination in
|
|
try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
|
|
try data.write(to: destination, options: [.atomic, .completeFileProtectionUntilFirstUserAuthentication])
|
|
}
|
|
return Self(read: {
|
|
guard FileManager.default.fileExists(atPath: manifest.path) else { return AgentDraftSnapshot() }
|
|
let data = try Data(contentsOf: manifest)
|
|
guard data.count <= 2 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The saved attachment draft is too large.") }
|
|
let snapshot = try JSONDecoder().decode(AgentDraftSnapshot.self, from: data)
|
|
guard snapshot.items.count <= 8, Set(snapshot.items.map(\.id)).count == snapshot.items.count,
|
|
snapshot.items.allSatisfy({ UUID(uuidString: $0.id) != nil }) else { throw AgentAttachmentIssue.local("The saved attachment draft is invalid.") }
|
|
return snapshot
|
|
}, write: { snapshot in
|
|
let data = try JSONEncoder().encode(snapshot)
|
|
guard data.count <= 2 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The saved attachment draft is too large.") }
|
|
try write(data, manifest)
|
|
}, stage: { id, data in
|
|
guard data.count <= 8 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The staged file exceeds the size limit.") }
|
|
try write(data, file(id))
|
|
}, bytes: { id in
|
|
let url = try file(id)
|
|
let size = try url.resourceValues(forKeys: [.fileSizeKey]).fileSize ?? Int.max
|
|
guard size <= 8 * 1024 * 1024 else { throw AgentAttachmentIssue.local("The staged file exceeds the size limit.") }
|
|
return try Data(contentsOf: url)
|
|
}, remove: { id in if let url = try? file(id) { try? FileManager.default.removeItem(at: url) } })
|
|
}
|
|
}
|
|
|
|
@MainActor @Observable
|
|
final class AgentAttachmentDraft {
|
|
private(set) var id = UUID().uuidString
|
|
var caption = "" { didSet { save() } }
|
|
private(set) var items: [AgentDraftItem] = []
|
|
private(set) var capability: AgentAttachmentCapabilities?
|
|
private(set) var error: String?
|
|
private(set) var persistenceAvailable = true
|
|
@ObservationIgnored private var client: (any AgentAPIClienting)?
|
|
@ObservationIgnored private var storage: AgentDraftStorage?
|
|
@ObservationIgnored private var prepared: [String: AgentPreparedAttachment] = [:]
|
|
@ObservationIgnored private var loaders: [String: @Sendable () async throws -> AgentPreparedAttachment] = [:]
|
|
@ObservationIgnored private var worker: Task<Void, Never>?
|
|
@ObservationIgnored private var generation = UUID()
|
|
@ObservationIgnored private var restoring = false
|
|
@ObservationIgnored private var active = true
|
|
var supportsAttachments: Bool { capability?.supported == true && persistenceAvailable }
|
|
var ready: [AgentAttachment] { items.compactMap(\.attachment) }
|
|
var canSend: Bool {
|
|
caption.utf16.count <= 40_000 && (!caption.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty || !items.isEmpty)
|
|
&& (items.isEmpty || (supportsAttachments && items.allSatisfy { $0.state == .ready && $0.attachment != nil }))
|
|
&& persistenceAvailable
|
|
}
|
|
|
|
func configure(client: any AgentAPIClienting, storage: AgentDraftStorage?) {
|
|
suspend()
|
|
self.client = client; self.storage = storage
|
|
prepared = [:]; loaders = [:]; capability = nil; error = nil; persistenceAvailable = true
|
|
restoring = true
|
|
defer { restoring = false; active = true }
|
|
do {
|
|
let saved = try storage?.read() ?? AgentDraftSnapshot()
|
|
id = saved.id ?? UUID().uuidString
|
|
caption = saved.caption
|
|
items = saved.items.map { item in
|
|
var item = item
|
|
if item.state != .ready {
|
|
item.state = .failed
|
|
item.issue = .local("Upload interrupted. Retry with the same upload ID. If selection did not finish, remove and select the file again.", code: "upload_interrupted")
|
|
}
|
|
return item
|
|
}
|
|
} catch {
|
|
caption = ""; items = []; persistenceAvailable = false
|
|
self.error = "Could not restore attachment draft: \(error.localizedDescription)"
|
|
}
|
|
}
|
|
func updateCapability(_ value: AgentAttachmentCapabilities?) {
|
|
capability = value?.supported == true ? value : nil
|
|
if capability == nil { suspend(); active = true }
|
|
}
|
|
func resume() { active = true; pump() }
|
|
func suspend() {
|
|
active = false; generation = UUID(); worker?.cancel(); worker = nil
|
|
for index in items.indices where items[index].state == .loading || items[index].state == .uploading {
|
|
items[index].state = .failed
|
|
items[index].issue = .local("Transfer interrupted. Retry to continue.", code: "upload_interrupted")
|
|
}
|
|
save()
|
|
}
|
|
@discardableResult
|
|
func add(filename: String, load: @escaping @Sendable () async throws -> AgentPreparedAttachment) -> String? {
|
|
guard supportsAttachments, let capability else { error = AgentAttachmentIssue.unavailable.localizedDescription; return nil }
|
|
guard items.count < capability.countLimit else { error = "You can attach up to \(capability.countLimit) files."; return nil }
|
|
let item = AgentDraftItem(filename: filename)
|
|
items.append(item); loaders[item.id] = load; error = nil
|
|
save(); pump()
|
|
return item.id
|
|
}
|
|
func remove(_ id: String) {
|
|
// Restart the serial worker; completion from the old generation can never reinsert a removed item.
|
|
let interrupted = items.filter { $0.state == .loading || $0.state == .uploading }.map(\.id)
|
|
generation = UUID(); worker?.cancel(); worker = nil
|
|
items.removeAll { $0.id == id }; loaders[id] = nil; prepared[id] = nil
|
|
save()
|
|
if persistenceAvailable { storage?.remove(id) }
|
|
for other in interrupted where other != id {
|
|
if let index = items.firstIndex(where: { $0.id == other }) { items[index].state = .loading }
|
|
}
|
|
pump()
|
|
}
|
|
func retry(_ id: String) {
|
|
guard supportsAttachments else { error = AgentAttachmentIssue.unavailable.localizedDescription; return }
|
|
guard let index = items.firstIndex(where: { $0.id == id }), items[index].state == .failed else { return }
|
|
items[index].state = .loading; items[index].issue = nil; error = nil
|
|
save(); pump()
|
|
}
|
|
func validateReady() throws -> [AgentAttachment] {
|
|
guard items.isEmpty || supportsAttachments else { throw AgentAttachmentIssue.unavailable }
|
|
guard items.allSatisfy({ $0.state == .ready && $0.attachment != nil }) else { throw AgentAttachmentIssue.local("Wait for every attachment to finish, or retry/remove failed files.") }
|
|
if let capability { try capability.validate(ready) }
|
|
return ready
|
|
}
|
|
@discardableResult func clearAfterSubmission() -> Bool {
|
|
guard persistenceAvailable else { return false }
|
|
let nextID = UUID().uuidString
|
|
do { try storage?.write(AgentDraftSnapshot(id: nextID)) }
|
|
catch {
|
|
persistenceAvailable = false
|
|
self.error = "Message is saved in the outbox, but its draft could not be cleared. Reconnect before retrying: \(error.localizedDescription)"
|
|
return false
|
|
}
|
|
let ids = items.map(\.id)
|
|
generation = UUID(); worker?.cancel(); worker = nil
|
|
restoring = true; caption = ""; items = []; id = nextID; restoring = false
|
|
prepared = [:]; loaders = [:]
|
|
ids.forEach { storage?.remove($0) }
|
|
return true
|
|
}
|
|
func save() {
|
|
guard !restoring, persistenceAvailable else { return }
|
|
do { try storage?.write(AgentDraftSnapshot(caption: caption, items: items, id: id)) }
|
|
catch { persistenceAvailable = false; self.error = "Could not save attachment draft: \(error.localizedDescription)" }
|
|
}
|
|
private func pump() {
|
|
guard active, supportsAttachments, let client, let limits = capability, worker == nil,
|
|
let item = items.first(where: { $0.state == .loading }) else { return }
|
|
let run = generation
|
|
worker = Task { [weak self] in
|
|
guard let self else { return }
|
|
defer { if self.generation == run { self.worker = nil; self.pump() } }
|
|
do {
|
|
let file: AgentPreparedAttachment
|
|
if let existing = self.prepared[item.id] { file = existing }
|
|
else if let mime = item.mimeType, let storage = self.storage {
|
|
file = try AgentAttachmentPreparation.prepare(storage.bytes(item.id), filename: item.filename, mimeType: mime, limits: limits)
|
|
} else if let loader = self.loaders[item.id] { file = try await loader() }
|
|
else { throw AgentAttachmentIssue.local("Selection did not finish before the app closed. Remove this item and select it again.", code: "selection_unavailable") }
|
|
try Task.checkCancellation()
|
|
guard self.generation == run, let index = self.items.firstIndex(where: { $0.id == item.id }) else { return }
|
|
let otherBytes = self.items.filter { $0.id != item.id }.reduce(0) { $0 + ($1.sizeBytes ?? 0) }
|
|
guard otherBytes + file.data.count <= limits.totalLimit else { throw AgentAttachmentIssue.local("Attachments exceed the total message size limit.", code: "attachment_limit_exceeded") }
|
|
self.prepared[item.id] = file
|
|
try self.storage?.stage(item.id, file.data)
|
|
self.items[index].filename = file.filename; self.items[index].mimeType = file.mimeType
|
|
self.items[index].sizeBytes = file.data.count; self.items[index].state = .uploading
|
|
self.save()
|
|
guard self.persistenceAvailable else { throw AgentAttachmentIssue.local("Save failed. Reconnect before uploading.") }
|
|
let attachment = try await client.agentUpload(file: file, uploadID: item.id)
|
|
try attachment.validate()
|
|
guard attachment.filename == file.filename, attachment.mimeType == file.mimeType,
|
|
attachment.sizeBytes == file.data.count, attachment.sha256 == file.sha256 else {
|
|
throw AgentAttachmentIssue.local("Upload receipt does not match the original file. Retry to confirm.", code: "invalid_upstream_response")
|
|
}
|
|
try Task.checkCancellation()
|
|
guard self.generation == run, let index = self.items.firstIndex(where: { $0.id == item.id }) else { return }
|
|
self.items[index].attachment = attachment; self.items[index].state = .ready; self.items[index].issue = nil
|
|
self.save()
|
|
self.loaders[item.id] = nil; self.prepared[item.id] = nil
|
|
// Keep staged bytes until the message is durably handed to the outbox.
|
|
} catch {
|
|
guard self.generation == run, let index = self.items.firstIndex(where: { $0.id == item.id }) else { return }
|
|
self.items[index].state = .failed; self.items[index].issue = .wrapping(error); self.save()
|
|
}
|
|
}
|
|
}
|
|
}
|