104 lines
5.6 KiB
Swift
104 lines
5.6 KiB
Swift
// SPDX-FileCopyrightText: 2026 Kaede Akino
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
import AgentKit
|
|
import Foundation
|
|
import SwiftData
|
|
import TermindApplication
|
|
import TermindPersistence
|
|
import Testing
|
|
|
|
@Suite @MainActor
|
|
struct AssistantModelSelectionPersistenceTests {
|
|
@Test func firstSendAndLaterManualChoiceRestoreForTheSameChat() async throws {
|
|
let container = try TermindPersistence.makeContainer(inMemory: true)
|
|
let repository = SwiftDataAgentRunRepository(context: container.mainContext)
|
|
let selections = SwiftDataAssistantModelSelectionStore(context: container.mainContext)
|
|
let provider = ModelProvider(name: "Models", models: [AIModel(id: "first"), AIModel(id: "second")])
|
|
let first = AIModelRef(modelProviderID: provider.id, modelID: "first")
|
|
let second = AIModelRef(modelProviderID: provider.id, modelID: "second")
|
|
let ai = AIConfigurationModel(modelProviders: [provider])
|
|
let controller = AssistantModelSelectionController(ai: ai, store: selections)
|
|
let chat = ConversationModel()
|
|
controller.seed(chat)
|
|
#expect(chat.model == first)
|
|
#expect(try controller.prepareForSend(chat) == first)
|
|
try await repository.save(
|
|
AgentRunSnapshot(
|
|
conversationID: chat.id, state: .completed, permissionMode: .askForApproval,
|
|
messages: [AgentTranscriptMessage(role: .user, text: "Hello")]))
|
|
|
|
let loaded = try await repository.loadConversations()
|
|
let restoredStore = AssistantConversationStore()
|
|
restoredStore.restore(loaded, modelSelections: try controller.selectionsForRestore(loaded))
|
|
let reopened = try #require(restoredStore.conversation(id: chat.id))
|
|
#expect(reopened.model == first)
|
|
try controller.select(second, for: reopened)
|
|
#expect(try selections.load()[chat.id] == second)
|
|
|
|
let reopenedAgain = AssistantConversationStore()
|
|
reopenedAgain.restore(loaded, modelSelections: try controller.selectionsForRestore(loaded))
|
|
#expect(reopenedAgain.conversation(id: chat.id)?.model == second)
|
|
}
|
|
|
|
@Test func v38ConversationGetsOneDurableModelAndDeletionRemovesIt() async throws {
|
|
let directory = FileManager.default.temporaryDirectory
|
|
.appendingPathComponent("TermindChatModelMigration-\(UUID().uuidString)")
|
|
try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
|
|
defer { try? FileManager.default.removeItem(at: directory) }
|
|
let url = directory.appendingPathComponent("Termind.store")
|
|
let oldSchema = Schema(versionedSchema: TermindSchemaV38.self)
|
|
let id = UUID()
|
|
var oldContainer: ModelContainer? = try ModelContainer(
|
|
for: oldSchema,
|
|
configurations: [localModelConfiguration("Old", schema: oldSchema, url: url)])
|
|
oldContainer?.mainContext.insert(PersistentAgentConversation(id: id, title: "Old Chat"))
|
|
try oldContainer?.mainContext.save()
|
|
oldContainer = nil
|
|
|
|
let container = try TermindPersistence.makeContainer(url: url)
|
|
let repository = SwiftDataAgentRunRepository(context: container.mainContext)
|
|
let selections = SwiftDataAssistantModelSelectionStore(context: container.mainContext)
|
|
let provider = ModelProvider(name: "Models", models: [AIModel(id: "current")])
|
|
let model = AIModelRef(modelProviderID: provider.id, modelID: "current")
|
|
let ai = AIConfigurationModel(modelProviders: [provider])
|
|
let controller = AssistantModelSelectionController(ai: ai, store: selections)
|
|
let oldChats = try await repository.loadConversations()
|
|
|
|
#expect(try controller.selectionsForRestore(oldChats)[id] == model)
|
|
#expect(try selections.load()[id] == model)
|
|
ai.assign(nil, to: .chat)
|
|
#expect(try controller.selectionsForRestore(oldChats)[id] == model)
|
|
|
|
try await repository.deleteConversation(id)
|
|
#expect(try selections.load()[id] == nil)
|
|
}
|
|
|
|
@Test func forkKeepsTheSourceModelInItsOwnDurableRecord() async throws {
|
|
let container = try TermindPersistence.makeContainer(inMemory: true)
|
|
let repository = SwiftDataAgentRunRepository(context: container.mainContext)
|
|
let selections = SwiftDataAssistantModelSelectionStore(context: container.mainContext)
|
|
let provider = ModelProvider(name: "Models", models: [AIModel(id: "first"), AIModel(id: "second")])
|
|
let first = AIModelRef(modelProviderID: provider.id, modelID: "first")
|
|
let second = AIModelRef(modelProviderID: provider.id, modelID: "second")
|
|
let controller = AssistantModelSelectionController(
|
|
ai: AIConfigurationModel(modelProviders: [provider]), store: selections)
|
|
let user = AgentTranscriptMessage(role: .user, text: "Hello")
|
|
let answer = AgentTranscriptMessage(role: .assistant, text: "World")
|
|
let source = ConversationModel(transcript: [user, answer], runState: .completed)
|
|
source.model = second
|
|
try selections.save(second, for: source.id)
|
|
let operations = AssistantConversationOperations(repository: repository, modelSelections: controller)
|
|
|
|
let fork = try #require(try await operations.fork(source, at: answer.id, title: "Fork"))
|
|
#expect(fork.model == second)
|
|
#expect(try selections.load()[fork.id] == second)
|
|
#expect(try selections.load()[source.id] == second)
|
|
#expect(try await repository.loadConversations().contains { $0.id == fork.id })
|
|
|
|
try controller.select(first, for: source)
|
|
#expect(try selections.load()[source.id] == first)
|
|
#expect(try selections.load()[fork.id] == second)
|
|
}
|
|
}
|