903 lines
44 KiB
Swift
903 lines
44 KiB
Swift
// SPDX-FileCopyrightText: 2026 Kaede Akino
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
import AgentKit
|
|
import Foundation
|
|
import SwiftData
|
|
import TermindApplication
|
|
|
|
@MainActor
|
|
public final class SwiftDataAgentRunRepository: AgentRunPersisting, Sendable {
|
|
private let context: ModelContext
|
|
private var deletedConversationIDs = Set<UUID>()
|
|
private var deletedRunIDs = Set<UUID>()
|
|
|
|
public init(context: ModelContext) {
|
|
self.context = context
|
|
interruptUnfinishedRuns()
|
|
}
|
|
|
|
@MainActor
|
|
public func save(_ snapshot: AgentRunSnapshot) async throws {
|
|
guard !deletedConversationIDs.contains(snapshot.conversationID) else {
|
|
deletedRunIDs.insert(snapshot.id)
|
|
return
|
|
}
|
|
guard !deletedRunIDs.contains(snapshot.id) else { return }
|
|
let previousDeletedRunIDs = deletedRunIDs
|
|
do {
|
|
let runID = snapshot.id
|
|
let runDescriptor = FetchDescriptor<PersistentAgentRun>(predicate: #Predicate { $0.id == runID })
|
|
let runRow: PersistentAgentRun
|
|
if let row = try context.fetch(runDescriptor).first {
|
|
let persistedStartedAt = row.startedAt
|
|
row.update(snapshot)
|
|
if row.startedAt < persistedStartedAt {
|
|
row.startedAt = persistedStartedAt
|
|
}
|
|
if let finishedAt = row.finishedAt, finishedAt < row.startedAt {
|
|
row.finishedAt = row.startedAt
|
|
}
|
|
runRow = row
|
|
} else {
|
|
let row = PersistentAgentRun(snapshot)
|
|
context.insert(row)
|
|
runRow = row
|
|
}
|
|
|
|
let conversationID = snapshot.conversationID
|
|
let conversationDescriptor = FetchDescriptor<PersistentAgentConversation>(
|
|
predicate: #Predicate { $0.id == conversationID }
|
|
)
|
|
// No title seeding here. `title` means "the generated summary", and
|
|
// empty means there is not one yet — which is what lets the UI decide to
|
|
// fall back to the opening line rather than having that line frozen into
|
|
// storage as though it were a name someone chose.
|
|
if let conversation = try context.fetch(conversationDescriptor).first {
|
|
conversation.updatedAt = .now
|
|
} else {
|
|
context.insert(
|
|
PersistentAgentConversation(
|
|
id: conversationID, title: "", createdAt: snapshot.startedAt
|
|
))
|
|
}
|
|
|
|
// Matched on message id, not on this run's id. A conversation's earlier
|
|
// turns are replayed into every later run, so filtering by `runID` made
|
|
// each save re-insert the whole history — colliding with the unique
|
|
// constraint and rewriting rows that were already correct.
|
|
let durableMessages = snapshot.messages.filter { !$0.isCompaction }
|
|
let ids = Set(durableMessages.map(\.id))
|
|
let stored = try context.fetch(
|
|
FetchDescriptor<PersistentAgentMessage>(
|
|
predicate: #Predicate { ids.contains($0.id) }
|
|
))
|
|
var storedByID = Dictionary(uniqueKeysWithValues: stored.map { ($0.id, $0) })
|
|
if let prompt = durableMessages.last(where: { $0.role == .user }),
|
|
let stagedMessage = storedByID[prompt.id],
|
|
stagedMessage.runID == prompt.id, prompt.id != runID,
|
|
try isStagingReplacementRun(prompt.id, message: stagedMessage, conversationID: conversationID)
|
|
{
|
|
stagedMessage.runID = runID
|
|
stagedMessage.sequence = durableMessages.firstIndex(where: { $0.id == prompt.id }) ?? 0
|
|
let stagingID = prompt.id
|
|
let stagingDescriptor = FetchDescriptor<PersistentAgentRun>(
|
|
predicate: #Predicate { $0.id == stagingID }
|
|
)
|
|
if let stagingRun = try context.fetch(stagingDescriptor).first {
|
|
if runRow.startedAt < stagingRun.startedAt {
|
|
runRow.startedAt = stagingRun.startedAt
|
|
}
|
|
if let finishedAt = runRow.finishedAt, finishedAt < runRow.startedAt {
|
|
runRow.finishedAt = runRow.startedAt
|
|
}
|
|
context.delete(stagingRun)
|
|
}
|
|
deletedRunIDs.insert(stagingID)
|
|
}
|
|
let storedIDs = Set(storedByID.keys)
|
|
var inserted = false
|
|
for (sequence, message) in durableMessages.enumerated()
|
|
where !storedIDs.contains(message.id) {
|
|
let row = PersistentAgentMessage(message, runID: runID, sequence: sequence)
|
|
context.insert(row)
|
|
storedByID[message.id] = row
|
|
insertContextAttachments(for: message)
|
|
inserted = true
|
|
try insertOrigin(for: message)
|
|
}
|
|
// A user row may already exist (editing or interrupted preparation) before
|
|
// the run freezes its context. Add its backing once without rewriting history.
|
|
let replayIDs = Set(
|
|
try context.fetch(FetchDescriptor<PersistentAgentMessageReplay>()).map(\.messageID))
|
|
for message in durableMessages where !replayIDs.contains(message.id) {
|
|
insertReplay(for: message)
|
|
}
|
|
guard inserted || context.hasChanges else { return }
|
|
try context.save()
|
|
} catch {
|
|
context.rollback()
|
|
deletedRunIDs = previousDeletedRunIDs
|
|
throw error
|
|
}
|
|
}
|
|
|
|
@MainActor
|
|
public func persistCompaction(
|
|
_ record: AgentCompactionRecord, conversationID: UUID
|
|
) async throws {
|
|
guard !deletedConversationIDs.contains(conversationID) else {
|
|
throw AgentConversationPersistenceError.conversationNotFound
|
|
}
|
|
let conversations = try context.fetch(
|
|
FetchDescriptor<PersistentAgentConversation>(
|
|
predicate: #Predicate { $0.id == conversationID }
|
|
))
|
|
guard conversations.first != nil else {
|
|
throw AgentConversationPersistenceError.conversationNotFound
|
|
}
|
|
if let boundary = record.compactedThroughMessageID {
|
|
let runIDs = Set(
|
|
try context.fetch(FetchDescriptor<PersistentAgentRun>())
|
|
.filter { $0.conversationID == conversationID }.map(\.id))
|
|
let rows = try context.fetch(
|
|
FetchDescriptor<PersistentAgentMessage>(
|
|
predicate: #Predicate { $0.id == boundary }
|
|
))
|
|
guard rows.contains(where: { runIDs.contains($0.runID) }) else {
|
|
throw AgentConversationPersistenceError.messageNotFound
|
|
}
|
|
}
|
|
let recordID = record.id
|
|
let existing = try context.fetch(
|
|
FetchDescriptor<PersistentAgentCompaction>(
|
|
predicate: #Predicate { $0.id == recordID }
|
|
))
|
|
guard existing.isEmpty else { return }
|
|
context.insert(PersistentAgentCompaction(record, conversationID: conversationID))
|
|
try context.savePendingChanges()
|
|
}
|
|
|
|
@MainActor
|
|
public func record(_ event: AgentEvent, runID: UUID) async throws {
|
|
guard !deletedRunIDs.contains(runID) else { return }
|
|
switch event {
|
|
case .toolProposed(let invocation, _):
|
|
try upsertInvocation(invocation, state: "proposed")
|
|
case .toolStarted(let invocation):
|
|
try upsertInvocation(invocation, state: "started")
|
|
case .toolFinished(let invocation, let result):
|
|
try upsertInvocation(invocation, state: "finished")
|
|
context.insert(PersistentAgentToolResult(result, invocation: invocation))
|
|
case .reviewStarted(let request):
|
|
context.insert(
|
|
PersistentAgentReviewRecord(
|
|
runID: runID, invocationID: request.invocation.id, decision: nil
|
|
))
|
|
case .reviewFinished(let request, let decision):
|
|
context.insert(
|
|
PersistentAgentReviewRecord(
|
|
runID: runID, invocationID: request.invocation.id, decision: decision
|
|
))
|
|
case .reviewFailed(let request, let reason, let timedOut):
|
|
context.insert(
|
|
PersistentAgentReviewRecord(
|
|
runID: runID, invocationID: request.invocation.id,
|
|
failure: reason, timedOut: timedOut
|
|
))
|
|
default:
|
|
// Nothing to write, so nothing to save. `AgentRuntime` already
|
|
// filters on `isDurable`; this is the second half of the same
|
|
// guarantee, for any caller that does not.
|
|
return
|
|
}
|
|
try context.save()
|
|
}
|
|
|
|
@MainActor
|
|
public func loadConversations() async throws -> [StoredAgentConversation] {
|
|
let conversations = try context.fetch(
|
|
FetchDescriptor<PersistentAgentConversation>(
|
|
sortBy: [SortDescriptor(\.updatedAt, order: .reverse)]
|
|
))
|
|
let orderRows = try context.fetch(FetchDescriptor<PersistentAgentConversationOrder>())
|
|
let sortIndexByConversation = Dictionary(
|
|
uniqueKeysWithValues: orderRows.map { ($0.conversationID, $0.sortIndex) }
|
|
)
|
|
let runs = try context.fetch(FetchDescriptor<PersistentAgentRun>())
|
|
let messages = try context.fetch(FetchDescriptor<PersistentAgentMessage>())
|
|
let replayByMessage = Dictionary(
|
|
uniqueKeysWithValues:
|
|
try context.fetch(FetchDescriptor<PersistentAgentMessageReplay>()).map { ($0.messageID, $0) })
|
|
let originByMessage = Dictionary(
|
|
uniqueKeysWithValues:
|
|
try context.fetch(FetchDescriptor<PersistentAgentMessageOrigin>()).map { ($0.messageID, $0.value) })
|
|
let compactionRows = try context.fetch(FetchDescriptor<PersistentAgentCompaction>())
|
|
let attachmentRows = try context.fetch(
|
|
FetchDescriptor<PersistentAgentContextAttachments>()
|
|
)
|
|
let attachmentsByMessage = Dictionary(
|
|
uniqueKeysWithValues: attachmentRows.map { ($0.messageID, $0.attachments) }
|
|
)
|
|
let invocations = try context.fetch(FetchDescriptor<PersistentAgentToolInvocation>())
|
|
let results = try context.fetch(FetchDescriptor<PersistentAgentToolResult>())
|
|
let reviews = try context.fetch(FetchDescriptor<PersistentAgentReviewRecord>())
|
|
return conversations.map { conversation in
|
|
let matchingRuns = runs.filter { $0.conversationID == conversation.id }
|
|
.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
|
|
// their message timestamps were close, which reorders a transcript
|
|
// the model is about to be shown.
|
|
let runOrder = Dictionary(
|
|
uniqueKeysWithValues: matchingRuns.enumerated().map { ($0.element.id, $0.offset) }
|
|
)
|
|
let transcript = messages.filter { runIDs.contains($0.runID) }
|
|
.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 =
|
|
(try? JSONDecoder().decode(
|
|
[AgentToolCall].self, from: row.toolCallsData
|
|
)) ?? []
|
|
let images =
|
|
(try? JSONDecoder().decode(
|
|
[AgentImageAttachment].self, from: row.imagesData
|
|
)) ?? []
|
|
let contextAttachments = attachmentsByMessage[row.id] ?? []
|
|
let reasoning =
|
|
(try? JSONDecoder().decode(
|
|
[AgentReasoningBlock].self, from: row.reasoningData
|
|
)) ?? []
|
|
var text = row.text
|
|
var modelText = replayByMessage[row.id]?.modelText
|
|
if role == .tool {
|
|
let result = results.last { $0.runID == row.runID && $0.callID == row.toolCallID }?
|
|
.restoredResult
|
|
modelText =
|
|
modelText.map {
|
|
AgentPersistenceRedactor.restoredToolResultText($0, fallbacks: [result?.modelContent])
|
|
} ?? result?.modelContent
|
|
text = AgentPersistenceRedactor.restoredToolResultText(
|
|
text, fallbacks: [result?.content, modelText])
|
|
}
|
|
return AgentTranscriptMessage(
|
|
id: row.id, role: role, text: text,
|
|
authoredText: row.authoredText, images: images,
|
|
contextAttachments: contextAttachments,
|
|
reasoning: reasoning,
|
|
toolCalls: calls,
|
|
toolCallID: row.toolCallID, toolName: row.toolName,
|
|
isError: row.isError, isCompaction: row.isCompaction,
|
|
contextSnapshot: replayByMessage[row.id]?.contextSnapshot,
|
|
modelText: modelText,
|
|
origin: originByMessage[row.id],
|
|
createdAt: row.createdAt
|
|
)
|
|
}
|
|
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 { 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: recoveredTranscript,
|
|
records: compactionRows.filter { $0.conversationID == conversation.id }
|
|
.map(\.value)
|
|
)
|
|
return StoredAgentConversation(
|
|
id: conversation.id,
|
|
messages: projection.messages,
|
|
modelTranscript: projection.modelTranscript,
|
|
compactions: projection.records,
|
|
messageRunIDs: messageRunIDs,
|
|
runs: matchingRuns.compactMap { run in
|
|
guard let state = AgentRunState(rawValue: run.state) else { return nil }
|
|
return AgentRunSummary(
|
|
id: run.id, state: state, startedAt: run.startedAt,
|
|
finishedAt: run.finishedAt, failure: run.failure
|
|
)
|
|
},
|
|
toolCards: cards,
|
|
lastRunState: state, title: conversation.title,
|
|
metadata: StoredAgentConversationMetadata(
|
|
customTitle: conversation.customTitle,
|
|
isPinned: conversation.isPinned,
|
|
sortIndex: sortIndexByConversation[conversation.id],
|
|
archivedAt: conversation.archivedAt,
|
|
createdAt: conversation.createdAt,
|
|
updatedAt: conversation.updatedAt
|
|
)
|
|
)
|
|
}
|
|
}
|
|
|
|
@MainActor
|
|
public func setTitle(_ title: String, conversationID: UUID) async throws {
|
|
guard !deletedConversationIDs.contains(conversationID) else { return }
|
|
let descriptor = FetchDescriptor<PersistentAgentConversation>(
|
|
predicate: #Predicate { $0.id == conversationID }
|
|
)
|
|
// A title is model output about remote-influenced text, so it goes
|
|
// through the same redaction every other stored string does.
|
|
guard let row = try context.fetch(descriptor).first else { return }
|
|
row.title = AgentPersistenceRedactor.visibleText(String(title.prefix(80)))
|
|
row.updatedAt = .now
|
|
try context.save()
|
|
}
|
|
|
|
@MainActor
|
|
public func setConversationMetadata(
|
|
_ metadata: StoredAgentConversationMetadata, conversationID: UUID
|
|
) async throws {
|
|
guard !deletedConversationIDs.contains(conversationID) else {
|
|
throw AgentConversationPersistenceError.conversationNotFound
|
|
}
|
|
let descriptor = FetchDescriptor<PersistentAgentConversation>(
|
|
predicate: #Predicate { $0.id == conversationID }
|
|
)
|
|
guard let row = try context.fetch(descriptor).first else {
|
|
throw AgentConversationPersistenceError.conversationNotFound
|
|
}
|
|
row.customTitle = AgentPersistenceRedactor.visibleText(
|
|
String(metadata.customTitle.prefix(80))
|
|
)
|
|
row.isPinned = metadata.isPinned
|
|
row.archivedAt = metadata.archivedAt
|
|
row.updatedAt = metadata.updatedAt
|
|
try context.savePendingChanges()
|
|
}
|
|
|
|
@MainActor
|
|
public func reorderConversations(_ updates: [StoredAgentConversationOrder]) async throws {
|
|
let ids = Set(updates.map(\.id))
|
|
guard updates.count == ids.count else {
|
|
throw AgentConversationPersistenceError.conversationNotFound
|
|
}
|
|
let rows = try context.fetch(
|
|
FetchDescriptor<PersistentAgentConversation>(
|
|
predicate: #Predicate { ids.contains($0.id) }
|
|
))
|
|
guard rows.count == ids.count else {
|
|
throw AgentConversationPersistenceError.conversationNotFound
|
|
}
|
|
let updatesByID = Dictionary(uniqueKeysWithValues: updates.map { ($0.id, $0) })
|
|
let orderRows = try context.fetch(
|
|
FetchDescriptor<PersistentAgentConversationOrder>(
|
|
predicate: #Predicate { ids.contains($0.conversationID) }
|
|
))
|
|
var orderByID = Dictionary(
|
|
uniqueKeysWithValues: orderRows.map { ($0.conversationID, $0) }
|
|
)
|
|
for row in rows {
|
|
guard let update = updatesByID[row.id] else { continue }
|
|
row.isPinned = update.isPinned
|
|
if let order = orderByID[row.id] {
|
|
order.sortIndex = update.sortIndex
|
|
} else {
|
|
let order = PersistentAgentConversationOrder(
|
|
conversationID: row.id, sortIndex: update.sortIndex
|
|
)
|
|
context.insert(order)
|
|
orderByID[row.id] = order
|
|
}
|
|
}
|
|
try context.savePendingChanges()
|
|
}
|
|
|
|
@MainActor
|
|
public func truncateConversation(
|
|
_ id: UUID, fromUserMessageID messageID: UUID,
|
|
replacingWith replacement: AgentTranscriptMessage
|
|
) async throws {
|
|
try Task.checkCancellation()
|
|
guard replacement.role == .user, !replacement.isCompaction,
|
|
replacement.toolCalls.isEmpty, replacement.toolCallID == nil,
|
|
replacement.toolName == nil, !replacement.isError,
|
|
replacement.id != messageID
|
|
else { throw AgentConversationPersistenceError.messageNotFound }
|
|
let allRuns = try context.fetch(FetchDescriptor<PersistentAgentRun>())
|
|
let runs = allRuns.filter { $0.conversationID == id }
|
|
.sorted(by: Self.runPrecedes)
|
|
let allMessages = try context.fetch(FetchDescriptor<PersistentAgentMessage>())
|
|
let runIDs = Set(runs.map(\.id))
|
|
let runOrder = Dictionary(
|
|
uniqueKeysWithValues: runs.enumerated().map { ($0.element.id, $0.offset) })
|
|
let conversationMessages = allMessages.filter { runIDs.contains($0.runID) }
|
|
.sorted { Self.messagePrecedes($0, $1, runOrder: runOrder) }
|
|
guard
|
|
!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 target = conversationMessages[boundary]
|
|
let keptMessages = Array(conversationMessages[..<boundary])
|
|
let removedMessages = Array(conversationMessages[boundary...])
|
|
let removedMessageIDs = Set(removedMessages.map(\.id))
|
|
let laterRuns = Array(runs.dropFirst(targetRunIndex + 1))
|
|
let laterRunIDs = Set(laterRuns.map(\.id))
|
|
let keptTargetMessages = keptMessages.filter { $0.runID == target.runID }
|
|
let originRows = try context.fetch(FetchDescriptor<PersistentAgentMessageOrigin>())
|
|
guard !originRows.contains(where: { $0.messageID == target.id }) else {
|
|
throw AgentConversationPersistenceError.messageNotFound
|
|
}
|
|
let removesTargetRun = keptTargetMessages.isEmpty
|
|
let runsToDelete = removesTargetRun ? [runs[targetRunIndex]] + laterRuns : laterRuns
|
|
let runIDsToDelete = Set(runsToDelete.map(\.id))
|
|
let retiredRunIDs = laterRunIDs.union([target.runID])
|
|
let previousDeletedRunIDs = deletedRunIDs
|
|
deletedRunIDs.formUnion(retiredRunIDs)
|
|
let removesOpeningTurn = !keptMessages.contains {
|
|
$0.role == AgentTranscriptRole.user.rawValue && !$0.isCompaction
|
|
}
|
|
do {
|
|
try deleteGraph(runIDs: runIDsToDelete, runs: runsToDelete)
|
|
if !removesTargetRun,
|
|
let lastKept = keptTargetMessages.last(where: {
|
|
$0.role != AgentTranscriptRole.tool.rawValue
|
|
}) ?? keptTargetMessages.last
|
|
{
|
|
try deleteTargetRunTools(
|
|
runID: target.runID, keptMessages: keptTargetMessages,
|
|
removedMessages: removedMessages.filter { $0.runID == target.runID },
|
|
boundary: target.createdAt)
|
|
try deleteMessageBacking(removedMessages.filter { $0.runID == target.runID })
|
|
let targetRun = runs[targetRunIndex]
|
|
targetRun.state = AgentRunState.completed.rawValue
|
|
// Tool results are omitted from the visible transcript after
|
|
// restore. End the neutral prefix summary at the last retained
|
|
// non-tool message so live editing and reopening agree.
|
|
targetRun.finishedAt = lastKept.createdAt
|
|
targetRun.failure = nil
|
|
}
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentCompaction>())
|
|
where row.conversationID == id
|
|
&& row.compactedThroughMessageID.map(removedMessageIDs.contains) == true
|
|
{
|
|
context.delete(row)
|
|
}
|
|
|
|
let lastKeptRun = removesTargetRun ? runs[..<targetRunIndex].last : runs[targetRunIndex]
|
|
let replacementStartedAt =
|
|
if let lastKeptRun, replacement.createdAt <= lastKeptRun.startedAt {
|
|
lastKeptRun.startedAt.addingTimeInterval(0.001)
|
|
} else {
|
|
replacement.createdAt
|
|
}
|
|
let replacementRun = AgentRunSnapshot(
|
|
id: replacement.id, conversationID: id, state: .completed,
|
|
permissionMode: .askForApproval, messages: [replacement],
|
|
startedAt: replacementStartedAt, finishedAt: replacementStartedAt
|
|
)
|
|
context.insert(PersistentAgentRun(replacementRun))
|
|
context.insert(
|
|
PersistentAgentMessage(
|
|
replacement, runID: replacementRun.id, sequence: 0
|
|
))
|
|
insertContextAttachments(for: replacement)
|
|
insertReplay(for: replacement)
|
|
try insertOrigin(for: replacement)
|
|
|
|
let descriptor = FetchDescriptor<PersistentAgentConversation>(
|
|
predicate: #Predicate { $0.id == id }
|
|
)
|
|
if let conversation = try context.fetch(descriptor).first {
|
|
if removesOpeningTurn { conversation.title = "" }
|
|
conversation.updatedAt = .now
|
|
}
|
|
try context.save()
|
|
} catch {
|
|
context.rollback()
|
|
deletedRunIDs = previousDeletedRunIDs
|
|
throw error
|
|
}
|
|
}
|
|
|
|
@MainActor
|
|
public func createConversation(_ value: StoredAgentConversation) async throws {
|
|
guard !deletedConversationIDs.contains(value.id) else {
|
|
throw AgentConversationPersistenceError.conversationAlreadyExists
|
|
}
|
|
let conversationID = value.id
|
|
let descriptor = FetchDescriptor<PersistentAgentConversation>(
|
|
predicate: #Predicate { $0.id == conversationID }
|
|
)
|
|
guard try context.fetch(descriptor).isEmpty else {
|
|
throw AgentConversationPersistenceError.conversationAlreadyExists
|
|
}
|
|
|
|
let metadata = value.metadata
|
|
let conversation = PersistentAgentConversation(
|
|
id: value.id, title: value.title,
|
|
createdAt: metadata.createdAt, updatedAt: metadata.updatedAt
|
|
)
|
|
conversation.customTitle = metadata.customTitle
|
|
conversation.isPinned = metadata.isPinned
|
|
conversation.archivedAt = metadata.archivedAt
|
|
try context.commit {
|
|
context.insert(conversation)
|
|
if let sortIndex = metadata.sortIndex {
|
|
context.insert(
|
|
PersistentAgentConversationOrder(
|
|
conversationID: value.id, sortIndex: sortIndex
|
|
))
|
|
}
|
|
|
|
let summaries =
|
|
value.runs.isEmpty
|
|
? [
|
|
AgentRunSummary(
|
|
id: UUID(), state: value.lastRunState,
|
|
startedAt: metadata.createdAt, finishedAt: metadata.updatedAt
|
|
)
|
|
]
|
|
: value.runs.sorted { $0.startedAt < $1.startedAt }
|
|
for summary in summaries {
|
|
context.insert(
|
|
PersistentAgentRun(
|
|
AgentRunSnapshot(
|
|
id: summary.id, conversationID: value.id, state: summary.state,
|
|
permissionMode: .askForApproval, messages: [],
|
|
startedAt: summary.startedAt, finishedAt: summary.finishedAt,
|
|
failure: summary.failure
|
|
)))
|
|
}
|
|
let validRunIDs = Set(summaries.map(\.id))
|
|
let fallbackRunID = summaries[0].id
|
|
for (sequence, message) in value.messages.enumerated() {
|
|
guard !message.isCompaction else { continue }
|
|
context.insert(
|
|
PersistentAgentMessage(
|
|
message,
|
|
runID: value.messageRunIDs[message.id].flatMap {
|
|
validRunIDs.contains($0) ? $0 : nil
|
|
} ?? fallbackRunID,
|
|
sequence: sequence
|
|
))
|
|
insertContextAttachments(for: message)
|
|
insertReplay(for: message)
|
|
try insertOrigin(for: message)
|
|
}
|
|
for record in value.compactions {
|
|
context.insert(PersistentAgentCompaction(record, conversationID: value.id))
|
|
}
|
|
for card in value.toolCards {
|
|
var invocation = card.invocation
|
|
if !validRunIDs.contains(invocation.runID) {
|
|
invocation.runID = fallbackRunID
|
|
}
|
|
context.insert(PersistentAgentToolInvocation(invocation, state: card.state))
|
|
if let result = card.result {
|
|
context.insert(PersistentAgentToolResult(result, invocation: invocation))
|
|
}
|
|
if let review = card.review {
|
|
context.insert(
|
|
PersistentAgentReviewRecord(
|
|
runID: invocation.runID, invocationID: invocation.id, decision: review
|
|
))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
@MainActor
|
|
public func deleteConversation(_ id: UUID) async throws {
|
|
deletedConversationIDs.insert(id)
|
|
let runs = try context.fetch(FetchDescriptor<PersistentAgentRun>())
|
|
.filter { $0.conversationID == id }
|
|
let runIDs = Set(runs.map(\.id))
|
|
deletedRunIDs.formUnion(runIDs)
|
|
|
|
try deleteGraph(runIDs: runIDs, runs: runs)
|
|
|
|
let conversations = try context.fetch(FetchDescriptor<PersistentAgentConversation>())
|
|
for row in conversations where row.id == id { context.delete(row) }
|
|
let orderRows = try context.fetch(FetchDescriptor<PersistentAgentConversationOrder>())
|
|
for row in orderRows where row.conversationID == id { context.delete(row) }
|
|
let selectionRows = try context.fetch(FetchDescriptor<PersistentAssistantModelSelection>())
|
|
for row in selectionRows where row.conversationID == id { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentCompaction>())
|
|
where row.conversationID == id { context.delete(row) }
|
|
try context.save()
|
|
}
|
|
|
|
private func deleteGraph(runIDs: Set<UUID>, runs: [PersistentAgentRun]) throws {
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentToolResult>())
|
|
where runIDs.contains(row.runID) { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentReviewRecord>())
|
|
where runIDs.contains(row.runID) { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentToolInvocation>())
|
|
where runIDs.contains(row.runID) { context.delete(row) }
|
|
let messages = try context.fetch(FetchDescriptor<PersistentAgentMessage>())
|
|
.filter { runIDs.contains($0.runID) }
|
|
let messageIDs = Set(messages.map(\.id))
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentContextAttachments>())
|
|
where messageIDs.contains(row.messageID) { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentMessageReplay>())
|
|
where messageIDs.contains(row.messageID) { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentMessageOrigin>())
|
|
where messageIDs.contains(row.messageID) { context.delete(row) }
|
|
for row in messages { context.delete(row) }
|
|
for row in runs { context.delete(row) }
|
|
}
|
|
|
|
private func deleteMessageBacking(_ messages: [PersistentAgentMessage]) throws {
|
|
let messageIDs = Set(messages.map(\.id))
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentContextAttachments>())
|
|
where messageIDs.contains(row.messageID) { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentMessageReplay>())
|
|
where messageIDs.contains(row.messageID) { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentMessageOrigin>())
|
|
where messageIDs.contains(row.messageID) { context.delete(row) }
|
|
for row in messages { context.delete(row) }
|
|
}
|
|
|
|
private func deleteTargetRunTools(
|
|
runID: UUID, keptMessages: [PersistentAgentMessage],
|
|
removedMessages: [PersistentAgentMessage], boundary: Date
|
|
) throws {
|
|
let removedMessageIDs = Set(removedMessages.map(\.id))
|
|
let keptMessageIDs = Set(keptMessages.map(\.id))
|
|
let keptCallIDs = Set(keptMessages.flatMap(Self.toolCallIDs))
|
|
let removedCallIDs = Set(removedMessages.flatMap(Self.toolCallIDs))
|
|
let invocations = try context.fetch(FetchDescriptor<PersistentAgentToolInvocation>())
|
|
let removedInvocations = invocations.filter { invocation in
|
|
guard invocation.runID == runID else { return false }
|
|
if let source = invocation.sourceMessageID {
|
|
if removedMessageIDs.contains(source) { return true }
|
|
if keptMessageIDs.contains(source) { return false }
|
|
}
|
|
if removedCallIDs.contains(invocation.callID) { return true }
|
|
if keptCallIDs.contains(invocation.callID) { return false }
|
|
// Legacy rows with no recoverable source or call ownership belong
|
|
// to the edited suffix at the boundary as well as after it.
|
|
return invocation.createdAt >= boundary
|
|
}
|
|
let invocationIDs = Set(removedInvocations.map(\.id))
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentToolResult>())
|
|
where row.runID == runID && invocationIDs.contains(row.invocationID) { context.delete(row) }
|
|
for row in try context.fetch(FetchDescriptor<PersistentAgentReviewRecord>())
|
|
where row.runID == runID && invocationIDs.contains(row.invocationID) { context.delete(row) }
|
|
for row in removedInvocations { context.delete(row) }
|
|
}
|
|
|
|
private func isStagingReplacementRun(
|
|
_ id: UUID, message: PersistentAgentMessage, conversationID: UUID
|
|
) throws -> Bool {
|
|
let runDescriptor = FetchDescriptor<PersistentAgentRun>(predicate: #Predicate { $0.id == id })
|
|
guard let run = try context.fetch(runDescriptor).first,
|
|
run.conversationID == conversationID, run.state == AgentRunState.completed.rawValue,
|
|
run.failure == nil, message.id == id, message.runID == id,
|
|
message.role == AgentTranscriptRole.user.rawValue, !message.isCompaction
|
|
else { return false }
|
|
let messages = try context.fetch(FetchDescriptor<PersistentAgentMessage>())
|
|
guard messages.filter({ $0.runID == id }).count == 1 else { return false }
|
|
let invocations = try context.fetch(FetchDescriptor<PersistentAgentToolInvocation>())
|
|
let results = try context.fetch(FetchDescriptor<PersistentAgentToolResult>())
|
|
let reviews = try context.fetch(FetchDescriptor<PersistentAgentReviewRecord>())
|
|
return !invocations.contains { $0.runID == id }
|
|
&& !results.contains { $0.runID == id }
|
|
&& !reviews.contains { $0.runID == id }
|
|
}
|
|
|
|
private static func runPrecedes(_ lhs: PersistentAgentRun, _ rhs: PersistentAgentRun) -> Bool {
|
|
if lhs.startedAt != rhs.startedAt { return lhs.startedAt < rhs.startedAt }
|
|
return lhs.id.uuidString < rhs.id.uuidString
|
|
}
|
|
|
|
private static func messagePrecedes(
|
|
_ lhs: PersistentAgentMessage, _ rhs: PersistentAgentMessage, runOrder: [UUID: Int]
|
|
) -> Bool {
|
|
let left = runOrder[lhs.runID] ?? 0
|
|
let right = runOrder[rhs.runID] ?? 0
|
|
if left != right { return left < right }
|
|
if lhs.sequence != rhs.sequence { return lhs.sequence < rhs.sequence }
|
|
if lhs.createdAt != rhs.createdAt { return lhs.createdAt < rhs.createdAt }
|
|
let rank: (String) -> Int = { role in
|
|
switch AgentTranscriptRole(rawValue: role) {
|
|
case .user: 0
|
|
case .assistant: 1
|
|
case .tool: 2
|
|
case nil: 3
|
|
}
|
|
}
|
|
if rank(lhs.role) != rank(rhs.role) { return rank(lhs.role) < rank(rhs.role) }
|
|
return lhs.id.uuidString < rhs.id.uuidString
|
|
}
|
|
|
|
private static func toolCallIDs(_ message: PersistentAgentMessage) -> [String] {
|
|
((try? JSONDecoder().decode([AgentToolCall].self, from: message.toolCallsData)) ?? []).map(\.id)
|
|
}
|
|
|
|
private func insertReplay(for message: AgentTranscriptMessage) {
|
|
guard message.contextSnapshot != nil || message.modelText != nil else { return }
|
|
context.insert(PersistentAgentMessageReplay(message))
|
|
}
|
|
|
|
private func insertContextAttachments(for message: AgentTranscriptMessage) {
|
|
guard !message.contextAttachments.isEmpty else { return }
|
|
context.insert(
|
|
PersistentAgentContextAttachments(
|
|
messageID: message.id, attachments: message.contextAttachments
|
|
))
|
|
}
|
|
|
|
private func insertOrigin(for message: AgentTranscriptMessage) throws {
|
|
guard let origin = message.origin else { return }
|
|
context.insert(PersistentAgentMessageOrigin(messageID: message.id, origin: origin))
|
|
let messageID = message.id
|
|
let receipts = try context.fetch(
|
|
FetchDescriptor<PersistentAgentDeliveryReceipt>(
|
|
predicate: #Predicate { $0.messageID == messageID }))
|
|
for receipt in receipts where receipt.state == AgentMessageDeliveryState.queued.rawValue {
|
|
receipt.state = AgentMessageDeliveryState.delivered.rawValue
|
|
receipt.failure = nil
|
|
receipt.updatedAt = .now
|
|
}
|
|
}
|
|
|
|
/// Puts back tool results that were recorded but never made it into the
|
|
/// transcript.
|
|
///
|
|
/// The two are written on different schedules: a result row lands as its
|
|
/// tool finishes, while its transcript row waits for the run's next
|
|
/// snapshot save. A process killed in between — force quit, crash, the OS
|
|
/// reclaiming the app — leaves the call answered in storage and unanswered
|
|
/// in history, and an assistant turn with an unanswered call is rejected
|
|
/// outright on the next request.
|
|
///
|
|
/// `AgentUnansweredToolCallRepair` would invent a placeholder for the same
|
|
/// gap, and does for anything still missing after this. Doing it here first
|
|
/// is what gives the model the output the tool actually produced instead of
|
|
/// being told the outcome is unknown when it is sitting in the database.
|
|
public static func recovering(
|
|
_ transcript: [AgentTranscriptMessage],
|
|
results: [PersistentAgentToolResult],
|
|
invocations: [PersistentAgentToolInvocation]
|
|
) -> [AgentTranscriptMessage] {
|
|
guard transcript.contains(where: { $0.role == .assistant && !$0.toolCalls.isEmpty })
|
|
else { return transcript }
|
|
|
|
let names = Dictionary(
|
|
invocations.map { ($0.callID, $0.toolName) }, uniquingKeysWith: { first, _ in first }
|
|
)
|
|
var recovered: [AgentTranscriptMessage] = []
|
|
recovered.reserveCapacity(transcript.count)
|
|
var index = transcript.startIndex
|
|
while index < transcript.endIndex {
|
|
let message = transcript[index]
|
|
recovered.append(message)
|
|
index += 1
|
|
guard message.role == .assistant, !message.toolCalls.isEmpty else { continue }
|
|
|
|
var answered: Set<String> = []
|
|
while index < transcript.endIndex, transcript[index].role == .tool {
|
|
if let id = transcript[index].toolCallID { answered.insert(id) }
|
|
recovered.append(transcript[index])
|
|
index += 1
|
|
}
|
|
for call in message.toolCalls where !answered.contains(call.id) {
|
|
// The last recorded one: a call id is unique to its invocation,
|
|
// so more than one row can only mean a rewritten result.
|
|
guard let stored = results.last(where: { $0.callID == call.id }) else { continue }
|
|
let result = stored.restoredResult
|
|
recovered.append(
|
|
AgentTranscriptMessage.toolResult(
|
|
result, runID: stored.runID, toolName: names[call.id] ?? call.name
|
|
))
|
|
}
|
|
}
|
|
return recovered
|
|
}
|
|
|
|
private func upsertInvocation(
|
|
_ invocation: AgentToolInvocation, state: String
|
|
) throws {
|
|
let id = invocation.id
|
|
let fetch = FetchDescriptor<PersistentAgentToolInvocation>(
|
|
predicate: #Predicate { $0.id == id }
|
|
)
|
|
if let row = try context.fetch(fetch).first {
|
|
let rank = ["proposed": 0, "started": 1, "finished": 2]
|
|
if rank[state, default: 0] > rank[row.state, default: 0] { row.state = state }
|
|
if row.sourceMessageID == nil { row.sourceMessageID = invocation.sourceMessageID }
|
|
} else {
|
|
context.insert(PersistentAgentToolInvocation(invocation, state: state))
|
|
}
|
|
}
|
|
|
|
private func storedToolCard(
|
|
_ row: PersistentAgentToolInvocation,
|
|
results: [PersistentAgentToolResult],
|
|
reviews: [PersistentAgentReviewRecord]
|
|
) -> StoredAgentToolCard? {
|
|
guard let arguments = try? AgentJSONValue.decode(row.argumentsData) else { return nil }
|
|
let invocation = AgentToolInvocation(
|
|
id: row.id, runID: row.runID,
|
|
call: AgentToolCall(id: row.callID, name: row.toolName, arguments: arguments),
|
|
sourceMessageID: row.sourceMessageID,
|
|
targetHostID: row.targetHostID, createdAt: row.createdAt
|
|
)
|
|
let result = results.last(where: { $0.invocationID == row.id })?.restoredResult
|
|
let descriptor = AgentToolDescriptor(
|
|
name: row.toolName, summary: row.toolName,
|
|
inputSchema: .object(["type": .string("object")]),
|
|
target: row.targetHostID == nil ? .local : .host,
|
|
approvalPolicy: .ask, presentation: result?.toolPresentation
|
|
)
|
|
let review = reviews.filter { $0.invocationID == row.id }
|
|
.sorted { $0.createdAt < $1.createdAt }
|
|
.last.flatMap { stored -> GuardianDecision? in
|
|
guard let verdict = stored.verdict.flatMap(GuardianVerdict.init),
|
|
let risk = stored.risk.flatMap(GuardianRisk.init),
|
|
let authorization = stored.userAuthorization.flatMap({ raw in
|
|
GuardianAuthorization(rawValue: raw)
|
|
?? (raw == "sufficient" ? .high : raw == "insufficient" ? .low : nil)
|
|
})
|
|
else { return nil }
|
|
return GuardianDecision(
|
|
verdict: verdict, risk: risk, userAuthorization: authorization,
|
|
reason: stored.reason ?? ""
|
|
)
|
|
}
|
|
return StoredAgentToolCard(
|
|
invocation: invocation, descriptor: descriptor, state: row.state,
|
|
result: result, review: review
|
|
)
|
|
}
|
|
|
|
private func interruptUnfinishedRuns() {
|
|
do {
|
|
let rows = try context.fetch(FetchDescriptor<PersistentAgentRun>())
|
|
var changed = false
|
|
for row in rows {
|
|
guard let state = AgentRunState(rawValue: row.state), !state.isTerminal else { continue }
|
|
row.state = AgentRunState.interrupted.rawValue
|
|
row.finishedAt = .now
|
|
changed = true
|
|
}
|
|
if changed { try context.save() }
|
|
} catch { context.rollback() }
|
|
}
|
|
}
|